Core.hs

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}

-- NOTE: This file needs to be updated to accommodate new AST shapes.
-- A portion of instances have been updated to include the Err functor;
-- remaining instances are to be updated once this is stable.

module Language.Python.Core
( toplevelCompile
, Bindings
) where

import Prelude hiding (fail)

import           AST.Element
import           Control.Algebra hiding ((:+:))
import           Control.Effect.Reader
import           Core.Core as Core
import           Core.Name as Name
import           Data.Coerce
import           Data.Foldable
import           Data.Function
import           Data.List.NonEmpty (NonEmpty (..))
import           Data.Maybe
import           GHC.Records
import qualified Language.Python.AST as Py
import           Language.Python.Failure
import           Language.Python.Patterns
import           Source.Span (Span)
import           Syntax.Stack (Stack (..))
import qualified Syntax.Stack as Stack

-- | Keeps track of the current scope's bindings (so that we can, when
-- compiling a class or module, return the list of bound variables as
-- a Core record so that all immediate definitions are exposed)
newtype Bindings = Bindings { unBindings :: Stack Name }
  deriving (Eq, Monoid, Semigroup, Show)

def :: Name -> Bindings -> Bindings
def n = coerce (Stack.:> n)

prelude :: Has Core sig t => [Name] -> t Name
prelude = foldl' (...) (pure "__semantic_prelude")

-- We leave the representation of Core syntax abstract so that it's not
-- possible for us to 'cheat' by pattern-matching on or eliminating a
-- compiled term.
type CoreSyntax sig t = ( Has Core sig t
                        , Has (Ann Span) sig t
                        , Has Failure sig t
                        , Foldable t
                        )

class Compile (py :: * -> *) where
  compile :: ( CoreSyntax syn t
             , Has (Reader Bindings) sig m
             )
          => py Span
          -> (t Name -> m (t Name))
          -> (t Name -> m (t Name))

  default compile :: (Applicative m, Has Failure syn t, Show (py Span)) => py Span -> (t Name -> m (t Name)) -> (t Name -> m (t Name))
  compile a _ _ = defaultCompile a

toplevelCompile :: ( CoreSyntax syn t
                   , Has (Reader Bindings) sig m
                   )
                => Py.Module Span
                -> m (t Name)
toplevelCompile py = compile py pure none

-- | TODO: This is not right, it should be a reference to a Preluded
-- NoneType instance, but it will do for now.
none :: Has Core sig t => t Name
none = unit

locate :: ( HasField "ann" syntax Span
          , CoreSyntax syn t
          )
       => syntax
       -> t a
       -> t a
locate syn  = Core.annAt (getField @"ann" syn)

defaultCompile :: (Applicative m, Has Failure syn t, Show py) => py -> m (t Name)
defaultCompile = pure . unimplemented


instance (Compile l, Compile r) => Compile (l :+: r) where
  compile (L1 l) cc = compile l cc
  compile (R1 r) cc = compile r cc

instance Compile Py.AssertStatement
instance Compile Py.Attribute

-- Assignment compilation. Assignments are an uneasy hybrid of expressions
-- (since they appear to have values, i.e. `a = b = c`) and statements (because
-- they introduce bindings). For that reason, they deserve special attention.
--
-- The correct desugaring for the expression above looks like, given a continuation @cont@:
-- @
--  (b :<- c) >>>= (a :<- b) >>>= cont
-- @
-- The tree structure that we get out of tree-sitter is not particularly conducive to expressing
-- this naturally, so we engage in a small desugaring step so that we can turn a list [a, b, c]
-- into a sequenced Core expression using >>>= and a fold through which information—specifically
-- the LHS to assign—flows.

-- RHS represents the right-hand-side of an assignment that we get out of tree-sitter.
-- Desugared is the "terminal" node in a sequence of assignments, i.e. given a = b = c,
-- c will be the terminal node. It is never an assignment.
type RHS = (Py.Assignment :+: Py.AugmentedAssignment) :+: Desugared
type Desugared = Py.ExpressionList :+: Py.Yield

-- We have to pair locations and names, and tuple syntax is harder to
-- read in this case than a happy little constructor.
data Located a = Located Span a

-- Desugaring an RHS involves walking as deeply as possible into an
-- assignment, storing the names we encounter as we go and eventually
-- returning a terminal expression. We have to keep track of which
desugar :: [Located Name]
        -> RHS Span
        -> Either String ([Located Name], Desugared Span)
desugar acc = \case
  Prj Py.Assignment { left = SingleIdentifier name, right = Just rhs, ann} ->
    desugar (Located ann name : acc) rhs
  R1 any -> pure (acc, any)
  other -> Left ("desugar: couldn't desugar RHS " <> show other)

-- This is an algebra that is invoked from a left fold but that
-- returns a function (the 'difference' pattern) so that we can pass
-- information about what RHS we need down the chain: unlike most fold
-- functions, it has four parameters, not three (since our fold
-- returns a function). There's some pun to be made on "collapsing
-- sugar", like "icing" or "sugar water" but I'll leave that as an
-- exercise to the reader.
collapseDesugared :: (CoreSyntax syn t, Has (Reader Bindings) sig m)
                  => Located Name           -- The current LHS to which to assign
                  -> (t Name -> m (t Name)) -- A meta-continuation: it takes a name and returns a continuation
                  -> t Name                 -- The current RHS to which to assign, yielded from an outer continuation
                  -> m (t Name)             -- The properly-sequenced resolut
collapseDesugared (Located loc n) cont rem =
  let assigning = fmap (Core.annAt loc . ((Name.named' n :<- rem) >>>=))
  in assigning (local (def n) (cont (pure n))) -- gotta call local here to record this assignment

instance Compile Py.Assignment where
  compile it@Py.Assignment
    { left = SingleIdentifier name
    , right = Just rhs
    , ann
    } cc next = case desugar [Located ann name] rhs of
      Right (names, val) -> compile val pure next >>= foldr collapseDesugared cc names >>= pure . locate it
      Left msg           -> pure $ unimplemented msg



  compile other _ _ = pure $ invariantViolated ("Unhandled assignment case: " <> show other)

-- End assignment compilation

instance Compile Py.AugmentedAssignment
instance Compile Py.Await
instance Compile Py.BinaryOperator

instance Compile Py.Block where
  compile it@Py.Block{ Py.extraChildren = body} cc
    = fmap (locate it)
    . foldr compile cc body

instance Compile Py.BooleanOperator
instance Compile Py.BreakStatement

instance Compile Py.Call where
  compile it@Py.Call
    { function
    , arguments = L1 Py.ArgumentList { extraChildren = args }
    } cc next = do
    func <- compile function pure next
    let compileArg = \case
          Prj expr -> compile (expr :: Py.Expression Span) pure next
          other    -> pure . invariantViolated $ "Can't compile non-expression function argument: " <> show other

    -- Python function arguments are defined to evaluate left to right.
    args <- traverse compileArg args
    locate it (func $$* args) & cc
  compile it _ _ = pure . invariantViolated $ "can't compile Call node with generator expression: " <> show it

instance Compile Py.ClassDefinition where
  compile it@Py.ClassDefinition { body = pybody, name = Py.Identifier _ann (Name.name -> n) } cc next = do
    let buildTypeCall _ = do
          bindings <- asks @Bindings (toList . unBindings)
          let buildName n = (n, pure n)
              contents = record . fmap buildName $ bindings
              typefn = prelude ["type"]
              object = prelude ["object"]

          pure (typefn $$ Core.string (formatName n) $$ object $$ contents)

    body <- compile pybody buildTypeCall next
    let coreName = Name.named' n
        assignClass = coreName :<- rec coreName body
        continuing = fmap (locate it . (assignClass >>>=))
    continuing (local (def n) (cc next))

instance Compile Py.ComparisonOperator

deriving instance Compile Py.CompoundStatement

instance Compile Py.ConcatenatedString
instance Compile Py.ConditionalExpression
instance Compile Py.ContinueStatement

instance Compile Py.DecoratedDefinition where
  compile it@Py.DecoratedDefinition
    { definition
    , extraChildren = [ Py.Decorator { extraChildren } ]
    } cc next = do
    let thenReassign item = do
          bindings <- asks unBindings
          case bindings of
            _ :> lastbound -> do
              tocall <- compile extraChildren pure next
              let callit go = (pure lastbound .= (tocall $$ pure lastbound)) >>> go
              fmap callit (cc item)
            _ -> pure . invariantViolated $ "Encountered a decorated definition without a corresponding function"
    locate it <$> compile definition thenReassign next
  compile it _ _ = pure . invariantViolated $ "Can't figure out decorated definition " <> show it
instance Compile Py.DeleteStatement
instance Compile Py.Dictionary
instance Compile Py.DictionaryComprehension

instance Compile Py.DottedName where
  compile it@Py.DottedName
    { extraChildren = Py.Identifier { text } :| rest
    } cc _next = do
    let aggregate Py.Identifier { text = inner } x = x ... Name.name inner
        composite = foldr aggregate (pure (Name.name text)) rest
    locate it composite & cc


instance Compile Py.Ellipsis
instance Compile Py.ExecStatement

deriving instance Compile Py.Expression

instance Compile Py.ExpressionStatement where
  compile it@Py.ExpressionStatement { Py.extraChildren = children } cc
    = fmap (locate it)
    . foldr compile cc children

instance Compile Py.ExpressionList where
  compile it@Py.ExpressionList { Py.extraChildren = [child] } cc
    = fmap (locate it)
    . compile child cc
  compile Py.ExpressionList { Py.extraChildren = items } _
    = const . pure . unimplemented $ "ExpressionList of length " <> show items


instance Compile Py.False where
  compile it cc _ = cc $ locate it (bool False)

instance Compile Py.Float
instance Compile Py.ForStatement

instance Compile Py.FunctionDefinition where
  compile it@Py.FunctionDefinition
    { name       = Py.Identifier _ann1 name
    , parameters = Py.Parameters _ann2 parameters
    , body
    } cc next = do
      -- Compile each of the parameters, then the body.
      let parameterMs = fmap param parameters
      if any isNothing parameterMs
        then pure . invariantViolated $ "Couldn't extract parameters"
        else do
          let parameters' = catMaybes parameterMs
          body' <- compile body pure next
          -- Build a lambda.
          let located = locate it (rec (Name.named' (Name.name name)) (lams parameters' body'))
          -- Give it a name (below), then augment the current continuation
          -- with the new name (with 'def'), so that calling contexts know
          -- that we have built an exportable definition.
          assigning located <$> local (def (Name.name name)) (cc next)
    where param (Py.Parameter (Prj (Py.Identifier _pann pname))) = Just . named' . Name.name $ pname
          param _                                                = Nothing
          assigning item f = (Name.named' (Name.name name) :<- item) >>>= f

instance Compile Py.FutureImportStatement
instance Compile Py.GeneratorExpression
instance Compile Py.GlobalStatement

instance Compile Py.Identifier where
  compile Py.Identifier { text } cc _ = cc . pure . Name.name $ text

instance Compile Py.IfStatement where
  compile it@Py.IfStatement{ condition, consequence, alternative} cc next =
    locate it <$> (if' <$> compile condition pure next
                       <*> compile consequence cc next
                       <*> foldr clause (cc next) alternative)
    where clause (R1 Py.ElseClause{ body }) _ = compile body cc next
          clause (L1 Py.ElifClause{ condition, consequence }) rest  =
            if' <$> compile condition pure next <*> compile consequence cc next <*> rest


instance Compile Py.ImportFromStatement
instance Compile Py.ImportStatement
instance Compile Py.Integer

instance Compile Py.Lambda where
  compile it@Py.Lambda
    { body
    , parameters
    } cc next = do
      let unparams (Py.LambdaParameters _ ps) = toList ps
          unparam (Py.Parameter (Prj (Py.Identifier _pann pname))) = Just . named' . Name.name $ pname
          unparam _                                                = Nothing
      body' <- compile body cc next
      let params = maybe [] unparams parameters
      pure . locate it . lams (catMaybes (fmap unparam params)) $ body'

instance Compile Py.List
instance Compile Py.ListComprehension
instance Compile Py.ListSplat

instance Compile Py.Module where
  compile it@Py.Module { Py.extraChildren = stmts } _cc =
    -- This action gets passed to compile, which means it is the
    -- final action taken after the compiling fold finishes. It takes
    -- care of listening for the current set of bound variables (which
    -- is augmented by assignments and function definitions) and
    -- creating a record corresponding to those bindings.
    let buildRecord _ = do
          bindings <- asks @Bindings (toList . unBindings)
          let buildName n = (n, pure n)
          pure . record . fmap buildName $ bindings
    in fmap (locate it) . foldr compile buildRecord stmts

instance Compile Py.NamedExpression

instance Compile Py.None where
  -- None is not an lvalue, and thus always points to the prelude's None.
  compile _it cc _ = cc (prelude ["None"])

instance Compile Py.NonlocalStatement

instance Compile Py.NotOperator where
  compile _it@Py.NotOperator{ argument } cc next = do
    val <- compile argument pure next
    cc (prelude ["not"] $$ val)

instance Compile Py.ParenthesizedExpression where
  compile it@Py.ParenthesizedExpression { extraChildren } cc
    = fmap (locate it)
    . compile extraChildren cc

instance Compile Py.PassStatement where
  compile it@Py.PassStatement {} cc _ = cc $ locate it Core.unit

deriving instance Compile Py.PrimaryExpression

instance Compile Py.PrintStatement

instance Compile Py.ReturnStatement where
  compile it@Py.ReturnStatement { Py.extraChildren = vals } _ next = locate it <$> case vals of
    Nothing -> pure none
    Just Py.ExpressionList { extraChildren = [val] } -> compile val pure next
    Just Py.ExpressionList { extraChildren = vals  } -> pure (invariantViolated ("unimplemented: return statement returning " <> show (length vals) <> " values"))


instance Compile Py.RaiseStatement
instance Compile Py.Set
instance Compile Py.SetComprehension

deriving instance Compile Py.SimpleStatement

instance Compile Py.String where
  compile it@Py.String { extraChildren } cc _ = do
    let extract = \case
          Prj Py.EscapeSequence { text } -> Just text
          _other                         -> Nothing

    let contents = fmap extract extraChildren

    if any isNothing contents
      then pure . invariantViolated $ "Couldn't string-desugar " <> show it
      else let new = prelude ["str", "__slots", "__new__"]
           in cc $ locate it (new $$ Core.string (mconcat (catMaybes contents)))

instance Compile Py.Subscript

instance Compile Py.True where
  compile it cc _next = cc $ locate it (bool True)

instance Compile Py.TryStatement

instance Compile Py.Tuple where
  compile it@Py.Tuple { Py.extraChildren = [] } cc _ = cc $ locate it unit

  compile it _ _                                     = pure $ unimplemented it

instance Compile Py.UnaryOperator
instance Compile Py.WhileStatement
instance Compile Py.WithStatement
instance Compile Py.Yield