correctly handle push-calls!!!!!!!!

src/Gyehoek/CPS/Stackify.hs

@@ -22,6 +22,7 @@ import Data.HashMap.Strict (HashMap)
 import qualified Data.HashMap.Strict as H
 import Data.HashSet.Lens (hashMap)
 import Data.List (List)
+import GHC.Exts (IsList(fromList))
 
 
 type Stackify = Writer Stk.Program
@@ -29,26 +30,30 @@ 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 && x /= g.returnLabel
+
 stackify
   :: (GenSym :> es, Stackify :> es)
   => Env -> Exp -> Eff es (Seq Stk.Instr)
 
 stackify g (ExpLetRec [(f, kap@(AbsKappa' xs m))] e) = do
   let vs = (f, Stk.ValLabel f) : (bindReg <$> xs)
-  let live = free' kap & filter (`H.member` g.bound)
-  m' <- stackify (g & #bound .~ H.fromList (vs ++ (bindReg <$> live))) m
+  let ls = live g kap
+  m' <- stackify (g & #bound .~ H.fromList (vs ++ (bindReg <$> ls))) m
   tell [Stk.MkBlock f xs $
-        [Stk.Pop x | x <- live] <> toList m']
-  stackify (g & #bound . at f ?~ Stk.ValLabel f) e
+        [Stk.Pop x | x <- ls] <> toList 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)
   lam_body <- gensym' "lambda-body"
   m' <- stackify (g & #bound .~ H.fromList vs
-                  & #bound . at f ?~ Stk.ValLabel lam_body) m
-  tell [Stk.MkBlock lam_body xs . toList $
-       -- this is probably evil and wrong.
-        [Stk.PopCont k] <> m']
+                    & #bound . at f ?~ Stk.ValLabel lam_body
+                    & #returnLabel .~ k) m
+  tell [Stk.MkBlock lam_body xs . toList $ m']
   stackify (g & #bound . at f ?~ Stk.ValLabel lam_body) e
 
 stackify g (ExpIf c t f) = do
@@ -57,13 +62,24 @@ stackify g (ExpIf c t f) = do
   pure [ Stk.If (stackifyVal g c) (toList t') (toList f') ]
 
 stackify g (ExpApply f xs ktail) = do
-  pure [ Stk.PushCont (var g ktail)
-       , Stk.Call (stackifyVal g f) (stackifyVal g <$> xs)
-       ]
+  pure $
+    [ Stk.PushCont (Stk.ValLabel k) ]
+    <> fromList [ Stk.Push (Stk.ValReg l) | l <- ls ]
+    <> [ Stk.Call (stackifyVal g f) (stackifyVal g <$> xs) ]
+  where
+    k = case var g ktail of
+      Stk.ValLabel x -> x
+      x -> error [i|expected a label, got #{x} (i guess)|]
+    ls = fold $ g ^. #liveness . at k
 
-stackify g (ExpContinue k xs) = pure
-  [ Stk.Call (var g k) (stackifyVal g <$> xs)
-  ]
+-- this probably won't work for call/cc, for cps-converted code it'll
+-- be fine i think. notice how, instead of calling `var g k`, we just
+-- assume it's the return continuation on top of the stack.
+stackify g (ExpContinue k xs) = do
+  ktail <- gensym' $ k ^. _Wrapped'
+  pure [ Stk.PopCont ktail
+       , Stk.Call (Stk.ValReg ktail) (stackifyVal g <$> xs)
+       ]
 
 stackify g (ExpPrim p (MkKappa [x] e)) = do
   e' <- stackify (g & #bound . at x ?~ Stk.ValReg x) e
@@ -90,11 +106,16 @@ bindReg x = (x, Stk.ValReg x)
 
 data Env = MkEnv
   { bound :: HashMap Name Stk.Val
+  , returnLabel :: Name
+  -- | 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
+emptyEnv = MkEnv mempty "halt" mempty
 
 
 
@@ -105,3 +126,23 @@ stackifyExp lbl e = do
 
 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))
+|]

src/Gyehoek/CPS/Syntax.hs

@@ -231,6 +231,7 @@ instance CPS Val where toCPS = Gyehoek.Sexp.fromSexp
 instance CPS Kappa where toCPS = Gyehoek.Sexp.fromSexp
 instance CPS Lambda where toCPS = Gyehoek.Sexp.fromSexp
 instance CPS Abs where toCPS = Gyehoek.Sexp.fromSexp
+instance CPS Program where toCPS = Gyehoek.Sexp.fromSexp
 
 cps :: QuasiQuoter
 cps = Gyehoek.Sexp.makeSx' [| toCPS |]

test/Gyehoek/Test/CPS/Stackify.hs

@@ -63,7 +63,7 @@ condition = testCase "if" do
     [cps|(if #f (continue halt 123) (continue halt 456))|]
 
 procedure = testGroup "procedure"
-  [ expectFail $ testCase "factorial" do
+  [ testCase "factorial" do
       evalsTo [ObjImm (ImmInt 720)]
         [cps|(letrec ((fac (λ (n ktail)
                              (prim (zero? n)