shared closures maybe

.dir-locals.el

@@ -9,6 +9,9 @@
       . (progn (defun apply-cabal-fmt-h ()
                  (haskell-mode-buffer-apply-command "cabal-fmt"))
                (add-hook 'before-save-hook #'apply-cabal-fmt-h nil t)))))
+ (scheme-mode
+  . ((eval . (dolist (s '(kappa κ prim))
+               (put s 'scheme-indent-function 1)))))
  (nil
   . ((eval
       . (progn (defun display-ansi ()

doc/closure-conversion.org

@@ -132,3 +132,34 @@ multiple ~env-ref~ calls could probably be replaced with a primitive that loads
                                  $code)
                      1))))
 #+end_src
+
+** example
+
+#+begin_src scheme
+  (λ (n m ktail)
+    (letrec ((f (λ (x ktail-0) (+ x n ktail-0)))
+             (g (λ (y ktail-1) (+ y g ktail-1))))
+      (prim (cons f g) ktail)))
+#+end_src
+
+#+begin_src scheme
+  (λ (n m ktail)
+    (letrec ((f-code (λ (x ktail-0)
+                       (prim (env-get 2)
+                         (κ (n)
+                           (+ x n ktail-0)))))
+             (g-code (λ (y ktail-1)
+                       (prim (env-get 3)
+                         (κ (m)
+                           (+ y m ktail-1))))))
+      (letrec ((with-closure-code
+                (κ (f g)
+                  (prim (get-env 0)
+                    (κ (ktail)
+                      (prim cons f g ktail))))))
+        (prim (make-shared-closure (with-closure-code)
+                                   ktail)
+          (κ (with-closure)
+            (prim (make-shared-closure (f-code g-code) n m)
+              with-closure))))))
+#+end_src

gyehoek.cabal

@@ -58,14 +58,16 @@ library
   -- cabal-fmt: expand src
   exposed-modules:
     Gyehoek.CPS.Close
+    Gyehoek.CPS.Contify
     Gyehoek.CPS.Convert
     Gyehoek.CPS.Eval
+    Gyehoek.CPS.Hoist
     Gyehoek.CPS.Stackify
     Gyehoek.CPS.Syntax
     Gyehoek.Driver
-    Gyehoek.Language
     Gyehoek.GenSym
     Gyehoek.Jalmot
+    Gyehoek.Language
     Gyehoek.Lift1
     Gyehoek.Options
     Gyehoek.Prelude
@@ -100,6 +102,7 @@ library
     , hashable
     , invertible-grammar
     , lens
+    , lucid
     , megaparsec
     , mtl
     , optparse-applicative
@@ -107,6 +110,7 @@ library
     , pretty-simple
     , prettyprinter
     , prettyprinter-ansi-terminal
+    , prettyprinter-lucid
     , process
     , recursion-schemes
     , scientific
@@ -117,8 +121,7 @@ library
     , typed-process
     , unordered-containers
     , vector
-    , lucid
+    , tardis
-    , prettyprinter-lucid
 
   hs-source-dirs:   src
   default-language: GHC2024
@@ -171,6 +174,7 @@ test-suite doctest
   build-depends:
     , base
     , gyehoek
+
   default-extensions: CPP
   main-is:            doctest.hs
 

src/Gyehoek/CPS/Close.hs

@@ -7,40 +7,86 @@ import Gyehoek.CPS.Syntax
 import Data.List (nub)
 import Gyehoek.GenSym
 import Gyehoek.Prelude
+import Debug.Pretty.Simple
+import Gyehoek.Sexp qualified as S
+import Data.HashSet.Lens
+import Data.Traversable
 
 
-close :: GenSym :> es => Exp -> Eff es Exp
+genCodeName :: GenSym :> es => Name -> Eff es Name
-close = transformM \case
+genCodeName f = gensym' @Name $ f ^. _Wrapped' . to (<> "-code")
-  ExpLetRec [(f, AbsLambda lam@(MkLambda bs kb m))] e -> do
-    f_code <- gensym' @Name $ f ^. _Wrapped' . to (<> "-code")
-    -- it would probably be most sane to generate a symbol for `env`,
-    -- 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 ->
-               let p = if x == f then PrimEnv @Val else PrimEnvRef n
-               in [cps|
-                 (prim #{p}
-                   (κ (#{x}) #{q}))
-             |])
-             m frees
-    pure [cps|
-      (letrec ((#{f_code} (λ (##{bs} #{kb})
-                            #{m'})))
-        (prim (make-closure #{f_code} ##{frees})
-          (κ (#{f}) #{e})))
-    |]
 
-  -- ExpApply f xs ktail -> do
+bindEnv :: List Name -> Exp -> Exp
-  --   code <- gensym' @Name "code"
+bindEnv frees m = [cps|
-  --   pure [cps|
+  (prim (get-env) (κ #{frees} #{m}))
-  --     (prim (env-code #{f})
+|]
-  --       (κ (#{code})
+
-  --         (#{code} #{f} ##{xs} #{ktail})))
+close1 :: forall es. GenSym :> es => Exp -> Eff es Exp
-  --   |]
+close1 = \case
+  lr@(ExpLetRec bs e) -> do
+    let boundNames = bs ^.. each . _1
+    let boundNames' = setOf each boundNames
+    let frees = bs
+          & foldMapOf
+            (each . _2)
+            (freeWithBound' boundNames')
+          & nub
+    pTraceShowM frees
+    env_cont_l <- gensym' @Name "env-cont"
+    e_l <- gensym' @Name "letrec-body-cont"
+    bs' <- for bs \(f,ab) -> do
+      f_code_l <- genCodeName f
+      pure ( f_code_l
+           , ab & absBody %~ bindEnv (boundNames ++ frees)
+           )
+    let codes = bs' ^.. each . _1 . to MkLabel
+    pure [cps|
+      (letrec #{bs'}
+        (prim (make-shared-closure #{codes} #{frees})
+          (κ #{boundNames}
+            #{e})))
+    |]
 
   e -> pure e
 
+close :: forall es. GenSym :> es => Exp -> Eff es Exp
+close = transformM close1
+
 closeProgram :: GenSym :> es => Program -> Eff es Program
 closeProgram = traverseOf (#body . #body) close
+
+prog :: Program
+prog = [cps|
+(λ (start-ktail0)
+  (letrec ((fac
+            (λ (n lambda-tail1)
+              (letrec ((prim-k3
+                        (κ (r2)
+                          (letrec ((truthy-cont4
+                                    (κ ()
+                                      (continue lambda-tail1 1)))
+                                   (falsey-cont5
+                                    (κ ()
+                                      (letrec ((prim-k7
+                                                (κ (r6)
+                                                  (letrec
+                                                      ((r8
+                                                        (κ (x9)
+                                                          (letrec
+                                                              ((prim-k11
+                                                                (κ (r10)
+                                                                  (continue
+                                                                   lambda-tail1
+                                                                   r10))))
+                                                            (prim
+                                                             (* n x9)
+                                                             prim-k11)))))
+                                                    (fac r6 r8)))))
+                                        (prim (- n 1) prim-k7)))))
+                            (if r2
+                              truthy-cont4
+                              falsey-cont5)))))
+                (prim (zero? n) prim-k3)))))
+    (letrec ((r12 (κ (x13) (continue start-ktail0 x13))))
+      (fac 20 r12))))
+|]

src/Gyehoek/CPS/Contify.hs

@@ -0,0 +1,67 @@
+{-# LANGUAGE ApplicativeDo #-}
+module Gyehoek.CPS.Contify
+  ( contifyProgram
+  ) where
+
+import Control.Monad.Tardis
+import Gyehoek.CPS.Syntax
+import Gyehoek.Prelude
+import qualified Data.HashSet as HS
+import Control.Lens.Unsound (adjoin)
+import Debug.Pretty.Simple
+import qualified Data.HashMap.Strict as H
+import Control.Monad.Writer.Lazy
+import Control.Monad.Trans.Tardis (liftTardisT)
+
+
+-- | ain't no way...
+-- type T = WriterT (HashSet Name) (Tardis (HashSet Name) (HashSet Name))
+type T = TardisT (HashSet Name) (HashSet Name) (Writer (HashSet Name))
+
+evalT :: T a -> a
+-- evalT = (`evalTardis` (mempty,mempty)) . fmap fst . runWriterT
+evalT = fst . runWriter . (`evalTardisT` (mempty,mempty))
+
+runT :: T a -> (a, HashSet Name)
+-- runT = (`evalTardis` (mempty,mempty)) . runWriterT
+runT = runWriter . (`evalTardisT` (mempty,mempty))
+
+-- | inline function if it hasn't been used in the past, and won't
+-- be used in the future.
+tryInline :: Name -> Kappa -> T Kexp
+tryInline kname kap = do
+  modifyBackwards (HS.insert kname)
+  p <- getsPast (HS.member kname)
+  modifyForwards (HS.insert kname)
+  q <- getsFuture (HS.member kname)
+  let c = p || q
+  liftTardisT . tell $ if c then HS.singleton kname else mempty
+  pure $ if c
+    then KexpVar kname
+    else KexpKappa kap
+
+getKap :: HashMap Name Abs -> Name -> Maybe Kappa
+getKap g kname = g ^? ix kname . #AbsKappa
+
+contify :: HashMap Name Abs -> Exp -> T Exp
+contify g = transformM \case
+  ExpApply f xs (KexpVar kname) | Just kap <- getKap g kname
+   -> ExpApply f xs <$> tryInline kname kap
+  ExpPrim p (KexpVar kname) | Just kap <- getKap g kname
+   -> ExpPrim p <$> tryInline kname kap
+  e -> pure e
+
+contifyProgram :: HoistedProgram -> Eff es HoistedProgram
+contifyProgram p = do
+  let g = p.bindings
+  let (p',contifiedVars) =
+        runT $
+          traverseOf
+            (adjoin
+              (#bindings . each . body)
+              (#body . body))
+            (contify g)
+            p
+  pTraceShowM contifiedVars
+  -- pure $ p' & #bindings %~ H.filterWithKey \k _ -> HS.member k contifiedVars
+  pure p'

src/Gyehoek/CPS/Convert.hs

@@ -46,26 +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 <- gensym' @Name "cc"
---     r <- gensym' "r"
---     m <- k . one $ ValVar r
---     ccish <- gensym' @Name "cc-ish"
---     x <- gensym' @Name "x"
---     pure [cps|
---       (letrec ((#{cc} (κ (#{r}) #{m})))
---         (letrec ((#{ccish} (λ (#{x} _) (continue #{cc} #{x}))))
---           (#{withcc'} #{ccish} #{cc})))
---     |]
-
--- ...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"
+    r_l <- gensym' "r"
-    ExpPrim p' . MkKappa [r] <$> k [ValVar 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"
@@ -78,17 +67,25 @@ convert (Scm.ExpLambda xs e) k = do
 
 convert (Scm.ExpApply f xs) k =
   telescope (convert1 @es) (f:|xs) \(f':|xs') -> do
-    r <- gensym' "r"
+    r <- gensym' @Name "r"
     x <- gensym' "x"
     m <- k [ValVar x]
     pure $ ExpLetRec [(r, AbsKappa' [x] m)] $
-      ExpApply f' xs' r
+      ExpApply f' xs' (KexpVar r)
 
 convert (Scm.ExpBegin xs) k = telescope (convert @es) xs (k . NE.last)
 
 convert (Scm.ExpIf c t f) k =
-  convert1 c \c' ->
+  convert1 c \c' -> do
-    ExpIf c' <$> convert t k <*> convert f k
+    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/Eval.hs

@@ -59,11 +59,6 @@ eval g (ExpPrim p (MkKappa bs e)) = case evalVal g <$> p of
       lbl = case x of
         ObjImm (ImmLabel l) -> l
         _ -> error [i|expected label, got #{x}|]
-  PrimEnvCode x -> ret . (:[]) . ObjImm . ImmLabel $ code
-    where
-      code = case x of
-        ObjHob (HobClosure lbl _) -> lbl
-        _ -> error [i|expected closure, got #{x}|]
   _ -> error [i|unhandled prim: #{p}|]
   where
     ret rs = eval

src/Gyehoek/CPS/Hoist.hs

@@ -0,0 +1,24 @@
+module Gyehoek.CPS.Hoist
+  ( hoistProgram
+  ) where
+
+import Gyehoek.CPS.Syntax
+import Gyehoek.Prelude
+import qualified Data.HashMap.Strict as H
+import Effectful.Writer.Static.Local
+import Data.Foldable
+
+
+type Hoist = Writer (HashMap Name Abs)
+
+hoist :: Hoist :> es => Exp -> Eff es Exp
+hoist = transformM \case
+  ExpLetRec bs m -> do
+    traverse_ (\(k,v) -> tell $ H.singleton k v) bs
+    pure m
+  e -> pure e
+
+hoistProgram :: Program -> Eff es HoistedProgram
+hoistProgram p = do
+  (body,bindings) <- runWriter $ traverseOf #body hoist p.body
+  pure $ MkHoistedProgram {body,bindings}

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
@@ -44,34 +54,32 @@ emitRoutine :: Stackify :> es => Stk.Routine -> Eff es ()
 emitRoutine rt = tell [rt]
 
 stackify
-  :: (GenSym :> es, Stackify :> es)
+  :: forall es. (GenSym :> es, Stackify :> es)
   => Env -> Exp -> Eff es BlockBuilder
 
-stackify g (ExpLetRec [(f, AbsKappa kap)] e) = do
+stackify g (ExpLetRec bs e) = do
-  kap' <- stackifyKappa g kap
+  for_ bs \(f,a) ->
-  emitRoutine (Stk.MkRoutine (MkLabel f) . buildBlock $ kap')
+    emitRoutine =<< case a of
-  stackify g e
+      AbsKappa kap -> stackifyKappa g (MkLabel f) kap
-
+      AbsLambda lam -> stackifyLambda g (MkLabel f) lam
-stackify g (ExpLetRec [(f, AbsLambda lam)] e) = do
-  lam' <- stackifyLambda g (MkLabel f) lam
-  emitRoutine lam'
   stackify g e
 
 stackify g (ExpIf c t f) = do
   let c' = stackifyVal g c
-  t' <- buildBlock <$> stackify g t
+  let jump l =
-  f' <- buildBlock <$> stackify g f
+       Stk.MkBlock
-  pure . Tail $ Stk.If c' t' f'
+         [Stk.Push . Stk.ValLabel . MkLabel $ l]
+         (Stk.TailCall 0)
+  pure . Tail $ Stk.If c' (jump t) (jump f)
 
 stackify g (ExpApply f xs ktail) = do
   pure $
-    Code [ Stk.Push $ stackifyVal g (ValVar ktail)
+    Code [ Stk.Push $ stackifyVal g (ValVar $ ktail ^?! #KexpVar)
          , Stk.Push $ stackifyVal g f
          ] $
     Code (pushArgs g xs) $
     Tail (Stk.Call (length xs))
 
--- assume that `k` is the continuation on top of the stack lol.
 stackify g e@(ExpContinue k xs)
   | isn't (#_ValVar . only g.tail) k = pure $
     Code [ Stk.Push (stackifyVal g k) ] $
@@ -81,9 +89,20 @@ stackify g e@(ExpContinue k xs)
     Code (pushArgs g xs) $
     Tail (Stk.Return (length xs))
 
-stackify g (ExpPrim p kap) = do
+stackify g (ExpPrim (PrimCallCC withcc) cc) = do
-  kap' <- stackifyKappa g kap
+  let cc' = cc ^?! #KexpVar . to MkLabel
-  pure $ Code [ Stk.Prim (stackifyVal g <$> p) ] kap'
+  cc_l <- gensym' @Label "cc"
+  pure $
+    Code [ Stk.Push $ stackifyVal g withcc
+         , Stk.Push $ stackifyVal g (ValLabel cc')
+         ] $
+    Tail Stk.CallCC
+
+stackify g (ExpPrim p cc) = pure $
+  Code [ Stk.Push (Stk.ValLabel . MkLabel $ cc ^?! #KexpVar)
+       , Stk.Prim (stackifyVal g <$> p)
+       ] $
+  Tail $ Stk.TailCall 1
 
 stackify _ e = error [i|unimplemented exp: #{e}|]
 
@@ -99,12 +118,20 @@ _ValName = failing #_ValVar (#_ValImm . #_ImmLabel . #_MkLabel)
 
 stackifyKappa
   :: (Stackify :> es, GenSym :> es)
-  => Env -> Kappa
+  => Env -> Label -> Kappa
-  -> Eff es BlockBuilder
+  -> Eff es Stk.Routine
-stackifyKappa g (MkKappa xs m) = do
+stackifyKappa g kname (MkKappa xs m) = do
   let g' = g & #bound <>:~ xs
-  Code (loadArgs g'.bound)
+  m' <- stackify g' m
-    <$> stackify g' m
+  pure $
+    Stk.MkRoutine kname . buildBlock $
+    Code (loadArgs xs) $
+    Code [ Stk.Load (MkReg r) j
+         | v <- g ^.. #liveness . ix kname . each
+         , (j,r) <- itoListOf (#bound . itraversed) g
+         , r == v
+         ] $
+    m'
 
 stackifyLambda
   :: (Stackify :> es, GenSym :> es)
@@ -115,7 +142,8 @@ stackifyLambda g name (MkLambda xs k m) = do
   pure $
     Stk.MkRoutine name . buildBlock $
     Code (loadArgs xs) $
-    Code [Stk.Load (MkReg k) (length xs + 1)] m'
+    -- Code [Stk.Load (MkReg k) (length xs + 1)] $
+    m'
 
 stackifyVal :: Env -> Val -> Stk.Val
 stackifyVal g = \case
@@ -152,40 +180,101 @@ emptyEnv = MkEnv
 
 
 
-stackifyProgram :: GenSym :> es => Program -> Eff es Stk.Program
+stackifyProgram :: GenSym :> es => HoistedProgram -> Eff es Stk.Program
-stackifyProgram (MkProgram lam) = do
+stackifyProgram p = do
-  let g = emptyEnv
+  let liveness = p & foldMapOf
-  (_,p) <- runStackify $ emitRoutine =<< stackifyLambda g "start" lam
+        (#bindings . itraversed . withIndex . aside #AbsKappa)
-  pure p
+        \(kname,kap) -> H.singleton
+          (MkLabel kname)
+          (nub $ freeWithBound' (H.keysSet p.bindings) kap)
+  let g = MkEnv
+        { bound = mempty
+        , liveness
+        , tail = p.body.ktail }
+  let e = p.body & #body %~ ExpLetRec (H.toList p.bindings)
+  (_,p') <- runStackify $ emitRoutine =<< stackifyLambda g "start" e
+  pure p'
 
-letfn :: Program
+blah :: HoistedProgram
-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))))
-|]
-
-blah :: Program
 blah = [cps|
-(λ (ktail0)
+(letrec ((prim-k3 (κ (r2) (if r2 truthy-cont4 falsey-cont5)))
-          (letrec ((fac (λ (n ktail)
+         (prim-k7 (κ (r6) (fac r6 r8)))
-                             (prim (zero? n)
+         (make-closure-cont15 (κ (fac) (fac 20 r12)))
-                               (κ (x0)
+         (falsey-cont5 (κ () (prim (- n 1) prim-k7)))
-                                 (if x0
+         (r12 (κ (x13) (continue start-ktail0 x13)))
-                                   (continue ktail 1)
+         (truthy-cont4 (κ () (continue lambda-tail1 1)))
-                                   (prim (- n 1)
+         (prim-k11 (κ (r10) (continue lambda-tail1 r10)))
-                                     (κ (x1)
+         (r8 (κ (x9) (prim (* n x9) prim-k11)))
-                                       (letrec ((fac-k0
+         (fac-code14 (λ (n lambda-tail1) (prim (zero? n) prim-k3))))
-                                                  (κ (x2)
+  (λ (start-ktail0)
-                                                    (prim (* n x2)
+    (prim (make-closure $fac-code14) make-closure-cont15)))
-                                                      (κ (x3)
+|]
-                                                        (continue ktail x3))))))
+
-                                         (fac x1 fac-k0))))))))))
+p :: HoistedProgram
-               (fac 6 halt)))
+p = [cps|
+(letrec ((r12-code32 (κ (r12 start-ktail0 x13) (continue start-ktail0 x13)))
+         (prim-k7-code22
+          (κ (prim-k7 fac r6 r8 n x9 prim-k11 lambda-tail1 r10)
+            (prim
+             (make-shared-closure (r8) (n x9 prim-k11 lambda-tail1 r10))
+             letrec-body-cont18)))
+         (letrec-body-cont24
+          (κ (truthy-cont4 falsey-cont5)
+            (if r2
+              truthy-cont4
+              falsey-cont5)))
+         (letrec-body-cont18 (κ (r8) (fac r6 r8)))
+         (prim-k11-code16
+          (κ (prim-k11 lambda-tail1 r10)
+            (continue lambda-tail1 r10)))
+         (falsey-cont5-code26
+          (κ (truthy-cont4 falsey-cont5 lambda-tail1 n prim-k7
+              fac r6 r8 x9 prim-k11 r10)
+            (prim
+             (make-shared-closure
+              (prim-k7)
+              (fac r6 r8 n x9 prim-k11 lambda-tail1 r10))
+             letrec-body-cont21)))
+         (letrec-body-cont31 (κ (r12) (fac 20 r12)))
+         (r8-code19
+          (κ (r8 n x9 prim-k11 lambda-tail1 r10)
+            (prim
+             (make-shared-closure (prim-k11) (lambda-tail1 r10))
+             letrec-body-cont15)))
+         (letrec-body-cont28 (κ (prim-k3) (prim (zero? n) prim-k3)))
+         (truthy-cont4-code25
+          (κ (truthy-cont4 falsey-cont5 lambda-tail1 n prim-k7 fac
+              r6 r8 x9 prim-k11 r10)
+            (continue lambda-tail1 1)))
+         (fac-code35
+          (κ (fac n prim-k3 r2 truthy-cont4 falsey-cont5 lambda-tail1
+              prim-k7 r6 r8 x9 prim-k11 r10)
+            (prim
+             (make-shared-closure
+              (prim-k3)
+              (r2 truthy-cont4 falsey-cont5 lambda-tail1 n prim-k7
+               fac r6 r8 x9 prim-k11 r10))
+             letrec-body-cont28)))
+         (prim-k3-code29
+          (κ (prim-k3 r2 truthy-cont4 falsey-cont5 lambda-tail1 n prim-k7
+              fac r6 r8 x9 prim-k11 r10)
+            (prim
+             (make-shared-closure
+              (truthy-cont4 falsey-cont5)
+              (lambda-tail1 n prim-k7 fac r6 r8 x9 prim-k11 r10))
+             letrec-body-cont24)))
+         (letrec-body-cont21 (κ (prim-k7) (prim (- n 1) prim-k7)))
+         (letrec-body-cont15 (κ (prim-k11) (prim (* n x9) prim-k11)))
+         (letrec-body-cont34
+          (κ (fac)
+            (prim
+             (make-shared-closure (r12) (start-ktail0 x13))
+             letrec-body-cont31))))
+  (λ (start-ktail0)
+    (prim
+     (make-shared-closure
+      (fac)
+      (n prim-k3 r2 truthy-cont4 falsey-cont5 lambda-tail1
+       prim-k7 r6 r8 x9 prim-k11 r10))
+     letrec-body-cont34)))
 |]

src/Gyehoek/CPS/Syntax.hs

@@ -10,10 +10,12 @@ module Gyehoek.CPS.Syntax
   , Kappa(..)
   , Lambda(..)
   , Exp(..)
+  , Kexp(..)
   , ExpF(..)
   , Name(..)
   , Prim(..)
   , Program(..)
+  , HoistedProgram(..)
   , Lit(..)
   , Imm(..)
   , Obj(..)
@@ -39,6 +41,7 @@ module Gyehoek.CPS.Syntax
   , Free(..)
   , pattern ValLabel
   , pattern ObjLabel
+  , absBody
   )
   where
 
@@ -56,6 +59,9 @@ import Gyehoek.Sexp (G, (:-)(..))
 import qualified Data.InvertibleGrammar.Base as IG
 import Gyehoek.GenSym (Gen)
 import Data.String (IsString)
+import Control.Applicative
+import qualified Data.HashMap.Strict as H
+import GHC.Records (HasField (..))
 
 -- Data types
 
@@ -96,6 +102,8 @@ pattern ObjLabel l = ObjImm (ImmLabel l)
 -- | a heap object.
 data Hob
   = HobClosure { label :: Label, env :: List Obj }
+  -- should a continuation have a label, or an Obj?
+  | HobContinuation { cont :: Obj, stack :: NonEmpty (List Obj) }
   | HobPair Obj Obj
   deriving stock (Show, Generic, Data, Eq)
   deriving anyclass (NFData)
@@ -111,24 +119,30 @@ 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) Kexp
   | 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
-    , cont :: Name
+    , cont :: Kexp
     }
   deriving (Show, Generic, Data, Eq)
 
+data Kexp
+  = KexpVar Name
+  -- | Only to be used after contification.
+  | KexpKappa Kappa
+  deriving (Show, Generic, Data, Eq)
+
 pattern Halt :: List Val -> Exp
 pattern Halt xs = ExpContinue (ValLabel "halt") xs
 
@@ -143,6 +157,22 @@ data Program = MkProgram
   }
   deriving (Show, Generic, Data)
 
+data HoistedProgram = MkHoistedProgram
+  { bindings :: HashMap Name Abs
+  , body :: Lambda
+  }
+  deriving stock (Show, Generic, Data)
+
+type instance Index HoistedProgram = Name
+type instance IxValue HoistedProgram = Abs
+
+instance Ixed HoistedProgram where ix j = #bindings . ix j
+
+instance At HoistedProgram where at j = #bindings . at j
+
+instance Each HoistedProgram HoistedProgram Abs Abs where
+  each = #bindings . each
+
 makePrisms ''Kappa
 makePrisms ''Exp
 makeFieldsId ''Exp
@@ -166,6 +196,16 @@ _AbsLambda' = prism'
 
 instance Plated Exp where plate = uniplate
 
+
+absBody :: Lens' Abs Exp
+absBody = lens
+  (\case
+    AbsLambda lam -> lam.body
+    AbsKappa kap -> kap.body)
+  (\cases
+    (AbsLambda lam) b -> AbsLambda $ lam & #body .~ b
+    (AbsKappa kap) b -> AbsKappa $ kap & #body .~ b)
+
 
 -- DatumIso instances
 
@@ -208,6 +248,7 @@ instance S.DatumIso Reg where
 instance S.DatumIso Hob where
   datumIso = S.match
     $ S.With (. closure)
+    $ S.With (. cont)
     $ S.With (. conspair)
     $ S.End
     where
@@ -215,7 +256,13 @@ instance S.DatumIso Hob where
       -- closures can be printed, but not parsed.
       closure :: G (Datum :- t) (List Obj :- Label :- t)
       closure = IG.Flip $ IG.PartialIso
-        (\(env:-code:-t) -> S.Unreadable "#<procedure>" :- t)
+        (\(env:-code:-t) -> S.Unreadable [i|\#<procedure $#{code}>|] :- t)
+        (const . Left $ mempty)
+      cont :: G (Datum :- t) (NonEmpty (List Obj) :- _ :- t)
+      cont = IG.Flip $ IG.PartialIso
+        (\(_ :- l :- t) -> 
+          let x = S.encodeOrShow' @Text S.datumIso l
+          in S.Unreadable [i|\#<continuation #{x}>|] :- t)
         (const . Left $ mempty)
 
 instance S.DatumIso Lambda where
@@ -262,30 +309,51 @@ instance S.DatumIso Exp where
       if_ = S.ifLike "if"
         S.datumIso S.datumIso S.datumIso
       app :: forall t.
-        G (Datum :- t) (Name :- ([Val] :- (Val :- t)))
+        G (Datum :- t) (Kexp :- List Val :- Val :- t)
-      app = S.list $ S.el (S.datumIso @Val)
+      app = S.list $
-        -- >>> S.flipped Gyehoek.Datum.nonEmptyGrammar
+        S.flipped (S.PartialIso
+          (\(S.MkListContext ctx :- t) ->
+            case ctx of
+              f:kexp:xs -> S.MkListContext (f : snoc xs kexp) :- t
+              _ -> error "unreachable")
+          (\(S.MkListContext ctx :- t) ->
+            case unsnoc ctx of
+              Just (f:xs,kexp) -> Right $ S.MkListContext (f:kexp:xs) :- t
+              _ -> Left $ S.expected "continuation arg"))
+        >>> S.el (S.datumIso @Val)
+        >>> S.el (S.datumIso @Kexp)
         >>> S.rest (S.datumIso @Val)
-        -- >>> _
+        >>> S.onTail S.swap
-        >>> S.onTail (S.flipped $ IG.PartialIso
-                      (\(karg :- args :- op :- t) ->
-                         (args ++ [ValVar karg]) :- op :- t)
-                      (\(xs :- op :- t) -> case xs ^? _Snoc of
-                          Just (args,preview #ValVar -> Just karg) ->
-                            Right $ karg:- args :- op :- t
-                          _ -> Left $ S.expected "continuation arg"
-                      ))
-      -- prim = S.headTagged2 "prim"
-      --   (primDatumIso id (S.datumIso @Val))
-      --   (S.datumIso @Kappa)
       prim = S.list $
         S.el (S.decorate S.SynBuiltin >>> S.sym "prim")
         >>> S.el (primDatumIso id (S.datumIso @Val))
         >>> S.el S.datumIso
 
+instance S.DatumIso Kexp where
+  datumIso = S.match
+    $ S.With (S.datumIso @Name >>>)
+    $ S.With (S.datumIso @Kappa >>>)
+    $ S.End
+
 instance S.DatumIso Program where
   datumIso = S.with \prog -> S.datumIso @Lambda >>> prog
 
+-- the printed representation is pretty dishonest in its current
+-- state. consider the following hoisted program:
+-- 
+--  (letrec ((k (κ () (continue start-ktail 123))))
+--    (λ (start-ktail)
+--      (continue k)))
+-- 
+-- here, `start-ktail` is bound in `k`, but the printed representation
+-- fails to reflect that.
+instance S.DatumIso HoistedProgram where
+  datumIso = S.with \prog ->
+    S.letLike "letrec"
+      (S.datumIso @Name) (S.datumIso @Abs) (S.datumIso @Lambda)
+    >>> S.onTail (S.iso H.fromList H.toList)
+    >>> prog
+
 
 -- quasiquoters
 
@@ -298,6 +366,7 @@ instance CPS Kappa where toCPS = S.fromDatumUnsafe S.datumIso
 instance CPS Lambda where toCPS = S.fromDatumUnsafe S.datumIso
 instance CPS Abs where toCPS = S.fromDatumUnsafe S.datumIso
 instance CPS Program where toCPS = S.fromDatumUnsafe S.datumIso
+instance CPS HoistedProgram where toCPS = S.fromDatumUnsafe S.datumIso
 
 cps :: S.QuasiQuoter
 cps = S.makeSx' [| toCPS |]
@@ -320,7 +389,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
 
@@ -330,11 +400,21 @@ 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 Kexp where
+  freeWithBound' bound = \case
+    KexpVar x -> mif (`notElem` bound) x
+    KexpKappa kap -> freeWithBound' bound kap
+
 instance Free Exp where
   freeWithBound' bound = \case
     ExpPrim p k ->
-      p & toListOf (folded . #ValVar . filtered (`notElem` bound))
+      (p ^.. folded . #ValVar . filtered (`notElem` bound))
-        & (<> freeWithBound' bound k)
+      ++ freeWithBound' bound k
     ExpLetRec bs m ->
       foldMapOf (each . _2) (freeWithBound' bound') bs
       <> freeWithBound' bound' m
@@ -342,10 +422,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))
+      <> freeWithBound' bound k
 
 instance Free Kappa where
   freeWithBound' bound (MkKappa xs m) =

src/Gyehoek/Driver.hs

@@ -35,6 +35,8 @@ import Control.Arrow ((>>>))
 import Gyehoek.Prelude
 import Gyehoek.Jalmot
 import qualified Gyehoek.Sexp as S
+import Gyehoek.CPS.Hoist (hoistProgram)
+import Gyehoek.CPS.Contify (contifyProgram)
 
 
 main :: IO ()
@@ -122,9 +124,15 @@ driver opts = do
   closedCps <- closeProgram cps
   when opts.dumpClosed do
     hPutStrLn FS.stdout =<< S.encodeWith S.datumIso closedCps
+  hoistedCps <- hoistProgram closedCps
+  when opts.dumpHoisted do
+    hPutStrLn FS.stdout =<< S.encodeWith S.datumIso hoistedCps
+  -- contifiedCps <- contifyProgram hoistedCps
+  -- when opts.dumpContified do
+  --   hPutStrLn FS.stdout =<< S.encodeWith S.datumIso contifiedCps
   let rt_is p = is (_Just . p) opts.runtime
   dumpOrRun opts.dumpStackified (rt_is #Stackify)
-    (stackifyProgram closedCps)
+    (stackifyProgram hoistedCps)
     (hPutStrLn FS.stdout <=< S.encodeDataWith S.dataIso)
     (eval >=> fmap writeObj
          >>> T.unwords
@@ -140,7 +148,7 @@ driver opts = do
   --   inspectWasm
   --   (\wat -> withFile opts.output FS.WriteMode \h -> hPutStrLn h wat)
   when opts.traceStackified do
-    stackifyProgram closedCps >>= traceEval
+    stackifyProgram hoistedCps >>= traceEval
 
 parse_e2e :: FilePath -> IO Scm.Program
 parse_e2e = runJalmotIO . runFileSystem . readScm

src/Gyehoek/Options.hs

@@ -29,6 +29,8 @@ data Options = MkOptions
   , dumpCPS :: Bool
   , dumpParsed :: Bool
   , dumpStackified :: Bool
+  , dumpHoisted :: Bool
+  , dumpContified :: Bool
   , traceStackified :: Bool
   , runtime :: Maybe Runtime
   , inspectWasm :: Bool
@@ -61,6 +63,8 @@ parser = do
   dumpCPS <- switch (long "dump-cps")
   dumpStackified <- switch (long "dump-stackified")
   dumpParsed <- switch (long "dump-parsed")
+  dumpHoisted <- switch (long "dump-hoisted")
+  dumpContified <- switch (long "dump-contified")
   traceStackified <- switch (long "trace-stackified")
   inspectWasm <- switch $ long "inspect-wasm" <> short 'p'
   runtime <- option runtimeReader . fold $

src/Gyehoek/Scheme/Syntax.hs

@@ -81,11 +81,13 @@ data Prim e
   | PrimZeroP e
   | PrimNewline
   | PrimMakeClosure { code :: e, env :: List e }
+  | PrimMakeSharedClosure { codes :: List e, env :: List e }
+  | PrimGetEnv
   | PrimEnv
   | PrimEnvRef Int
-  | PrimEnvCode e
   | PrimCallCC e
   | PrimCaptureCC
+  | PrimInvokeCC e (List e)
   | PrimValues (List e)
   | PrimCallWithValues e e
   deriving stock (Show, Generic, Functor, Foldable, Traversable, Data, Eq)
@@ -168,11 +170,15 @@ primDatumIso namefn a = S.match
     $ S.With (. ht1 "zero?")
     $ S.With (. ht0 "newline")
     $ S.With (. ht1' "make-closure")
+    $ S.With (. S.headTagged2 (namefn "make-shared-closure")
+      (S.list $ S.rest a)
+      (S.list $ S.rest a))
+    $ S.With (. ht0 "get-env")
     $ S.With (. ht0 "env")
     $ S.With (. S.headTagged1 (namefn "env-ref") S.int)
-    $ S.With (. ht1 "env-code")
     $ S.With (. ht1 "call/cc")
     $ S.With (. ht0 "capture/cc")
+    $ S.With (. ht1' "invoke/cc")
     $ S.With (. ht0' "values")
     $ S.With (. ht2 "call-with-values")
     $ S.End

src/Gyehoek/Sexp/Grammar/Base.hs

@@ -9,7 +9,7 @@ module Gyehoek.Sexp.Grammar.Base
   , DatumGrammar
   , DataGrammar
   , Grammar
-  , ListContext
+  , ListContext(..)
   , (:-)((:-))
   -- * lists
   , list

src/Gyehoek/Stack/Syntax.hs

@@ -68,6 +68,7 @@ data Tail
   | Call Int
   | If Val Block Block
   | Return Int
+  | CallCC
   deriving stock (Show, Generic, Data)
   deriving anyclass (NFData)
 
@@ -116,6 +117,7 @@ instance S.DatumIso Tail where
     $ S.With (S.headTagged1 "call" S.datumIso >>>)
     $ S.With (if_ >>>)
     $ S.With (S.headTagged1 "return" S.datumIso >>>)
+    $ S.With (S.headTagged0 "call/cc" >>>)
     $ S.End
     where
       -- if_ = S.ifLike "if" (S.datumIso @Val) S.datumIso S.datumIso

src/Gyehoek/Stack/VM.hs

@@ -138,10 +138,14 @@ stepT g vm tc@(Call nargs) = do
   pure $ vm
     & jumpToRoutine rt
     & activeFrame .~ frm
-    & #stack %~ pushFrame newFrame
     -- it is not essential we clear the registers, but it'll
     -- make bugs more obvious.
     & #registers .~ mempty
+    & #stack %~ \stk ->
+      case f of
+        ObjHob (HobContinuation {stack}) ->
+          coerce $ stack & _NonEmpty . _1 <>:~ (args ++ [f])
+        _ -> pushFrame newFrame stk
 
 stepT g vm tc@(Return nret) = do
   (xs,_) <- splitAtExact nret (vm ^. activeFrame . #locals)
@@ -179,6 +183,21 @@ stepT g vm (If c t f) = do
       _ -> t
   pure $ jumpToBlock branch vm
 
+stepT g vm CallCC = do
+  (cc,withcc,frm) <- parseCallCC (vm ^. activeFrame)
+    & expectOf "bad call/cc" _Just
+  let stk = vm.stack & #frames . _NonEmpty . _1 .~ frm
+  let reified_cc = ObjHob $ HobContinuation cc (coerce stk)
+  let newFrame = MkFrame [reified_cc, withcc, cc]
+  rt <- getRoutine g withcc
+  pure $ vm
+    & jumpToRoutine rt
+    -- replace the active frame; don't push a new one.
+    & activeFrame .~ newFrame
+    -- it is not essential we clear the registers, but it'll make
+    -- bugs more obvious.
+    & #registers .~ mempty
+
 stepP :: Jalmot :> es => Env -> VM -> Prim Val -> Eff es VM
 stepP g vm p = traverse (evalVal g vm) p >>= \case
   PrimZeroP x -> case x of
@@ -192,10 +211,6 @@ stepP g vm p = traverse (evalVal g vm) p >>= \case
     case f of
       ObjImm (ImmLabel l) -> ret1 . ObjHob $ HobClosure l env
       _ -> vmerror [i|expected label, got #{f}|]
-  PrimEnvCode env ->
-    case env of
-      ObjHob (HobClosure l _) -> ret1 . ObjImm . ImmLabel $ l
-      _ -> vmerror [i|expected closure, got #{env}|]
   PrimEnv -> do
     x <- vm & expectOf "expected closure" (activeFrame . activeProcedure)
     ret1 x
@@ -211,6 +226,10 @@ stepP g vm p = traverse (evalVal g vm) p >>= \case
   PrimCdr x -> case x of
       ObjHob (HobPair _ cdr) -> ret1 cdr
       _ -> vmerror [i|expected pair, got ${x}|]
+  -- PrimCaptureCC -> do
+  --   label <- vm & expectOf [i|bad stack, no return addr|]
+  --     (activeFrame . returnAddress . #_ObjImm . #_ImmLabel)
+  --   ret1 . ObjHob $ HobContinuation { label }
   x -> vmerror [i|unimplemented prim: #{p}|]
   where
     ret vs = pure $ vm & activeFrame . #locals <>:~ vs
@@ -239,6 +258,7 @@ jumpToRoutine rt vm = vm
 getLabel :: Obj -> Maybe Label
 getLabel = \case
   ObjHob (HobClosure {label}) -> Just label
+  ObjHob (HobContinuation {cont}) -> getLabel cont
   ObjImm (ImmLabel label) -> Just label
   x -> Nothing
 
@@ -277,6 +297,11 @@ takeExact n xs = case compareLength xs n of
   (EQ;GT) -> Just $ take n xs
   LT -> Nothing
 
+parseCallCC :: Frame -> Maybe (Obj, Obj, Frame)
+parseCallCC frm = do
+  ([cc,withcc],ys) <- splitAtExact 2 (frm ^. #locals)
+  pure (cc,withcc,MkFrame ys)
+
 parseCall :: Int -> Frame -> Maybe (List Obj, Obj, Obj, Frame)
 parseCall nargs frm = do
   (xs,ys) <- splitAtExact (nargs+2) (frm ^. #locals)

t.scm

@@ -0,0 +1,5 @@
+((λ ()
+   (* 2 (call/cc
+         (λ (k)
+           (begin (k 6)
+                  3))))))

test/Gyehoek/Test/Golden.hs

@@ -29,16 +29,7 @@ brokenWasmTests =
 
 brokenStackifyTests :: List String
 brokenStackifyTests =
-  [ "callcc-early-exit-1"
+  [
-  , "callcc-early-exit-2"
-  , "callcc-early-exit-3"
-  , "callcc-early-exit-4"
-  , "callcc-early-exit-5"
-  , "callcc-early-exit-6"
-  , "callcc-nested-1"
-  , "callcc-nested-2"
-  , "callcc-discard"
-  , "callcc-constant"
   ]
 
 test_root :: IO TestTree