module Lang.Lam.Passes where
import FP
import Lang.CPS.Syntax
import Lang.Common
import Lang.Lam.Syntax
import qualified Lang.CPS.Syntax as C
import qualified Lang.Lam.Syntax as L
data StampSt = StampSt
{ stampExpID :: LocNum
, stampBdrID :: BdrNum
}
makeLenses ''StampSt
st0 :: StampSt
st0 = StampSt zer zer
data Env = Env { bdrEnv :: Map Name BdrNum }
makeLenses ''Env
env0 :: Env
env0 = Env mapEmpty
type StampM 𝓈 m = (HasLens 𝓈 StampSt, MonadStateE 𝓈 m, MonadReader Env m)
lookupName :: (StampM 𝓈 m) => Name -> m SName
lookupName x = do
xi <- maybeElim (nextL $ stampBdrIDL <.> view) return *$ lookup x ^$ askL bdrEnvL
return $ Stamped xi x
stampM :: (StampM 𝓈 m) => Exp -> m SExp
stampM pe = do
StampedFix ^@ nextL (stampExpIDL <.> view) <@> case runFix pe of
L.Lit l -> return $ L.Lit l
L.Var x -> L.Var ^$ lookupName x
Lam x e -> do
i <- nextL (stampBdrIDL <.> view)
se <- local (update bdrEnvL $ mapInsert x i) $ stampM e
return $ Lam (Stamped i x) se
L.Prim o e -> L.Prim o ^$ stampM e
L.Let x e b -> do
se <- stampM e
i <- nextL (stampBdrIDL <.> view)
sb <- local (update bdrEnvL $ mapInsert x i) $ stampM b
return $ L.Let (Stamped i x) se sb
App fe e -> App ^@ stampM fe <@> stampM e
L.If e te fe -> L.If ^@ stampM e <@> stampM te <@> stampM fe
stamp :: Exp -> SExp
stamp = evalState st0 . runReaderT env0 . stampM
data CPSSt = CPSSt
{ callID :: LocNum
, bdrID :: BdrNum
, genID :: Int
}
makeLenses ''CPSSt
stampL :: Lens CPSSt StampSt
stampL = lens lin lout
where
lin (CPSSt cid bid _) = StampSt cid bid
lout (CPSSt _ _ gid) (StampSt cid bid) = CPSSt cid bid gid
instance HasLens CPSSt StampSt where
view = stampL
cpsst0 :: CPSSt
cpsst0 = CPSSt zer zer zer
type CPSM m = (MonadOpaqueKon CPSKon SGCall m, MonadState CPSSt m)
fresh :: (CPSM m) => String -> m SGName
fresh x = do
i <- nextL bdrIDL
g <- nextL genIDL
return $ Stamped i $ GName (Just g) $ Name x
data CPSKon r m a where
MetaKon :: (a -> m r) -> CPSKon r m a
ObjectKon :: SGPico -> (SGPico -> m SGCall) -> CPSKon SGCall m SGPico
instance Morphism3 (KFun r) (CPSKon r) where
morph3 (KFun mk) = MetaKon mk
instance Morphism3 (CPSKon r) (KFun r) where
morph3 :: CPSKon r ~~> KFun r
morph3 (MetaKon mk) = KFun mk
morph3 (ObjectKon _ mk) = KFun mk
instance Isomorphism3 (KFun r) (CPSKon r) where
instance Balloon CPSKon SGCall where
inflate :: (Monad m) => CPSKon SGCall m ~> CPSKon SGCall (OpaqueKonT CPSKon SGCall m)
inflate (MetaKon mk) = MetaKon $ \ a -> makeMetaKonT $ \ k -> k *$ mk a
inflate (ObjectKon kx mk) = ObjectKon kx $ \ ax -> makeMetaKonT $ \ k -> k *$ mk ax
deflate :: (Monad m) => CPSKon SGCall (OpaqueKonT CPSKon SGCall m) ~> CPSKon SGCall m
deflate (MetaKon mk) = MetaKon $ \ a -> runMetaKonTWith return $ mk a
deflate (ObjectKon kx mk) = ObjectKon kx $ \ ax -> evalOpaqueKonT $ mk ax
letAtom :: (CPSM m) => String -> SGAtom -> m SGPico
letAtom n a = do
i <- nextL callIDL
x <- fresh n
modifyC (return . StampedFix i . C.Let x a) $
return $ C.Var x
reify :: (CPSM m) => CPSKon SGCall m SGPico -> m SGPico
reify (MetaKon mk) = do
x <- fresh "x"
c <- mk $ C.Var x
i <- nextL callIDL
letAtom "k" $ Stamped i $ LamK x c
reify (ObjectKon k _) = return k
reflect :: (CPSM m) => SGPico -> CPSKon SGCall m SGPico
reflect kx = ObjectKon kx $ \ ax -> do
i <- nextL callIDL
return $ StampedFix i $ AppK kx ax
cpsM :: (CPSM m) => SExp -> m SGPico
cpsM (StampedFix i e0) = case e0 of
L.Lit l ->
return $ C.Lit l
L.Var x -> do
let sx = sgNameFromSName x
return $ C.Var sx
L.Lam x e -> do
let sx = sgNameFromSName x
kx <- fresh "k"
c <- withOpaqueC (reflect $ C.Var kx) $ cpsM e
letAtom "f" $ Stamped i $ LamF sx kx c
L.Prim o e -> do
ex <- cpsM e
letAtom "a" $ Stamped i $ C.Prim o ex
L.Let x e b -> do
ea <- cpsAtomM e
let sx = sgNameFromSName x
callOpaqueCC $ \ (ko :: CPSKon SGCall m SGPico) -> do
bc <- withOpaqueC ko $ cpsM b
return $ StampedFix i $ C.Let sx ea bc
L.App f e -> do
callOpaqueCC $ \ (ko :: CPSKon SGCall m SGPico) -> do
fx <- cpsM f
ex <- cpsM e
ka <- reify ko
return $ StampedFix i $ AppF fx ex ka
L.If ce te fe -> do
callOpaqueCC $ \ (ko :: CPSKon SGCall m SGPico) -> do
cx <- cpsM ce
ko' <- reflect ^$ reify ko
tc <- withOpaqueC ko' $ cpsM te
fc <- withOpaqueC ko' $ cpsM fe
return $ StampedFix i $ C.If cx tc fc
cpsAtomM :: (CPSM m) => SExp -> m SGAtom
cpsAtomM se@(StampedFix i e0) = Stamped i ^$ case e0 of
L.Lam x e -> do
let sx = sgNameFromSName x
kx <- fresh "k"
c <- withOpaqueC (reflect $ C.Var kx) $ cpsM e
return $ LamF sx kx c
L.Prim o e -> do
ea <- cpsM e
return $ C.Prim o ea
_ -> do
ex <- cpsM se
return $ Pico ex
newtype StStateT m a = StStateT { unStStateT :: StateT CPSSt m a }
deriving
( Unit, Functor, Product, Applicative, Bind, Monad
, MonadStateI CPSSt, MonadStateE CPSSt, MonadState CPSSt
, MonadReaderI r, MonadReaderE r, MonadReader r
)
evalStStateT :: (Functor m) => CPSSt -> StStateT m a -> m a
evalStStateT s = evalStateT s . unStStateT
stampCPS :: Exp -> (SExp, SGCall)
stampCPS e = runReader env0 $ evalStStateT cpsst0 $ do
se <- stampM e
c <- runMetaKonT (cpsM se) $ \ ex -> do
i <- nextL callIDL
return $ StampedFix i $ Halt ex
return (se, c)
cps :: Exp -> SGCall
cps = snd . stampCPS