160 lines
4.6 KiB
Haskell
160 lines
4.6 KiB
Haskell
{-# LANGUAGE OverloadedLists #-}
|
|
module Gyehoek.CPS.Stackify
|
|
( stackifyExp
|
|
, stackifyProgram
|
|
, module Gyehoek.CPS.Syntax
|
|
) where
|
|
|
|
import Gyehoek.CPS.Syntax
|
|
import Gyehoek.Stack.Syntax qualified as Stk
|
|
import Data.Sequence (Seq)
|
|
import Data.Sequence qualified as Seq
|
|
import Gyehoek.GenSym
|
|
import Effectful.Writer.Static.Shared
|
|
import Data.Foldable
|
|
import qualified Data.HashMap.Strict as H
|
|
import Data.List (elemIndex)
|
|
import Gyehoek.Prelude
|
|
|
|
|
|
type Stackify = Writer Stk.Program
|
|
|
|
runStackify :: Eff (Stackify : es) a -> Eff es (a, Stk.Program)
|
|
runStackify = runWriter
|
|
|
|
live :: Free a => Env -> a -> List Name
|
|
live g e = free' e & filter \x ->
|
|
x `H.member` g.bound
|
|
&& not (x `elem` g.contStack)
|
|
|
|
data BlockBuilder
|
|
= Code (List Stk.Instr) BlockBuilder
|
|
| Tail Stk.Tail
|
|
deriving (Show, Generic)
|
|
|
|
buildBlock :: BlockBuilder -> Stk.Block
|
|
buildBlock = go [] where
|
|
go acc (Code xs bb) = go (acc ++ xs) bb
|
|
go acc (Tail t) = Stk.MkBlock acc t
|
|
|
|
emitRoutine :: Stackify :> es => Stk.Routine -> Eff es ()
|
|
emitRoutine rt = tell [rt]
|
|
|
|
stackify
|
|
:: (GenSym :> es, Stackify :> es)
|
|
=> Env -> Exp -> Eff es BlockBuilder
|
|
|
|
stackify g (ExpLetRec [(f, kap@(AbsKappa' xs m))] e) = do
|
|
let vs = (f, Stk.ValLabel f) : (bindReg <$> xs)
|
|
let ls = live g kap
|
|
m' <- stackify (g & #bound .~ H.fromList (vs ++ (bindReg <$> ls))) m
|
|
emitRoutine $
|
|
Stk.MkRoutine f xs . buildBlock $
|
|
Code [Stk.Pop x | x <- ls] m'
|
|
let g' = g & #bound . at f ?~ Stk.ValLabel f
|
|
& #liveness . at f ?~ ls
|
|
stackify g' e
|
|
|
|
stackify g (ExpLetRec [(f, AbsLambda' xs k m)] e) = do
|
|
let vs = (k:xs) <&> \x -> (x, Stk.ValReg x)
|
|
m' <- stackify (g & #bound .~ H.fromList vs
|
|
& #contStack %~ (k:)) m
|
|
emitRoutine $ Stk.MkRoutine f xs (buildBlock m')
|
|
stackify g e
|
|
|
|
stackify g (ExpIf c t f) = do
|
|
let c' = stackifyVal g c
|
|
t' <- buildBlock <$> stackify g t
|
|
f' <- buildBlock <$> stackify g f
|
|
pure . Tail $ Stk.If c' t' f'
|
|
|
|
stackify g (ExpApply f xs ktail) = pure $
|
|
Code [ Stk.PushCont k ] $
|
|
Code [ Stk.Push (Stk.ValReg l) | l <- ls ] $
|
|
Tail (Stk.TailCall (stackifyVal g f) (stackifyVal g <$> xs))
|
|
where
|
|
k = var g ktail
|
|
ls = fold $ (k ^? #ValImm . #ImmLabel)
|
|
>>= \klbl -> g ^. #liveness . at klbl
|
|
|
|
stackify g (ExpContinue k xs) =
|
|
-- return continuations require popping the stack. how do we know
|
|
-- when a continuation is a return continuation? is this a correct
|
|
-- test?
|
|
case elemIndex k g.contStack of
|
|
Nothing -> pure . Tail $ Stk.TailCall (Stk.ValLabel k) xs'
|
|
Just j -> do
|
|
ktail <- gensym' @Name $ k ^. _Wrapped'
|
|
pure $
|
|
Code (replicate j $ Stk.PopCont "_") $
|
|
Code [Stk.PopCont ktail] $
|
|
Tail (Stk.TailCall (Stk.ValReg ktail) xs')
|
|
where xs' = stackifyVal g <$> xs
|
|
|
|
stackify g (ExpPrim p (MkKappa [x] e)) = do
|
|
e' <- stackify (g & #bound . at x ?~ Stk.ValReg x) e
|
|
pure $
|
|
Code [ Stk.Prim x (stackifyVal g <$> p) ] $
|
|
e'
|
|
|
|
stackify _ e = error [i|unimplemented exp: #{e}|]
|
|
|
|
stackifyVal :: Env -> Val -> Stk.Val
|
|
stackifyVal g = \case
|
|
ValImm imm -> Stk.ValImm imm
|
|
ValVar v -> var g v
|
|
v -> error [i|unimplemented val: #{v}|]
|
|
|
|
var :: Env -> Name -> Stk.Val
|
|
var g v = case g ^. #bound . at v of
|
|
Just x -> x
|
|
Nothing -> Stk.ValLabel v
|
|
|
|
bindReg :: Name -> (Name, Stk.Val)
|
|
bindReg x = (x, Stk.ValReg x)
|
|
|
|
|
|
|
|
data Env = MkEnv
|
|
{ bound :: HashMap Name Stk.Val
|
|
-- | for each locally-bound continuation @k@, @liveness@ has an
|
|
-- entry @(k,ls)@ where @ls@ is the sequence of registers @k@
|
|
-- expects to find saved on the stack.
|
|
, liveness :: HashMap Name (List Name)
|
|
, contStack :: List Name
|
|
}
|
|
deriving (Show, Generic)
|
|
|
|
emptyEnv :: Env
|
|
emptyEnv = MkEnv mempty mempty ["halt"]
|
|
|
|
|
|
|
|
stackifyExp :: GenSym :> es => Name -> Exp -> Eff es Stk.Program
|
|
stackifyExp lbl e = do
|
|
(code,p) <- runStackify $ stackify emptyEnv e
|
|
pure $ p <> [ Stk.MkRoutine lbl [] (buildBlock code) ]
|
|
|
|
stackifyProgram :: GenSym :> es => Program -> Eff es Stk.Program
|
|
stackifyProgram (MkProgram e) = stackifyExp "main" e
|
|
|
|
|
|
|
|
fac :: Program
|
|
fac = [cps|
|
|
(letrec ((fac (λ (n ktail)
|
|
(prim (zero? n)
|
|
(κ (x0)
|
|
(if x0
|
|
(continue ktail 1)
|
|
(prim (- n 1)
|
|
(κ (x1)
|
|
(letrec ((fac-k0
|
|
(κ (x2)
|
|
(prim (* n x2)
|
|
(κ (x3)
|
|
(continue ktail x3))))))
|
|
(fac x1 fac-k0))))))))))
|
|
(fac 6 halt))
|
|
|]
|