171 lines
4.7 KiB
Haskell
171 lines
4.7 KiB
Haskell
{-# LANGUAGE OverloadedLists #-}
|
|
module Gyehoek.CPS.Stackify
|
|
( 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, nub)
|
|
import Data.Text qualified as T
|
|
import Gyehoek.Prelude
|
|
import Debug.Pretty.Simple
|
|
import qualified Gyehoek.Sexp as S
|
|
|
|
|
|
type Stackify = Writer Stk.Program
|
|
|
|
runStackify :: Eff (Stackify : es) a -> Eff es (a, Stk.Program)
|
|
runStackify = runWriter
|
|
|
|
live :: Free a => Env -> a -> List Name
|
|
-- TODO: free' should return an OSet lol
|
|
live g e = nub (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, AbsKappa kap)] e) = do
|
|
stackifyKappa g f kap \g' kap' -> do
|
|
emitRoutine kap'
|
|
stackify g' e
|
|
|
|
stackify g (ExpLetRec [(f, AbsLambda lam)] e) = do
|
|
stackifyLambda g f lam \g' lam' -> do
|
|
emitRoutine lam'
|
|
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.Push (Stk.ValReg l) | l <- ls ] $
|
|
-- Tail (Stk.TailCall (stackifyVal g f) (k : (stackifyVal g <$> xs)))
|
|
-- where
|
|
-- k = var g ktail
|
|
-- ls = fold $ (k ^? #ValImm . #ImmLabel)
|
|
-- >>= \klbl -> g ^. #liveness . at klbl
|
|
|
|
-- stackify g e@(ExpContinue k xs) = do
|
|
-- pure $
|
|
-- Code [ Stk.Push (Stk.ValReg l) | l <- ls ] $
|
|
-- Tail (Stk.TailCall k' (stackifyVal g <$> xs))
|
|
-- where
|
|
-- k' = stackifyVal g k
|
|
-- ls = fold $ (k' ^? #ValImm . #ImmLabel)
|
|
-- >>= \klbl -> g ^. #liveness . at klbl
|
|
|
|
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}|]
|
|
|
|
-- affine
|
|
_ValName :: Traversal' Val Name
|
|
_ValName = failing #ValVar (#ValImm . #ImmLabel)
|
|
|
|
stackifyKappa
|
|
:: (Stackify :> es, GenSym :> es)
|
|
=> Env -> Name -> Kappa
|
|
-> (Env -> Stk.Routine -> Eff es r)
|
|
-> Eff es r
|
|
stackifyKappa g name kap@(MkKappa xs m) w = _
|
|
-- stackifyKappa g name kap@(MkKappa xs m) w = do
|
|
-- let vs = (name, Stk.ValLabel name) : (bindReg <$> xs)
|
|
-- let ls = live g kap
|
|
-- m' <- stackify (g & #bound <>~ H.fromList (vs ++ (bindReg <$> ls))) m
|
|
-- let g' = g & #bound . at name ?~ Stk.ValLabel name
|
|
-- & #liveness . at name ?~ live g kap
|
|
-- let rt = Stk.MkRoutine name xs . buildBlock $
|
|
-- -- pop in the opposite order we push
|
|
-- Code [Stk.Pop x | x <- reverse ls] m'
|
|
-- w g' rt
|
|
|
|
stackifyLambda
|
|
:: (Stackify :> es, GenSym :> es)
|
|
=> Env -> Name -> Lambda
|
|
-> (Env -> Stk.Routine -> Eff es r)
|
|
-> Eff es r
|
|
stackifyLambda g name (MkLambda xs k m) w = do
|
|
let vs = [ (x, Stk.ValReg x) | x <- k:xs ]
|
|
m' <- stackify (g & #bound <>~ H.fromList vs) m
|
|
let g' = g & #bound . at name ?~ Stk.ValLabel name
|
|
w g' $ Stk.MkRoutine name (buildBlock m')
|
|
|
|
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)
|
|
}
|
|
deriving (Show, Generic)
|
|
|
|
emptyEnv :: Env
|
|
emptyEnv = MkEnv mempty mempty
|
|
|
|
|
|
|
|
stackifyProgram :: GenSym :> es => Program -> Eff es Stk.Program
|
|
stackifyProgram (MkProgram lam) = do
|
|
let g = emptyEnv
|
|
(_,p) <- runStackify $ stackifyLambda g "start" lam (const emitRoutine)
|
|
pure p
|
|
|
|
letfn :: Program
|
|
letfn = [cps|
|
|
(λ (start-ktail0)
|
|
(letrec ((lambda-body1
|
|
(λ (x lambda-tail2)
|
|
(prim (* x x) (κ (r3) (continue lambda-tail2 r3))))))
|
|
(letrec ((let-body6
|
|
(κ (square)
|
|
(letrec ((r4 (κ (x5) (continue start-ktail0 x5))))
|
|
(square 4 r4)))))
|
|
(continue let-body6 lambda-body1))))
|
|
|]
|
|
|