module Stg.Marshal.ToStg (
ToStg(..),
) where
import Control.Applicative
import Control.Monad.Trans.Class
import Control.Monad.Trans.Writer
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Map as M
import Data.Monoid
import Data.Text (Text)
import Stg.Language
import qualified Stg.Parser.QuasiQuoter as QQ
import qualified Stg.Prelude.List as Stg
import qualified Stg.Prelude.Maybe as Stg
import Stg.Util
genPrefix :: Text
genPrefix = "__"
class ToStg value where
toStg
:: Var
-> value
-> Program
toStg var val =
let (globals, actualDef) = runWriter (toStgWithGlobals var val)
in globals <> actualDef
toStgWithGlobals
:: Var
-> value
-> Writer Program Program
toStgWithGlobals var val = pure (toStg var val)
instance ToStg () where
toStg name _ = Program (Binds [(name, LambdaForm [] NoUpdate []
(AppC (Constr "Unit") []) )])
instance ToStg Integer where
toStg name i = Program (Binds [(name, LambdaForm [] NoUpdate []
(AppC (Constr "Int#") [AtomLit (Literal i)]) )])
instance ToStg Int where
toStg name i = toStg name (fromIntegral i :: Integer)
instance ToStg Bool where
toStg name b = Program (Binds [(name, LambdaForm [] NoUpdate []
(AppC (Constr (show' b)) []) )])
instance ToStg a => ToStg (Maybe a) where
toStgWithGlobals name Nothing = do
tell Stg.nothing
pure (Program (Binds [(name, [QQ.stg| \ => nothing |])]))
toStgWithGlobals name (Just x) = do
Program xBinding <- toStgWithGlobals justBindName x
pure (Program (Binds [
( name
, LambdaForm [] Update []
(Let NonRecursive
xBinding
(AppC "Just" [AtomVar justBindName]) ))]))
where
justBindName :: Var
justBindName = Var (genPrefix <> "justVal")
instance (ToStg a, ToStg b) => ToStg (Either a b) where
toStgWithGlobals name x = do
let bindName = Var (genPrefix <> chooseEither "left" "right" x <> "val")
Program xBinding <- case x of
Left l -> toStgWithGlobals bindName l
Right r -> toStgWithGlobals bindName r
pure (Program (Binds [
( name
, LambdaForm [] Update []
(Let NonRecursive
xBinding
(AppC (chooseEither "Left" "Right" x) [AtomVar bindName]) ))]))
where
chooseEither l _ (Left _) = l
chooseEither _ r (Right _) = r
instance ToStg a => ToStg [a] where
toStgWithGlobals name dataValues = do
tell Stg.nil
case dataValues of
(x:xs) -> do
(Just inExpression, letBindings)
<- mkListBinds Nothing (NonEmpty.zip [0..] (x :| xs))
let rec = if null xs then NonRecursive else Recursive
pure (Program (Binds [(name, LambdaForm [] Update []
(Let rec letBindings inExpression) )]))
_nil -> pure (Program (Binds [(name, [QQ.stg| \ => nil |])]))
where
mkConsVar :: Int -> Var
mkConsVar i = Var (genPrefix <> show' i <> "_cons")
mkListBinds
:: ToStg value
=> Maybe Expr
-> NonEmpty (Int, value)
-> Writer Program (Maybe Expr, Binds)
mkListBinds inExpression ((i, value) :| rest) = do
let valueVar = Var (genPrefix <> show' i <> "_value")
Program valueBind <- toStgWithGlobals valueVar value
(inExpression', restBinds) <- do
let consVar = mkConsVar i
nextConsVar = if null rest then Var "nil"
else mkConsVar (i+1)
consBind = case inExpression of
Nothing -> mempty
Just _ -> (Binds . M.singleton consVar) (LambdaForm
(valueVar : [nextConsVar | not (null rest)])
NoUpdate
[]
consExpr )
consExpr = AppC (Constr "Cons") (map AtomVar [valueVar, nextConsVar])
inExpression' = inExpression <|> Just consExpr
recursiveBinds <- case rest of
(i',v') : isvs -> fmap snd (mkListBinds inExpression' ((i',v') :| isvs))
_nil -> pure mempty
pure (inExpression', consBind <> recursiveBinds)
pure (inExpression', valueBind <> restBinds)
tupleEntry
:: ToStg value
=> Text
-> value
-> WriterT Binds (Writer Program) ()
tupleEntry name val = do
let bindName = Var (genPrefix <> name)
Program bind <- lift (toStgWithGlobals bindName val)
tell bind
tupleBinds
:: Var
-> Constr
-> Binds
-> Binds
tupleBinds name tupleCon binds =
let bindVars =
let Binds b = binds
in M.keys b
in Binds [(name,
LambdaForm [] Update []
(Let NonRecursive
binds
(AppC tupleCon (map AtomVar bindVars)) ))]
instance (ToStg a, ToStg b) => ToStg (a,b) where
toStgWithGlobals name (x,y) = do
binds <- execWriterT (do
tupleEntry "fst" x
tupleEntry "snd" y )
pure (Program (tupleBinds name (Constr "Pair") binds))
instance (ToStg a, ToStg b, ToStg c) => ToStg (a,b,c) where
toStgWithGlobals name (x,y,z) = do
binds <- execWriterT (do
tupleEntry "x" x
tupleEntry "y" y
tupleEntry "z" z )
pure (Program (tupleBinds name (Constr "Triple") binds))
instance (ToStg a, ToStg b, ToStg c, ToStg d) => ToStg (a,b,c,d) where
toStgWithGlobals name (w4,x4,y4,z4) = do
binds <- execWriterT (do
tupleEntry "w" w4
tupleEntry "x" x4
tupleEntry "y" y4
tupleEntry "z" z4 )
pure (Program (tupleBinds name (Constr "Quadruple") binds))
instance (ToStg a, ToStg b, ToStg c, ToStg d, ToStg e) => ToStg (a,b,c,d,e) where
toStgWithGlobals name (v5,w5,x5,y5,z5) = do
binds <- execWriterT (do
tupleEntry "v" v5
tupleEntry "w" w5
tupleEntry "x" x5
tupleEntry "y" y5
tupleEntry "z" z5 )
pure (Program (tupleBinds name (Constr "Quintuple") binds))