src/Gyehoek/CPS/Close.hs

@@ -7,28 +7,44 @@ 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
 
 
+genCodeName :: GenSym :> es => Name -> Eff es Name
+genCodeName f = gensym' @Name $ f ^. _Wrapped' . to (<> "-code")
+
+bindEnv :: Name -> List Name -> Exp -> Exp
+bindEnv l frees m = [cps|
+  (letrec ((#{l} (κ #{frees} #{m})))
+    (prim (get-env) #{l}))
+|]
+
 close :: forall es. GenSym :> es => Exp -> Eff es Exp
 close = transformM \case
-  ExpLetRec [(f, ab)] 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 $ freeWithBound' [f] ab
+  ExpLetRec bs e -> do
+    let boundNames = bs ^.. each . _1
+    let frees = bs
+          & foldMapOf
+            (each . _2 . absBody)
+            (freeWithBound' $ setOf (each . _1) bs)
+          & nub
     env_cont_l <- gensym' @Name "env-cont"
-    let m = ab ^. absBody
-    let m' = [cps|
-      (letrec ((#{env_cont_l} (κ #{frees} #{m})))
-        (prim (get-env) #{env_cont_l}))
-    |]
-    e_l <- gensym' @Name "make-closure-cont"
-    let ab' = ab & absBody .~ m'
+    e_l <- gensym' @Name "letrec-body-cont"
+    -- let ab' = ab & absBody .~ m'
+    bs' <- for bs \(f,ab) -> do
+      f_code_l <- genCodeName f
+      pure
+        ( f_code_l
+        , ab & absBody %~ bindEnv f_code_l (boundNames ++ frees)
+        )
+    let labels = MkLabel <$> boundNames
     pure [cps|
-      (letrec ((#{f_code} #{ab'}))
-        (letrec ((#{e_l} (κ (#{f}) #{e})))
-          (prim (make-closure ($ #{f_code}) ##{frees})
+      (letrec #{bs'}
+        (letrec ((#{e_l} (κ #{boundNames} #{e})))
+          (prim (make-shared-closure #{labels} #{frees})
             #{e_l})))
     |]
 

src/Gyehoek/CPS/Stackify.hs

@@ -212,48 +212,44 @@ blah = [cps|
 
 p :: HoistedProgram
 p = [cps|
-(letrec ((prim-k3-code26 (κ (r2) (prim (env-ref 0) env-cont29)))
-         (env-cont29 (κ (lambda-tail1) (prim (env-ref 1) env-cont28)))
+(letrec ((make-closure-cont22 (κ (prim-k7) (prim (- n 1) prim-k7)))
          (env-cont18
-          (κ (lambda-tail1)
+          (κ (n lambda-tail1)
             (prim
              (make-closure $prim-k11-code14 lambda-tail1)
              make-closure-cont16)))
-         (make-closure-cont20 (κ (r8) (fac r6 r8)))
          (env-cont15 (κ (lambda-tail1) (continue lambda-tail1 r10)))
          (falsey-cont5
           (κ ()
             (prim
-             (make-closure $prim-k7-code21 fac n lambda-tail1)
-             make-closure-cont25)))
-         (prim-k7-code21 (κ (r6) (prim (env-ref 0) env-cont24)))
-         (env-cont19 (κ (n) (prim (env-ref 1) env-cont18)))
+             (make-closure $prim-k7-code20 fac n lambda-tail1)
+             make-closure-cont22)))
          (truthy-cont4 (κ () (continue lambda-tail1 1)))
-         (env-cont28 (κ (n) (prim (env-ref 2) env-cont27)))
-         (fac-code34
-          (λ (n lambda-tail1)
-            (prim
-             (make-closure $prim-k3-code26 lambda-tail1 n fac)
-             make-closure-cont30)))
+         (r12-code26 (κ (x13) (prim (get-env) env-cont27)))
+         (make-closure-cont28 (κ (r12) (fac 20 r12)))
          (make-closure-cont16 (κ (prim-k11) (prim (* n x9) prim-k11)))
-         (env-cont32 (κ (start-ktail0) (continue start-ktail0 x13)))
-         (prim-k11-code14 (κ (r10) (prim (env-ref 0) env-cont15)))
-         (env-cont22
-          (κ (lambda-tail1)
+         (env-cont30
+          (κ ()
+            (prim
+             (make-closure $prim-k3-code23 lambda-tail1 n fac)
+             make-closure-cont25)))
+         (env-cont21
+          (κ (fac n lambda-tail1)
             (prim
              (make-closure $r8-code17 n lambda-tail1)
-             make-closure-cont20)))
-         (make-closure-cont25 (κ (prim-k7) (prim (- n 1) prim-k7)))
-         (env-cont23 (κ (n) (prim (env-ref 2) env-cont22)))
-         (r12-code31 (κ (x13) (prim (env-ref 0) env-cont32)))
-         (make-closure-cont33 (κ (r12) (fac 20 r12)))
-         (make-closure-cont35
+             make-closure-cont19)))
+         (make-closure-cont19 (κ (r8) (fac r6 r8)))
+         (prim-k7-code20 (κ (r6) (prim (get-env) env-cont21)))
+         (prim-k11-code14 (κ (r10) (prim (get-env) env-cont15)))
+         (make-closure-cont25 (κ (prim-k3) (prim (zero? n) prim-k3)))
+         (make-closure-cont31
           (κ (fac)
-            (prim (make-closure $r12-code31 start-ktail0) make-closure-cont33)))
-         (env-cont27 (κ (fac) (if r2 truthy-cont4 falsey-cont5)))
-         (make-closure-cont30 (κ (prim-k3) (prim (zero? n) prim-k3)))
-         (r8-code17 (κ (x9) (prim (env-ref 0) env-cont19)))
-         (env-cont24 (κ (fac) (prim (env-ref 1) env-cont23))))
+            (prim (make-closure $r12-code26 start-ktail0) make-closure-cont28)))
+         (prim-k3-code23 (κ (r2) (prim (get-env) env-cont24)))
+         (env-cont27 (κ (start-ktail0) (continue start-ktail0 x13)))
+         (r8-code17 (κ (x9) (prim (get-env) env-cont18)))
+         (env-cont24 (κ (lambda-tail1 n fac) (if r2 truthy-cont4 falsey-cont5)))
+         (fac-code29 (λ (n lambda-tail1) (prim (get-env) env-cont30))))
   (λ (start-ktail0)
-    (prim (make-closure $fac-code34) make-closure-cont35)))
+    (prim (make-closure $fac-code29) make-closure-cont31)))
 |]

src/Gyehoek/Scheme/Syntax.hs

@@ -81,6 +81,7 @@ data Prim e
   | PrimZeroP e
   | PrimNewline
   | PrimMakeClosure { code :: e, env :: List e }
+  | PrimMakeSharedClosure { codes :: List e, env :: List e }
   | PrimGetEnv
   | PrimEnv
   | PrimEnvRef Int
@@ -169,6 +170,9 @@ 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)