From cbb5eacfb2b0879db8e51348ae16be4c8b7654cf Mon Sep 17 00:00:00 2001
From: JoJo <jo@jo.zone>
Date: Tue, 12 May 2020 14:13:11 +0200
Subject: [PATCH] Fix regression where var bindings in let didn't have others
 in env

E.g. the following would not work, as `a` wasn't in the environment
when `b` was defined. This worked before, but I broke it when making
recursive function defs in `let` work. This commit makes it so both work.

(let ((a ...)
      (b a))
  ...)
---
 src/Codegen.hs | 21 ++++++++++++++++-----
 src/Gen.hs     |  3 +++
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/src/Codegen.hs b/src/Codegen.hs
index a2a228c..00f0fed 100644
--- a/src/Codegen.hs
+++ b/src/Codegen.hs
@@ -279,6 +279,8 @@ genCondBr predV genConseq genAlt = do
 
 genLet :: Defs -> Expr -> Gen Val
 genLet (Topo ds) letBody = do
+    -- For both function and variable bindings, we need separate the definition
+    -- into two passes, where the first pre-allocates some stuff.
     (binds, cs) <- fmap unzip $ forM ds $ \case
         (v, WithPos _ (_, Expr _ (Fun p (fb, fbt)))) -> do
             let fvXs = Set.toList (Set.delete p (freeVars fb))
@@ -287,11 +289,20 @@ genLet (Topo ds) letBody = do
                 (mapM (\(TypedVar _ t) -> genType t) fvXs)
             captures <- genHeapAllocGeneric tcaptures
             fbt' <- genType fbt
-            l <- genLambda' p (genExpr fb, fbt') (VLocal captures) fvXs
-            pure ((v, l), Just (captures, fvXs))
-        (v, WithPos _ (_, e)) -> genExpr e <&> \e' -> ((v, e'), Nothing)
-    withVals binds $ do
-        forM_ (catMaybes cs) (uncurry populateCaptures)
+            l <-
+                getVar
+                    =<< genLambda' p (genExpr fb, fbt') (VLocal captures) fvXs
+            pure ((v, l), Left (captures, fvXs))
+        (v@(TypedVar n t), WithPos _ (_, e)) -> do
+            t' <- genType t
+            mem <- emitReg n (alloca t')
+            pure ((v, mem), Right e)
+    withVars binds $ do
+        forM_ (zip binds cs) $ \case
+            (_, Left (captures, fvXs)) -> populateCaptures captures fvXs
+            ((_, mem), Right e) -> do
+                x <- getLocal =<< genExpr e
+                emitDo (store x mem)
         genExpr letBody
 
 genMatch :: Expr -> DecisionTree -> Type -> Gen Val
diff --git a/src/Gen.hs b/src/Gen.hs
index f970d25..434e98d 100644
--- a/src/Gen.hs
+++ b/src/Gen.hs
@@ -390,6 +390,9 @@ withLocal x v gen = do
     vPtr <- genStackAllocated v
     withVar x vPtr gen
 
+withVars :: [(TypedVar, Operand)] -> Gen a -> Gen a
+withVars = withXs withVar
+
 -- | Takes a local, stack allocated value, and runs a generator in the
 --   environment with the variable
 withVar :: TypedVar -> Operand -> Gen a -> Gen a
-- 
2.34.2