From 9e0eba7f8b5de8a686552f9d72b3208ed381cedb Mon Sep 17 00:00:00 2001
From: Matthew Plant <maplant@protonmail.com>
Date: Thu, 26 Dec 2024 21:07:47 -0500
Subject: [PATCH] Fix named let

---
 src/ast/parse.rs | 37 +++++++++++++++++++++++++++++++------
 src/env.rs       |  7 ++++++-
 tests/r6rs.scm   | 23 +++++++++++++++++++++++
 3 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/src/ast/parse.rs b/src/ast/parse.rs
index 21a5667..9af88d4 100644
--- a/src/ast/parse.rs
+++ b/src/ast/parse.rs
@@ -1,6 +1,6 @@
 use crate::{
     continuation::Continuation,
-    env::Env,
+    env::{Env, VarRef},
     error::RuntimeError,
     expand::{SyntaxRule, Transformer},
     gc::Gc,
@@ -579,8 +579,16 @@ async fn parse_let(
         }
     }
 
-    let new_env = env.push_lexical_contour(Gc::new(new_contour));
-    let body = Body::parse(body, &new_env, cont, span).await?;
+    let new_env = Gc::new(new_contour);
+    let new_exp_env = env.push_lexical_contour(new_env.clone());
+
+    if let Some(name) = name {
+        new_env
+            .write()
+            .def_local_var(name, Gc::new(Value::Undefined));
+    }
+
+    let mut ast_body = Body::parse(body, &new_exp_env, cont, span).await?;
 
     let mut bindings: Vec<_> = compiled_bindings
         .into_iter()
@@ -589,17 +597,34 @@ async fn parse_let(
 
     // If this is a named let, add a binding for a procedure with the same
     // body and args of the formals.
+    // This code is really bad, but I wanted to get this working
     if let Some(name) = name {
+        let mut new_new_env = new_env.new_lexical_contour();
+        for (binding, _) in &bindings {
+            new_new_env.def_local_var(binding, Gc::new(Value::Undefined));
+        }
+        let new_new_exp_env = new_exp_env.push_lexical_contour(Gc::new(new_new_env));
         let lambda = Lambda {
             args: Formals::FixedArgs(bindings.iter().map(|(ident, _)| ident.clone()).collect()),
-            body: body.clone(),
+            body: Body::parse(body, &new_new_exp_env, cont, span).await?,
+        };
+        ast_body = Body {
+            forms: ArcSlice::from(
+                vec![AstNode::Expression(Expression::Set(Set {
+                    var: Ref::Regular(VarRef::default().offset(bindings.len())),
+                    val: Arc::new(Expression::Lambda(lambda)),
+                }))]
+                .into_iter()
+                .chain(ast_body.forms.iter().map(|(x, _)| x).cloned())
+                .collect::<Vec<_>>(),
+            ),
         };
-        bindings.push((name.clone(), Expression::Lambda(lambda)));
+        bindings.push((name.clone(), Expression::Undefined));
     }
 
     Ok(Let {
         bindings: Arc::from(bindings),
-        body,
+        body: ast_body,
     })
 }
 
diff --git a/src/env.rs b/src/env.rs
index 304a8e2..36fee38 100644
--- a/src/env.rs
+++ b/src/env.rs
@@ -214,7 +214,7 @@ pub enum EvalError<'e> {
 
 /// Reference to a variable that is accessible via the current environment. Could be
 /// local or non-local depending on the depth field.
-#[derive(Debug, Copy, Clone, Trace)]
+#[derive(Debug, Copy, Clone, Trace, Default)]
 pub struct VarRef {
     /// Number of up environments we need to traverse in order to reach the defining
     /// scope of the variable. Variables with a depth of 0 are local.
@@ -238,6 +238,11 @@ impl VarRef {
         }
     }
 
+    pub fn offset(mut self, offset: usize) -> Self {
+        self.offset += offset;
+        self
+    }
+
     pub fn fetch(&self, env: &Gc<Env>) -> Gc<Value> {
         env.read().fetch_var(*self)
     }
diff --git a/tests/r6rs.scm b/tests/r6rs.scm
index a6d7933..dc40de9 100644
--- a/tests/r6rs.scm
+++ b/tests/r6rs.scm
@@ -201,3 +201,26 @@
 ;;                           ((x y) (values a b)))
 ;;               (list a b x y)))
 ;;           (x y x y))
+
+(assert-eq (let loop ((n 1))
+             (if (> n 10)
+                 '()
+                 (cons n (loop (+ n 1)))))
+           '(1 2 3 4 5 6 7 8 9 10))
+           
+(define-syntax loop
+  (lambda (x)
+    (syntax-case x ()
+      [(k e ...)
+       (with-syntax
+           ([break (datum->syntax #'k 'break)])
+         #'(call-with-current-continuation
+            (lambda (break)
+              (let f () e ... (f)))))])))
+
+(assert-eq (let ((n 3) (ls '()))
+             (loop
+              (if (= n 0) (break ls))
+              (set! ls (cons 'a ls))
+              (set! n (- n 1))))
+           '(a a a))