pass continuations as arguments, use normal stack

src/Gyehoek/CPS/Close.hs

@@ -38,4 +38,4 @@ close = transformM \case
   e -> pure e
 
 closeProgram :: GenSym :> es => Program -> Eff es Program
-closeProgram = traverseOf #body close
+closeProgram = traverseOf (#body . #body) close

src/Gyehoek/CPS/Convert.hs

@@ -109,8 +109,11 @@ convertLambda bs m = do
   pure [cps|(λ (##{bs} #{ktail}) #{m'})|]
 
 convertProgram :: forall es. (GenSym :> es) => Scm.Program -> Eff es Program
-convertProgram p =
-  MkProgram <$> telescope (convert @es) (p ^.. each . _Left) (pure . Halt)
+convertProgram p = do
+  ktail <- gensym' "start-ktail"
+  m <- telescope (convert @es) (p ^.. each . _Left)
+        (pure . ExpContinue (ValVar ktail))
+  pure . MkProgram $ MkLambda [] ktail m
 
 convertExp :: forall es. (GenSym :> es) => Scm.Exp -> Eff es Exp
 convertExp e = convert e (pure . Halt1)

src/Gyehoek/CPS/Eval.hs

@@ -2,6 +2,7 @@
 module Gyehoek.CPS.Eval
   ( evalProgram
   , module Gyehoek.CPS.Syntax
+  , evalExp
   ) where
 
 import Gyehoek.CPS.Syntax
@@ -22,7 +23,7 @@ eval :: Env -> Exp -> List Obj
 
 eval g (Halt xs) = evalVal g <$> xs
 
-eval g (ExpContinue k xs) =
+eval g (ExpContinue ((^?! #ValVar) -> k) xs) =
   case g ^. #labels . at k of
     Just (h, AbsKappa' bs m) -> eval h' m
       where h' = h & #vars <>~ envOfBinds bs (evalVal g <$> xs)
@@ -75,5 +76,11 @@ emptyEnv = MkEnv
     )
   }
 
+evalExp :: Exp -> List Obj
+evalExp = eval emptyEnv
+
 evalProgram :: Program -> List Obj
-evalProgram (MkProgram e) = eval emptyEnv e
+evalProgram (MkProgram lam) = eval emptyEnv [cps|
+  (letrec ((start #{lam}))
+    (apply start halt))
+|]

src/Gyehoek/CPS/Stackify.hs

@@ -1,7 +1,6 @@
 {-# LANGUAGE OverloadedLists #-}
 module Gyehoek.CPS.Stackify
-  ( stackifyExp
-  , stackifyProgram
+  ( stackifyProgram
   , module Gyehoek.CPS.Syntax
   ) where
 
@@ -13,8 +12,11 @@ import Gyehoek.GenSym
 import Effectful.Writer.Static.Shared
 import Data.Foldable
 import qualified Data.HashMap.Strict as H
-import Data.List (elemIndex)
+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
@@ -23,9 +25,10 @@ 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 ->
+-- 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)
+  -- && not (x `elem` g.contStack)
 
 data BlockBuilder
   = Code (List Stk.Instr) BlockBuilder
@@ -47,20 +50,18 @@ stackify
 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
+  m' <- stackify (g & #bound <>~ H.fromList (vs ++ (bindReg <$> ls))) m
   emitRoutine $
     Stk.MkRoutine f xs . buildBlock $
-      Code [Stk.Pop x | x <- ls] m'
+      -- pop in the opposite order we push
+      Code [Stk.Pop x | x <- reverse 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 (ExpLetRec [(f, AbsLambda lam)] e) = do
+  emitRoutine =<< stackifyLambda g f lam
+  stackify (g & #bound . at f ?~ Stk.ValLabel f) e
 
 stackify g (ExpIf c t f) = do
   let c' = stackifyVal g c
@@ -69,36 +70,41 @@ stackify g (ExpIf c t f) = do
   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))
+    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 (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 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'
+    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)
+
+stackifyLambda
+  :: (Stackify :> es, GenSym :> es)
+  => Env -> Name -> Lambda -> Eff es Stk.Routine
+stackifyLambda g name (MkLambda xs k m) = do
+  let vs = [ (x, Stk.ValReg x) | x <- k:xs ]
+  m' <- stackify (g & #bound <>~ H.fromList vs) m
+  pure $ Stk.MkRoutine name (k:xs) (buildBlock m')
+
 stackifyVal :: Env -> Val -> Stk.Val
 stackifyVal g = \case
   ValImm imm -> Stk.ValImm imm
@@ -121,39 +127,30 @@ data Env = MkEnv
   -- 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"]
+emptyEnv = MkEnv mempty mempty
 
 
 
-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
+stackifyProgram (MkProgram lam) = do
+  let g = emptyEnv
+  (start,p) <- runStackify $ stackifyLambda g "start" lam
+  pure $ p <> [ start ]
 
-
-
-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))
+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))))
 |]
+

src/Gyehoek/CPS/Syntax.hs

@@ -121,7 +121,7 @@ data Def = DefConstant Name Exp
   deriving (Show, Generic, Data)
 
 data Program = MkProgram
-  { body :: Exp
+  { body :: Lambda
   }
   deriving (Show, Generic, Data)
 
@@ -190,28 +190,23 @@ instance S.DatumIso Hob where
         (const . Left $ mempty)
 
 instance S.DatumIso Lambda where
-  datumIso = S.match
-    $ S.With (. lambda)
-    $ S.End
+  datumIso = S.with (lam >>>)
     where
-      lambda = S.list $
-        S.el S.lambdaKeyword
-        >>> S.el binders
-        >>> S.el S.datumIso
+      lam :: forall t. G (Datum :- t) (Exp :- Name :- List Name :- t)
+      lam = S.lambdaLike
+        S.lambdaKeyword
+        binders
+        (S.el $ S.datumIso @Exp)
       binders :: forall t. G (Datum :- t) (Name :- List Name :- t)
-      binders = S.list $
-        S.rest (S.datumIso @Name)
-        >>> S.onTail (S.flipped $ IG.PartialIso
-                     (\(ktail:-args:-t) -> (args ++ [ktail]) :- t)
-                     (\(args:-t) -> case args ^? _Snoc of
-                         Just (args',ktail) -> Right $ ktail :- args' :- t
-                         Nothing -> Left $ S.expected "cont param")
-                     )
+      binders =
+        S.list (S.rest $ S.datumIso @Name)
+        >>> S.flipped S.snoced
+        >>> S.swap
 
 instance S.DatumIso Kappa where
   datumIso = S.with \g ->
     S.lambdaLike S.kappaKeyword
-      (S.list $ S.rest (S.datumIso @Name))
+      (S.datumIso @(List Name))
       (S.el $ S.datumIso @Exp)
     >>> g
 
@@ -260,7 +255,7 @@ instance S.DatumIso Exp where
         >>> S.el S.datumIso
 
 instance S.DatumIso Program where
-  datumIso = S.with \prog -> S.datumIso @Exp >>> prog
+  datumIso = S.with \prog -> S.datumIso @Lambda >>> prog
 
 
 -- quasiquoters

src/Gyehoek/Sexp/Grammar/Base.hs

@@ -365,9 +365,9 @@ letLike kw name rhs e = listWithIndentation (NSpecial 1) $
 
 lambdaLike
   :: (forall t. G (Datum :- t) t)
-  -> DatumGrammar a
-  -> G (ListContext :- a :- t) (ListContext :- t')
-  -> G (Datum :- t) t'
+  -> G (Datum :- t1) (a :- t2)
+  -> G (ListContext :- a :- t2) (ListContext :- t3)
+  -> G (Datum :- t1) t3
 lambdaLike kw formals body = listWithIndentation (NSpecial 1) $
   el (decorate SynBuiltin >>> kw)
   >>> el formals

src/Gyehoek/Stack/Lower.hs

@@ -16,9 +16,9 @@ lowerBlock = _
 
 lowerInstr :: Instr -> Wasm.Expr
 lowerInstr = \case
-  PopCont ktail -> [wat|
+  -- PopCont ktail -> [wat|
     
-  |]
+  -- |]
 
 lowerProgram :: Program -> Eff es Wasm.Module
 lowerProgram p = pure [watM|

src/Gyehoek/Stack/Syntax.hs

@@ -60,6 +60,7 @@ data Block = MkBlock
 
 data Tail
   = TailCall Val (List Val)
+  | PushCall Val Val (List Val)
   | If Val Block Block
   deriving stock (Show, Generic, Data)
   deriving anyclass (NFData)
@@ -67,8 +68,6 @@ data Tail
 data Instr
   = Pop Name
   | Push Val
-  | PopCont Name
-  | PushCont Val
   | Prim Name (Prim Val)
   deriving stock (Show, Generic, Data)
   deriving anyclass (NFData)
@@ -91,8 +90,6 @@ instance S.DatumIso Instr where
   datumIso = S.match
     $ S.With (S.headTagged1 "pop!" regName >>>)
     $ S.With (S.headTagged1 "push!" S.datumIso >>>)
-    $ S.With (S.headTagged1 "pop-cont!" regName >>>)
-    $ S.With (S.headTagged1 "push-cont!" S.datumIso >>>)
     $ S.With (S.headTagged2 "prim" regName S.datumIso >>>)
     $ S.End
     where
@@ -108,6 +105,7 @@ instance S.DataIso Block where
 instance S.DatumIso Tail where
   datumIso = S.match
     $ S.With (S.headTagged1' "tail-call" S.datumIso S.datumIso >>>)
+    $ S.With (S.headTagged2' "push-call" S.datumIso S.datumIso S.datumIso >>>)
     $ S.With (if_ >>>)
     $ S.End
     where

src/Gyehoek/Stack/VM.hs

@@ -17,7 +17,6 @@ import Gyehoek.Prelude
 
 data VM = MkVM
   { stack :: List Obj
-  , kstack :: List Name
   , code :: List Instr
   , tail :: Tail
   , registers :: HashMap Name Obj
@@ -40,8 +39,6 @@ stepI :: Env -> VM -> Instr -> VM
 
 stepI e vm (Push v) = vm & #stack %~ (evalVal e vm v :)
 
-stepI e vm (PushCont k) = vm & #kstack %~ (evalToLabel e vm k :)
-
 stepI e vm (Prim r p) = case evalVal e vm <$> p of
   PrimZeroP x -> case x of
     ObjImm (ImmInt n) -> ret . ObjImm . ImmBool $ n == 0
@@ -74,11 +71,6 @@ stepI e vm (Pop r) = case vm ^. #stack of
   (x:xs) -> vm & #registers . at r ?~ x
                & #stack .~ xs
 
-stepI e vm ins@(PopCont r) = case vm ^. #kstack of
-  [] -> error [i|empty cont stack: #{ins}|]
-  (x:xs) -> vm & #registers . at r ?~ ObjImm (ImmLabel x)
-               & #kstack .~ xs
-
 stepI e vm ins = error [i|unimplemented instruction: #{ins}|]
 
 stepT :: Env -> VM -> Tail -> VM
@@ -95,6 +87,9 @@ stepT g vm (TailCall f xs) =
           Nothing -> error [i|undefined label: #{l}|]
           Just x -> x
 
+stepT g vm (PushCall k f xs) =
+  _
+
 stepT g vm (If c t f) = vm & #code .~ branch.code & #tail .~ branch.tail
   where
     branch = case evalVal g vm c of
@@ -116,9 +111,8 @@ evalVal e vm = \case
 initialVM :: VM
 initialVM = MkVM
   { stack = []
-  , kstack = ["halt"]
   , code = []
-  , tail = TailCall (ValLabel "main") []
+  , tail = TailCall (ValLabel "start") [ValLabel "halt"]
   , registers = mempty
   , stdout = ""
   , result = Nothing
@@ -154,3 +148,23 @@ writeObj (ObjImm im) = case im of
   ImmLabel l -> "#<procedure>"
 writeObj (ObjHob h) = case h of
   HobClosure code env -> "#<procedure>"
+
+blah = [stkP|
+(define ($lambda-body0-code7 %lambda-tail1 %lambda-body0 %x)
+  (prim %r2 (* %x %x))
+  (tail-call %lambda-tail1 %r2))
+
+(define ($r3 %x4)
+  (pop! %main-ktail)
+  (tail-call %main-ktail %x4))
+
+(define ($main %main-ktail)
+  (prim %lambda-body0 (make-closure $lambda-body0-code7))
+  (tail-call $let-body5 %lambda-body0))
+
+(define ($let-body5 %square)
+  (pop! %main-ktail)
+  (prim %code6 (env-code %square))
+  (push! %main-ktail)
+  (tail-call %code6 $r3 %square 4))
+|]