From 276c2c1249dc9c90941385a213664c993e3911b2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Madeleine=20Sydney=20=C5=9Alaga?= <msyds@deertopia.net>
Date: Sun, 30 Aug 2026 01:55:44 -0600
Subject: [PATCH] fix: closure-conversion of recursive functions

---
 src/Gyehoek/CPS/Close.hs     |  6 +++---
 src/Gyehoek/Scheme/Syntax.hs |  2 ++
 src/Gyehoek/Stack/VM.hs      | 24 ++++++++++++------------
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/src/Gyehoek/CPS/Close.hs b/src/Gyehoek/CPS/Close.hs
index 61991f5..e04e869 100644
--- a/src/Gyehoek/CPS/Close.hs
+++ b/src/Gyehoek/CPS/Close.hs
@@ -16,12 +16,12 @@ close = transformM \case
     -- 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] lam
+    let frees = nub $ free' lam
     let m' = ifoldr
              (\n x q ->
-               let n' = if x == f then 0 else n+1
+               let p = if x == f then PrimEnv @Val else PrimEnvRef n
                in [cps|
-                 (prim (env-ref #{n'})
+                 (prim #{p}
                    (κ (#{x}) #{q}))
              |])
              m frees
diff --git a/src/Gyehoek/Scheme/Syntax.hs b/src/Gyehoek/Scheme/Syntax.hs
index c302484..d928c58 100644
--- a/src/Gyehoek/Scheme/Syntax.hs
+++ b/src/Gyehoek/Scheme/Syntax.hs
@@ -81,6 +81,7 @@ data Prim e
   | PrimZeroP e
   | PrimNewline
   | PrimMakeClosure { code :: e, env :: List e }
+  | PrimEnv
   | PrimEnvRef Int
   | PrimEnvCode e
   | PrimCallCC e
@@ -167,6 +168,7 @@ primDatumIso namefn a = S.match
     $ S.With (. ht1 "zero?")
     $ S.With (. ht0 "newline")
     $ S.With (. ht1' "make-closure")
+    $ S.With (. ht0 "env")
     $ S.With (. S.headTagged1 (namefn "env-ref") S.int)
     $ S.With (. ht1 "env-code")
     $ S.With (. ht1 "call/cc")
diff --git a/src/Gyehoek/Stack/VM.hs b/src/Gyehoek/Stack/VM.hs
index 5cdc20b..85b0164 100644
--- a/src/Gyehoek/Stack/VM.hs
+++ b/src/Gyehoek/Stack/VM.hs
@@ -132,7 +132,7 @@ stepT :: Jalmot :> es => Env -> VM -> Tail -> Eff es VM
 
 stepT g vm tc@(Call nargs) = do
   (args,f,ret,frm) <- parseCall nargs (vm ^. activeFrame)
-    & expectOf [i|#{show tc}|] _Just 
+    & expectOf [i|bad call: #{show tc}|] _Just 
   rt <- getRoutine g f
   let newFrame = MkFrame $ args ++ [f,ret]
   pure $ vm
@@ -145,7 +145,7 @@ stepT g vm tc@(Call nargs) = do
 
 stepT g vm tc@(Return nret) = do
   (xs,_) <- splitAtExact nret (vm ^. activeFrame . #locals)
-    & expectOf [i|"#{show tc}"|] _Just
+    & expectOf [i|bad return: #{show tc}|] _Just
   expectOf [i|no return addr|] (activeFrame . returnAddress) vm >>= \case
     ObjLabel "halt" -> pure $ vm & #result ?~ xs
     ra -> do
@@ -159,7 +159,7 @@ stepT g vm tc@(Return nret) = do
 
 stepT g vm tc@(TailCall nargs) = do
   (args,f,ra) <- parseTailCall nargs (vm ^. activeFrame)
-    & expectOf [i|#{show tc}|] _Just
+    & expectOf [i|bad call: #{show tc}|] _Just
   case f of
     ObjLabel "halt" -> pure $ vm & #result ?~ args
     _ -> do
@@ -196,12 +196,14 @@ stepP g vm p = traverse (evalVal g vm) p >>= \case
     case env of
       ObjHob (HobClosure l _) -> ret1 . ObjImm . ImmLabel $ l
       _ -> vmerror [i|expected closure, got #{env}|]
-  PrimEnvRef n ->
-    case vm ^? activeFrame . activeProcedure of
-      Just (ObjHob (HobClosure {label,env})) -> ret1 case n of
-        0 -> ObjLabel label
-        n -> env ^?! ix (n-1)
-      x -> vmerror [i|expected closure, got #{x}|]
+  PrimEnv -> do
+    x <- vm & expectOf "expected closure" (activeFrame . activeProcedure)
+    ret1 x
+  PrimEnvRef n -> do
+    (label,env) <- vm & expectOf "expected closure"
+        (activeFrame . activeProcedure . #_ObjHob . #_HobClosure) 
+    x <- env & expectOf "expected upval" (ix n)
+    ret1 x
   PrimCons x y -> ret1 $ ObjHob $ HobPair x y
   PrimCar x -> case x of
       ObjHob (HobPair car _) -> ret1 car
@@ -250,9 +252,7 @@ getRoutine g f = do
 expectOf
   :: (HasCallStack, Jalmot :> es)
   => Text -> Getting (First a) s a -> s -> Eff es a
-expectOf msg l s = case s ^? l of
-  Just x -> pure x
-  Nothing -> vmerror [i|bad stack, expecting #{msg}|]
+expectOf msg l = maybe (vmerror msg) pure . preview l
 
 evalToLabel :: Jalmot :> es => Env -> VM -> Val -> Eff es Label
 evalToLabel e vm v =