src/Gyehoek/CPS/Close.hs

@@ -9,7 +9,7 @@ import Gyehoek.GenSym
 import Gyehoek.Prelude
 
 
-close :: GenSym :> es => Exp -> Eff es Exp
+close :: forall es. GenSym :> es => Exp -> Eff es Exp
 close = transformM \case
   ExpLetRec [(f, AbsLambda lam@(MkLambda bs kb m))] e -> do
     f_code <- gensym' @Name $ f ^. _Wrapped' . to (<> "-code")
@@ -17,19 +17,23 @@ close = transformM \case
     -- but we're reusing the lambda binding so we don't have to
     -- explicitly substitute recursive calls.
     let frees = nub $ free' lam
-    let m' = ifoldr
-             (\n x q ->
+    m' <- ifoldrM @_ @_ @(Eff es)
+             (\n x q -> do
+               q_l <- gensym' @Name "env-cont"
                let p = if x == f then PrimEnv @Val else PrimEnvRef n
-               in [cps|
-                 (prim #{p}
-                   (κ (#{x}) #{q}))
-             |])
+               pure [cps|
+                 (letrec ((#{q_l} (κ (#{x}) #{q})))
+                   (prim #{p}
+                     #{q_l}))
+               |])
              m frees
+    e_l <- gensym' @Name "make-closure-cont"
     pure [cps|
       (letrec ((#{f_code} (λ (##{bs} #{kb})
                             #{m'})))
-        (prim (make-closure #{f_code} ##{frees})
-          (κ (#{f}) #{e})))
+        (letrec ((#{e_l} (κ (#{f}) #{e})))
+          (prim (make-closure #{f_code} ##{frees})
+            #{e_l})))
     |]
 
   e -> pure e

src/Gyehoek/CPS/Convert.hs

@@ -46,29 +46,15 @@ convert (Scm.ExpLit l) k = k . one . ValImm $ case l of
   LitBool b -> ImmBool b
   _ -> _
 
--- special case: call/cc is desugared during cps-conversion...
--- convert (Scm.ExpPrim (PrimCallCC withcc)) k = do
---   convert1 withcc \withcc' -> do
---     cc_l <- gensym' @Name "cc"
---     r1_l <- gensym' @Name "r"
---     r2_l <- gensym' @Name "r"
---     ccish_l <- gensym' @Name "ccish"
---     reified_cc_l <- gensym' @Name "reified-cc"
---     m <- k [ValVar r1_l]
---     pure [cps|
--- (letrec ((#{cc_l} (κ (#{r1_l})
---                      #{m})))
---   (prim (capture/cc)
---        (κ (#{reified_cc_l})
---           (#{withcc'} #{reified_cc_l} #{cc_l}))))
---     |]
-
--- ...while all other prims are left as-is for later stages to
--- handle..
 convert (Scm.ExpPrim p) k =
   telescope (convert1 @es) p \p' -> do
-    r <- gensym' "r"
-    ExpPrim p' . MkKappa [r] <$> k [ValVar r]
+    r_l <- gensym' "r"
+    k_l <- gensym' @Name "prim-k"
+    m <- k [ValVar r_l]
+    pure [cps|
+      (letrec ((#{k_l} (κ (#{r_l}) #{m})))
+        (prim #{p'} #{k_l}))
+    |]
 
 convert (Scm.ExpLambda xs e) k = do
   f <- gensym' "lambda-body"
@@ -90,8 +76,16 @@ convert (Scm.ExpApply f xs) k =
 convert (Scm.ExpBegin xs) k = telescope (convert @es) xs (k . NE.last)
 
 convert (Scm.ExpIf c t f) k =
-  convert1 c \c' ->
-    ExpIf c' <$> convert t k <*> convert f k
+  convert1 c \c' -> do
+    t_l <- gensym' @Name "truthy-cont"
+    f_l <- gensym' @Name "falsey-cont"
+    t' <- convert t k
+    f' <- convert f k
+    pure [cps|
+      (letrec ((#{t_l} (κ () #{t'}))
+               (#{f_l} (κ () #{f'})))
+        (if #{c'} #{t_l} #{f_l}))
+    |]
 
 -- let-bindings are desugared into continuation calls whose parameters
 -- are the left-hand sides and whose arguments are the right-hand

src/Gyehoek/CPS/Stackify.hs

@@ -12,7 +12,7 @@ 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.List (elemIndex, nub, intersect)
 import Data.Text qualified as T
 import Gyehoek.Prelude
 import Debug.Pretty.Simple
@@ -30,6 +30,16 @@ live g e = nub (free' e) & filter \x ->
   x `elem` g.bound
   -- && not (x `elem` g.contStack)
 
+-- | The expression @load g e r n@ emits a 'Stk.Load' instruction if
+-- stack variable @n@ is live-out in expression @e@. Otherwise, a
+-- 'Stk.Pop' instruction is emitted.
+load :: Free a => Env -> a -> Reg -> Int -> Stk.Instr
+load g e r 0
+  | Just x <- g ^? #bound . _head
+  , x `elem` free e
+  = Stk.Pop r
+load g e r n = Stk.Load r n
+
 data BlockBuilder
   = Code (List Stk.Instr) BlockBuilder
   | Tail Stk.Tail
@@ -90,14 +100,14 @@ stackify g (ExpPrim (PrimCallCC withcc) cc) = do
          ] $
     Tail Stk.CallCC
 
-stackify g (ExpPrim p kap) = do
-  kap' <- stackifyKappa g kap
-  pure $ Code [ Stk.Prim (stackifyVal g <$> p) ] kap'
+stackify g (ExpPrim p (MkKappa rs e)) = do
+  _
+  pure $ Code [ Stk.Prim (stackifyVal g <$> p) ] _
 
 stackify _ e = error [i|unimplemented exp: #{e}|]
 
-loadArgs :: List Name -> List Stk.Instr
-loadArgs = imapOf itraversed \n x -> Stk.Load (MkReg x) n
+loadArgs :: Free a => Env -> a -> List Name -> List Stk.Instr
+loadArgs g e = imapOf itraversed \n x -> load g e (MkReg x) n
 
 pushArgs :: Env -> List Val -> List Stk.Instr
 pushArgs g args = [ Stk.Push $ stackifyVal g x | x <- reverse args ]
@@ -112,7 +122,7 @@ stackifyKappa
   -> Eff es BlockBuilder
 stackifyKappa g (MkKappa xs m) = do
   let g' = g & #bound <>:~ xs
-  Code (loadArgs g'.bound)
+  Code [ _ | x <- g'.bound `intersect` free' m ]
     <$> stackify g' m
 
 stackifyLambda

src/Gyehoek/CPS/Syntax.hs

@@ -56,6 +56,7 @@ import Gyehoek.Sexp (G, (:-)(..))
 import qualified Data.InvertibleGrammar.Base as IG
 import Gyehoek.GenSym (Gen)
 import Data.String (IsString)
+import Control.Applicative
 
 -- Data types
 
@@ -113,17 +114,17 @@ data Abs
   | AbsLambda Lambda
   deriving (Show, Generic, Data, Eq)
 
-pattern AbsKappa' :: [Name] -> Exp -> Abs
+pattern AbsKappa' :: List Name -> Exp -> Abs
 pattern AbsKappa' xs e = AbsKappa (MkKappa xs e)
 
-pattern AbsLambda' :: [Name] -> Name -> Exp -> Abs
+pattern AbsLambda' :: List Name -> Name -> Exp -> Abs
 pattern AbsLambda' xs e ktail = AbsLambda (MkLambda xs e ktail)
 
 data Exp
-  = ExpPrim (Prim Val) Kappa
+  = ExpPrim (Prim Val) Name
   | ExpLetRec { binders :: List (Name, Abs), body :: Exp }
   | ExpContinue Val (List Val)
-  | ExpIf Val Exp Exp
+  | ExpIf Val Name Name
   | ExpApply
     { op :: Val
     , args :: List Val
@@ -329,7 +330,8 @@ class Free a where
   freeWithBound :: HashSet Name -> a -> HashSet Name
   freeWithBound bound = HS.fromList . freeWithBound' bound
 
-  -- | Free variables given in the order of their appearance.
+  -- | Free variables given in the same left-to-right order they
+  -- appear.
   free' :: a -> List Name
   free' = freeWithBound' mempty
 
@@ -339,11 +341,16 @@ instance Free Abs where
   freeWithBound' bound (AbsKappa kap) = freeWithBound' bound kap
   freeWithBound' bound (AbsLambda lam) = freeWithBound' bound lam
 
+mif :: Alternative f => (a -> Bool) -> a -> f a
+mif p a
+  | p a = pure a
+  | otherwise = empty
+
 instance Free Exp where
   freeWithBound' bound = \case
     ExpPrim p k ->
-      p & toListOf (folded . #ValVar . filtered (`notElem` bound))
-        & (<> freeWithBound' bound k)
+      (p ^.. folded . #ValVar . filtered (`notElem` bound))
+      ++ mif (`notElem` bound) k
     ExpLetRec bs m ->
       foldMapOf (each . _2) (freeWithBound' bound') bs
       <> freeWithBound' bound' m
@@ -351,10 +358,10 @@ instance Free Exp where
     ExpContinue k xs -> filter (`notElem` bound) ((k:xs) ^.. each . #ValVar)
     ExpIf c t f ->
       (c ^.. #ValVar . filtered (`notElem` bound))
-      <> freeWithBound' bound t <> freeWithBound' bound f
+      <> mif (`notElem` bound) t <> mif (`notElem` bound) f
     ExpApply f xs k ->
       (f:xs) ^.. (each . #ValVar . filtered (`notElem` bound))
-      <> (k ^.. filtered (`notElem` bound))
+      <> mif (`notElem` bound) k
 
 instance Free Kappa where
   freeWithBound' bound (MkKappa xs m) =