shared closures maybe

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))))
+|]