{-# LANGUAGE FlexibleContexts #-}

module Transformations.Anf.AnfGen (genAnf) where

import Control.Monad.Cont (ContT, mapContT)
import Control.Monad.State (MonadState, MonadTrans (lift), State, evalState, get, modify)
import Control.Monad.Trans.Cont (evalContT)
import Data.Text (pack)
import qualified Transformations.Anf.Anf as Anf
import qualified Transformations.Ll.Lfr as Lfr
import qualified Trees.Common as Common

-- * ANF Generator

genAnf :: Lfr.Program -> Anf.Program
genAnf :: Program -> Program
genAnf (Lfr.Program [GlobalDeclaration]
gDecls Env
cnt) =
  let gDecls' :: [GlobalDeclaration]
gDecls' = State Env [GlobalDeclaration] -> Env -> [GlobalDeclaration]
forall s a. State s a -> s -> a
evalState ((GlobalDeclaration -> StateT Env Identity GlobalDeclaration)
-> [GlobalDeclaration] -> State Env [GlobalDeclaration]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM GlobalDeclaration -> StateT Env Identity GlobalDeclaration
genDecl [GlobalDeclaration]
gDecls) Env
cnt
   in [GlobalDeclaration] -> Program
Anf.Program [GlobalDeclaration]
gDecls'

-- * Internal

-- ** ANF Normalizer Continuation & Generator State

type NormCont r a = ContT r AnfGenState a

type AnfGenState = State Env

type Env = Common.IdCnt

-- ** ANF Generators

genDecl :: Lfr.GlobalDeclaration -> AnfGenState Anf.GlobalDeclaration
genDecl :: GlobalDeclaration -> StateT Env Identity GlobalDeclaration
genDecl (Lfr.GlobVarDecl (Lfr.VarDecl Identifier'
ident Expression
value)) = Identifier' -> Expression -> GlobalDeclaration
Anf.GlobVarDecl Identifier'
ident (Expression -> GlobalDeclaration)
-> StateT Env Identity Expression
-> StateT Env Identity GlobalDeclaration
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> StateT Env Identity Expression
genExpr Expression
value
genDecl (Lfr.GlobFunDecl Identifier'
ident [Identifier']
params Expression
body) = Identifier' -> [Identifier'] -> Expression -> GlobalDeclaration
Anf.GlobFunDecl Identifier'
ident [Identifier']
params (Expression -> GlobalDeclaration)
-> StateT Env Identity Expression
-> StateT Env Identity GlobalDeclaration
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expression -> StateT Env Identity Expression
genExpr Expression
body

genExpr :: Lfr.Expression -> AnfGenState Anf.Expression
genExpr :: Expression -> StateT Env Identity Expression
genExpr (Lfr.ExprId Identifier'
ident) = AtomicExpression -> StateT Env Identity Expression
forall (m :: * -> *).
MonadState Env m =>
AtomicExpression -> m Expression
returnAtom (AtomicExpression -> StateT Env Identity Expression)
-> AtomicExpression -> StateT Env Identity Expression
forall a b. (a -> b) -> a -> b
$ Identifier' -> AtomicExpression
Anf.AtomId Identifier'
ident
genExpr (Lfr.ExprPrimVal PrimitiveValue
val) = AtomicExpression -> StateT Env Identity Expression
forall (m :: * -> *).
MonadState Env m =>
AtomicExpression -> m Expression
returnAtom (AtomicExpression -> StateT Env Identity Expression)
-> AtomicExpression -> StateT Env Identity Expression
forall a b. (a -> b) -> a -> b
$ case PrimitiveValue
val of
  PrimitiveValue
Common.PrimValUnit -> AtomicExpression
Anf.AtomUnit
  Common.PrimValBool Bool
bool -> Bool -> AtomicExpression
Anf.AtomBool Bool
bool
  Common.PrimValInt Int64
int -> Int64 -> AtomicExpression
Anf.AtomInt Int64
int
genExpr (Lfr.ExprBinOp BinaryOperator
op Expression
lhs Expression
rhs) = ContT Expression (StateT Env Identity) Expression
-> StateT Env Identity Expression
forall (m :: * -> *) r. Monad m => ContT r m r -> m r
evalContT (ContT Expression (StateT Env Identity) Expression
 -> StateT Env Identity Expression)
-> ContT Expression (StateT Env Identity) Expression
-> StateT Env Identity Expression
forall a b. (a -> b) -> a -> b
$ do
  AtomicExpression
lhs' <- Expression -> NormCont Expression AtomicExpression
normalizeToAtom Expression
lhs
  AtomicExpression
rhs' <- Expression -> NormCont Expression AtomicExpression
normalizeToAtom Expression
rhs
  ComplexExpression
-> ContT Expression (StateT Env Identity) Expression
forall (m :: * -> *).
MonadState Env m =>
ComplexExpression -> m Expression
returnComplex (ComplexExpression
 -> ContT Expression (StateT Env Identity) Expression)
-> ComplexExpression
-> ContT Expression (StateT Env Identity) Expression
forall a b. (a -> b) -> a -> b
$ BinaryOperator
-> AtomicExpression -> AtomicExpression -> ComplexExpression
Anf.CompBinOp BinaryOperator
op AtomicExpression
lhs' AtomicExpression
rhs'
genExpr (Lfr.ExprUnOp UnaryOperator
op Expression
x) = ContT Expression (StateT Env Identity) Expression
-> StateT Env Identity Expression
forall (m :: * -> *) r. Monad m => ContT r m r -> m r
evalContT (ContT Expression (StateT Env Identity) Expression
 -> StateT Env Identity Expression)
-> ContT Expression (StateT Env Identity) Expression
-> StateT Env Identity Expression
forall a b. (a -> b) -> a -> b
$ do
  AtomicExpression
x' <- Expression -> NormCont Expression AtomicExpression
normalizeToAtom Expression
x
  ComplexExpression
-> ContT Expression (StateT Env Identity) Expression
forall (m :: * -> *).
MonadState Env m =>
ComplexExpression -> m Expression
returnComplex (ComplexExpression
 -> ContT Expression (StateT Env Identity) Expression)
-> ComplexExpression
-> ContT Expression (StateT Env Identity) Expression
forall a b. (a -> b) -> a -> b
$ UnaryOperator -> AtomicExpression -> ComplexExpression
Anf.CompUnOp UnaryOperator
op AtomicExpression
x'
genExpr (Lfr.ExprApp Expression
f Expression
arg) = ContT Expression (StateT Env Identity) Expression
-> StateT Env Identity Expression
forall (m :: * -> *) r. Monad m => ContT r m r -> m r
evalContT (ContT Expression (StateT Env Identity) Expression
 -> StateT Env Identity Expression)
-> ContT Expression (StateT Env Identity) Expression
-> StateT Env Identity Expression
forall a b. (a -> b) -> a -> b
$ do
  Identifier'
f' <- Expression -> NormCont Expression Identifier'
normalizeToId Expression
f
  AtomicExpression
arg' <- Expression -> NormCont Expression AtomicExpression
normalizeToAtom Expression
arg
  ComplexExpression
-> ContT Expression (StateT Env Identity) Expression
forall (m :: * -> *).
MonadState Env m =>
ComplexExpression -> m Expression
returnComplex (ComplexExpression
 -> ContT Expression (StateT Env Identity) Expression)
-> ComplexExpression
-> ContT Expression (StateT Env Identity) Expression
forall a b. (a -> b) -> a -> b
$ Identifier' -> AtomicExpression -> ComplexExpression
Anf.CompApp Identifier'
f' AtomicExpression
arg'
genExpr (Lfr.ExprIte Expression
c Expression
t Expression
e) = ContT Expression (StateT Env Identity) Expression
-> StateT Env Identity Expression
forall (m :: * -> *) r. Monad m => ContT r m r -> m r
evalContT (ContT Expression (StateT Env Identity) Expression
 -> StateT Env Identity Expression)
-> ContT Expression (StateT Env Identity) Expression
-> StateT Env Identity Expression
forall a b. (a -> b) -> a -> b
$ do
  AtomicExpression
c' <- Expression -> NormCont Expression AtomicExpression
normalizeToAtom Expression
c
  Expression
t' <- StateT Env Identity Expression
-> ContT Expression (StateT Env Identity) Expression
forall (m :: * -> *) a. Monad m => m a -> ContT Expression m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT Env Identity Expression
 -> ContT Expression (StateT Env Identity) Expression)
-> StateT Env Identity Expression
-> ContT Expression (StateT Env Identity) Expression
forall a b. (a -> b) -> a -> b
$ Expression -> StateT Env Identity Expression
genExpr Expression
t
  Expression
e' <- StateT Env Identity Expression
-> ContT Expression (StateT Env Identity) Expression
forall (m :: * -> *) a. Monad m => m a -> ContT Expression m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT Env Identity Expression
 -> ContT Expression (StateT Env Identity) Expression)
-> StateT Env Identity Expression
-> ContT Expression (StateT Env Identity) Expression
forall a b. (a -> b) -> a -> b
$ Expression -> StateT Env Identity Expression
genExpr Expression
e
  ComplexExpression
-> ContT Expression (StateT Env Identity) Expression
forall (m :: * -> *).
MonadState Env m =>
ComplexExpression -> m Expression
returnComplex (ComplexExpression
 -> ContT Expression (StateT Env Identity) Expression)
-> ComplexExpression
-> ContT Expression (StateT Env Identity) Expression
forall a b. (a -> b) -> a -> b
$ AtomicExpression -> Expression -> Expression -> ComplexExpression
Anf.CompIte AtomicExpression
c' Expression
t' Expression
e'
genExpr (Lfr.ExprLetIn (Lfr.VarDecl Identifier'
ident Expression
val) Expression
expr) = do
  Expression
val' <- Expression -> StateT Env Identity Expression
genExpr Expression
val
  Expression
expr' <- Expression -> StateT Env Identity Expression
genExpr Expression
expr
  Expression -> StateT Env Identity Expression
forall a. a -> StateT Env Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expression -> StateT Env Identity Expression)
-> Expression -> StateT Env Identity Expression
forall a b. (a -> b) -> a -> b
$ (Identifier', Expression) -> Expression -> Expression
Anf.ExprLetIn (Identifier'
ident, Expression
val') Expression
expr'

returnAtom :: (MonadState Env m) => Anf.AtomicExpression -> m Anf.Expression
returnAtom :: forall (m :: * -> *).
MonadState Env m =>
AtomicExpression -> m Expression
returnAtom = Expression -> m Expression
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expression -> m Expression)
-> (AtomicExpression -> Expression)
-> AtomicExpression
-> m Expression
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AtomicExpression -> Expression
Anf.ExprAtom

returnComplex :: (MonadState Env m) => Anf.ComplexExpression -> m Anf.Expression
returnComplex :: forall (m :: * -> *).
MonadState Env m =>
ComplexExpression -> m Expression
returnComplex = Expression -> m Expression
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Expression -> m Expression)
-> (ComplexExpression -> Expression)
-> ComplexExpression
-> m Expression
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ComplexExpression -> Expression
Anf.ExprComp

-- ** Normalizers

normalizeToAtom :: Lfr.Expression -> NormCont Anf.Expression Anf.AtomicExpression
normalizeToAtom :: Expression -> NormCont Expression AtomicExpression
normalizeToAtom Expression
expr = do
  Expression
expr' <- StateT Env Identity Expression
-> ContT Expression (StateT Env Identity) Expression
forall (m :: * -> *) a. Monad m => m a -> ContT Expression m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT Env Identity Expression
 -> ContT Expression (StateT Env Identity) Expression)
-> StateT Env Identity Expression
-> ContT Expression (StateT Env Identity) Expression
forall a b. (a -> b) -> a -> b
$ Expression -> StateT Env Identity Expression
genExpr Expression
expr
  case Expression
expr' of
    Anf.ExprAtom AtomicExpression
atom -> AtomicExpression -> NormCont Expression AtomicExpression
forall a. a -> ContT Expression (StateT Env Identity) a
forall (m :: * -> *) a. Monad m => a -> m a
return AtomicExpression
atom
    Expression
_ -> do
      Identifier'
ident <- AnfGenState Identifier' -> NormCont Expression Identifier'
forall (m :: * -> *) a. Monad m => m a -> ContT Expression m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift AnfGenState Identifier'
genId
      (StateT Env Identity Expression -> StateT Env Identity Expression)
-> NormCont Expression AtomicExpression
-> NormCont Expression AtomicExpression
forall {k} (m :: k -> *) (r :: k) a.
(m r -> m r) -> ContT r m a -> ContT r m a
mapContT
        (\StateT Env Identity Expression
e -> (Identifier', Expression) -> Expression -> Expression
Anf.ExprLetIn (Identifier'
ident, Expression
expr') (Expression -> Expression)
-> StateT Env Identity Expression -> StateT Env Identity Expression
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT Env Identity Expression
e)
        (AtomicExpression -> NormCont Expression AtomicExpression
forall a. a -> ContT Expression (StateT Env Identity) a
forall (m :: * -> *) a. Monad m => a -> m a
return (AtomicExpression -> NormCont Expression AtomicExpression)
-> AtomicExpression -> NormCont Expression AtomicExpression
forall a b. (a -> b) -> a -> b
$ Identifier' -> AtomicExpression
Anf.AtomId Identifier'
ident)

normalizeToId :: Lfr.Expression -> NormCont Anf.Expression Common.Identifier'
normalizeToId :: Expression -> NormCont Expression Identifier'
normalizeToId Expression
expr = do
  Expression
expr' <- StateT Env Identity Expression
-> ContT Expression (StateT Env Identity) Expression
forall (m :: * -> *) a. Monad m => m a -> ContT Expression m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT Env Identity Expression
 -> ContT Expression (StateT Env Identity) Expression)
-> StateT Env Identity Expression
-> ContT Expression (StateT Env Identity) Expression
forall a b. (a -> b) -> a -> b
$ Expression -> StateT Env Identity Expression
genExpr Expression
expr
  case Expression
expr' of
    Anf.ExprAtom (Anf.AtomId Identifier'
ident) -> Identifier' -> NormCont Expression Identifier'
forall a. a -> ContT Expression (StateT Env Identity) a
forall (m :: * -> *) a. Monad m => a -> m a
return Identifier'
ident
    Expression
_ -> do
      Identifier'
ident <- AnfGenState Identifier' -> NormCont Expression Identifier'
forall (m :: * -> *) a. Monad m => m a -> ContT Expression m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift AnfGenState Identifier'
genId
      (StateT Env Identity Expression -> StateT Env Identity Expression)
-> NormCont Expression Identifier'
-> NormCont Expression Identifier'
forall {k} (m :: k -> *) (r :: k) a.
(m r -> m r) -> ContT r m a -> ContT r m a
mapContT
        (\StateT Env Identity Expression
e -> (Identifier', Expression) -> Expression -> Expression
Anf.ExprLetIn (Identifier'
ident, Expression
expr') (Expression -> Expression)
-> StateT Env Identity Expression -> StateT Env Identity Expression
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT Env Identity Expression
e)
        (Identifier' -> NormCont Expression Identifier'
forall a. a -> ContT Expression (StateT Env Identity) a
forall (m :: * -> *) a. Monad m => a -> m a
return Identifier'
ident)

-- ** Identifier Generation

genId :: AnfGenState Common.Identifier'
genId :: AnfGenState Identifier'
genId = do
  Env
cnt <- StateT Env Identity Env
forall s (m :: * -> *). MonadState s m => m s
get
  (Env -> Env) -> StateT Env Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (Env -> Env -> Env
forall a. Num a => a -> a -> a
+ Env
1)
  Identifier' -> AnfGenState Identifier'
forall a. a -> StateT Env Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return (Identifier' -> AnfGenState Identifier')
-> Identifier' -> AnfGenState Identifier'
forall a b. (a -> b) -> a -> b
$ Env -> Identifier -> Identifier'
Common.Gen Env
cnt (Identifier -> Identifier') -> Identifier -> Identifier'
forall a b. (a -> b) -> a -> b
$ String -> Identifier
pack String
"anf"