fix: closure-conversion of recursive functions

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

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")

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 =