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

test/Gyehoek/Test/CPS/Eval.hs

@@ -34,8 +34,8 @@ test_cpsInterpreter = testGroup "cps interpreter" $
       |]
   ]
 
-evalsTo :: HasCallStack => List Obj -> Sut.Program -> Assertion
-evalsTo rs p = Sut.evalProgram p @?= rs
+evalsTo :: HasCallStack => List Obj -> Sut.Exp -> Assertion
+evalsTo rs e = Sut.evalExp e @?= rs
 
 primitives = testGroup "primitives"
   [ testGroup "arith"

test/Gyehoek/Test/CPS/Stackify.hs

@@ -4,10 +4,11 @@ import Test.Tasty (TestTree, testGroup)
 import Test.Tasty.HUnit
 import qualified Gyehoek.CPS.Stackify as Sut
 import Gyehoek.Stack.VM as Stk
-import Data.List (List)
 import Gyehoek.CPS.Syntax (cps)
+import Gyehoek.CPS.Syntax qualified as CPS
 import Gyehoek.GenSym (runGenSym)
 import Effectful
+import Gyehoek.Prelude
 
 
 test_stackify =
@@ -19,9 +20,11 @@ test_stackify =
   ]
 
 evalsTo :: List Obj -> Sut.Exp -> Assertion
-evalsTo rs e =
-  Stk.eval e' @?= rs
-  where e' = runPureEff . runGenSym $ Sut.stackifyExp "main" e
+evalsTo rs e = Stk.eval e' @?= rs
+  where
+    e' = e & CPS.MkLambda [] "_ktail"
+      & CPS.MkProgram
+      & Sut.stackifyProgram & runGenSym & runPureEff
 
 trivialReturn = testGroup "trivial return"
   [ testCase "return int" do

test/Gyehoek/Test/CPS/Syntax.hs

@@ -28,15 +28,20 @@ free = testGroup "free"
 qq :: TestTree
 qq = testGroup "parser"
   [ testCase "lambda" do
-      assertEqual "" (Sut.MkLambda ["x","y"] "ktail"
-               (Sut.ExpContinue (Sut.ValLabel "ktail") [Sut.ValVar "x"]))
+      assertEqual ""
+        (Sut.MkLambda ["x","y"] "ktail"
+         (Sut.ExpContinue (Sut.ValVar "ktail") [Sut.ValVar "x"]))
         [cps|(λ (x y ktail) (continue ktail x))|]
-      assertEqual "" (Sut.MkLambda [] "ktail"
-               (Sut.ExpContinue (Sut.ValLabel "ktail") [Sut.ValVar "x"]))
+      assertEqual ""
+        (Sut.MkLambda [] "ktail"
+         (Sut.ExpContinue (Sut.ValVar "ktail") [Sut.ValVar "x"]))
         [cps|(λ (ktail) (continue ktail x))|]
   , testCase "kappa" do
-      assertEqual "" (Sut.MkKappa ["x","y"]
-               (Sut.ExpContinue (Sut.ValLabel "k123") [Sut.ValVar "x", Sut.ValVar "y"]))
+      assertEqual ""
+        (Sut.MkKappa ["x","y"]
+         (Sut.ExpContinue
+          (Sut.ValVar "k123")
+          [Sut.ValVar "x", Sut.ValVar "y"]))
         [cps|(κ (x y) (continue k123 x y))|]
   , testCase "application" do
       assertEqual "" (Sut.ExpApply (Sut.ValVar "f")

test/Gyehoek/Test/Stack/VM.hs

@@ -14,58 +14,58 @@ evalsTo rs p = Sut.eval p @?= rs
 test_root = testGroup "stack machine"
   [ testCase "lit int" do
       evalsTo [ObjImm (ImmInt 3)] [stkP|
-        (define ($main)
-          (pop-cont! %ktail)
+        (define ($start %ktail)
           (tail-call %ktail 3))
       |]
   , testCase "return constant" do
       evalsTo [ObjImm (ImmInt 123)] [stkP|
-        (define ($main)
-          (tail-call $silly))
-        (define ($silly)
-          (pop-cont! %ktail)
+        (define ($start %ktail)
+          (tail-call $silly %ktail))
+        (define ($silly %ktail)
           (tail-call %ktail 123))
       |]
+  , testCase "identity continuation" do
+      evalsTo [ObjImm (ImmInt 45)] [stkP|
+        (define ($start %ktail)
+          (push! %ktail)
+          (tail-call $id 45))
+        (define ($id %x)
+          (pop! %ktail)
+          (tail-call %ktail %x))
+      |]
   , testCase "identity function" do
       evalsTo [ObjImm (ImmInt 45)] [stkP|
-        (define ($main)
-          (tail-call $id 45))
-        (define ($id %x)
-          (pop-cont! %ktail)
+        (define ($start %ktail)
+          (tail-call $id 45 %ktail))
+        (define ($id %x %ktail)
           (tail-call %ktail %x))
       |]
-  -- , testCase "square" do
-  --     evalsTo [ObjImm (ImmInt 16)] [stkP|
-  --       (define ($main))
-  --     |]
   , testCase "square" do
       evalsTo [ObjImm (ImmInt 16)] [stkP|
-        (define ($main)
-          (tail-call $square 4))
-        (define ($square %x)
+        (define ($start %ktail)
+          (tail-call $square 4 %ktail))
+        (define ($square %x %ktail)
           (prim %x2 (* %x %x))
-          (pop-cont! %ktail)
           (tail-call %ktail %x2))
       |]
   , testCase "factorial" do
       let hsfac (n :: Int) = foldr (*) (1) [1..n]
       let fac (n :: Int) = [stkP|
-            (define ($fac %n)
+            (define ($fac %n %ktail)
               (prim %x0 (zero? %n))
               (if %x0
-                (then (pop-cont! %ktail)
-                      (tail-call %ktail 1))
+                (then (tail-call %ktail 1))
                 (else (push! %n)
+                      (push! %ktail)
                       (prim %x1 (- %n 1))
-                      (push-cont! $fac-k0)
-                      (tail-call $fac %x1))))
+                      (tail-call $fac %x1 $fac-k0))))
             (define ($fac-k0 %x2)
+              (pop! %ktail)
               (pop! %n)
               (prim %x3 (* %x2 %n))
-              (pop-cont! %ktail)
               (tail-call %ktail %x3))
-            (define ($main)
-              (tail-call $fac #{n}))
+            (define ($start %ktail)
+              (tail-call $fac #{n} %ktail))
           |]
       evalsTo [ObjImm (ImmInt 1)] $ fac 0
       evalsTo [ObjImm (ImmInt 1)] $ fac 1