Files
gyehoek-hs/src/Gyehoek/CPS/Stackify.hs
T
2026-08-28 11:39:37 -06:00

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))))
|]