From dfb44f06baf2380052b48b6214cbafb366fb9295 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Madeleine=20Sydney=20=C5=9Alaga?= <msyds@deertopia.net>
Date: Wed, 2 Sep 2026 16:04:45 -0600
Subject: [PATCH] shared closures maybe

---
 .dir-locals.el             |  3 ++
 doc/closure-conversion.org | 31 ++++++++++++++++
 src/Gyehoek/CPS/Close.hs   | 73 +++++++++++++++++++++++++++++---------
 3 files changed, 90 insertions(+), 17 deletions(-)

diff --git a/.dir-locals.el b/.dir-locals.el
index fc1236c..6062139 100644
--- a/.dir-locals.el
+++ b/.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 ()
diff --git a/doc/closure-conversion.org b/doc/closure-conversion.org
index 56a7f1d..993a5be 100644
--- a/doc/closure-conversion.org
+++ b/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
diff --git a/src/Gyehoek/CPS/Close.hs b/src/Gyehoek/CPS/Close.hs
index d715c2c..3692bb5 100644
--- a/src/Gyehoek/CPS/Close.hs
+++ b/src/Gyehoek/CPS/Close.hs
@@ -16,38 +16,77 @@ 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}))
+bindEnv :: List Name -> Exp -> Exp
+bindEnv frees m = [cps|
+  (prim (get-env) (κ #{frees} #{m}))
 |]
 
-close :: forall es. GenSym :> es => Exp -> Eff es Exp
-close = transformM \case
-  ExpLetRec bs e -> do
+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 . absBody)
-            (freeWithBound' $ setOf (each . _1) bs)
+            (each . _2)
+            (freeWithBound' boundNames')
           & nub
+    pTraceShowM frees
     env_cont_l <- gensym' @Name "env-cont"
     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)
-        )
+      pure ( f_code_l
+           , ab & absBody %~ bindEnv (boundNames ++ frees)
+           )
+    let codes = bs' ^.. each . _1 . to MkLabel
     pure [cps|
       (letrec #{bs'}
-        (letrec ((#{e_l} (κ #{boundNames} #{e})))
-          (prim (make-shared-closure #{boundNames} #{frees})
-            #{e_l})))
+        (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))))
+|]