refactor mutable variables to boxes

crates/steel-core/src/compiler/compiler.rs

@@ -487,6 +487,8 @@ impl Compiler {
             .remove_unused_globals_with_prefix("mangler", &self.macro_env, &self.module_manager)
             .lift_pure_local_functions()
             .lift_all_local_functions();
+        // This might be sus, lets see!
+        // .replace_mutable_captured_variables_with_boxes();
 
         // TODO: Just run this... on each module in particular
         // .remove_unused_globals_with_prefix("mangler");
@@ -504,6 +506,11 @@ impl Compiler {
 
         semantic.refresh_variables();
 
+        semantic.populate_captures();
+        semantic.populate_captures();
+
+        semantic.replace_mutable_captured_variables_with_boxes();
+
         if log_enabled!(log::Level::Debug) {
             debug!(
                 "Successfully expanded defines: {:?}",
@@ -683,6 +690,12 @@ impl Compiler {
 
         semantic.refresh_variables();
 
+        // Replace mutation with boxes
+        semantic.populate_captures();
+        semantic.populate_captures();
+
+        semantic.replace_mutable_captured_variables_with_boxes();
+
         if log_enabled!(log::Level::Debug) {
             debug!(
                 "Successfully expanded defines: {:?}",

crates/steel-core/src/compiler/passes/analysis.rs

@@ -2398,6 +2398,122 @@ impl<'a> VisitorMutRefUnit for ElideSingleArgumentLambdaApplications<'a> {
     }
 }
 
+fn box_argument(ident: ExprKind) -> ExprKind {
+    ExprKind::List(List::new(vec![ExprKind::atom("#%box"), ident]))
+}
+
+fn unbox_argument(ident: ExprKind) -> ExprKind {
+    ExprKind::List(List::new(vec![ExprKind::atom("#%unbox"), ident]))
+}
+
+fn setbox_argument(ident: ExprKind, expr: ExprKind) -> ExprKind {
+    ExprKind::List(List::new(vec![ExprKind::atom("#%set-box!"), ident, expr]))
+}
+
+struct ReplaceSetOperationsWithBoxes<'a> {
+    analysis: &'a Analysis,
+}
+
+impl<'a> VisitorMutRefUnit for ReplaceSetOperationsWithBoxes<'a> {
+    fn visit(&mut self, expr: &mut ExprKind) {
+        match expr {
+            ExprKind::If(f) => self.visit_if(f),
+            ExprKind::Define(d) => self.visit_define(d),
+            ExprKind::LambdaFunction(l) => self.visit_lambda_function(l),
+            ExprKind::Begin(b) => self.visit_begin(b),
+            ExprKind::Return(r) => self.visit_return(r),
+            ExprKind::Quote(q) => self.visit_quote(q),
+            ExprKind::Macro(m) => self.visit_macro(m),
+            ExprKind::Atom(a) => {
+                if let Some(analysis) = self.analysis.get(&a.syn) {
+                    if analysis.kind == IdentifierStatus::HeapAllocated {
+                        let mut dummy = ExprKind::List(List::new(Vec::new()));
+                        std::mem::swap(&mut dummy, expr);
+
+                        *expr = unbox_argument(dummy);
+                    }
+                }
+            }
+            ExprKind::List(l) => self.visit_list(l),
+            ExprKind::SyntaxRules(s) => self.visit_syntax_rules(s),
+            ExprKind::Set(s) => {
+                if let Some(analysis) = self.analysis.get(s.variable.atom_syntax_object().unwrap())
+                {
+                    if analysis.kind != IdentifierStatus::HeapAllocated {
+                        self.visit(&mut s.expr);
+
+                        return;
+                    }
+                }
+
+                // Go ahead and drop down the expression
+                self.visit(&mut s.expr);
+
+                let mut set_expr = ExprKind::List(List::new(Vec::new()));
+                std::mem::swap(&mut s.expr, &mut set_expr);
+
+                let mut dummy_ident = ExprKind::List(List::new(Vec::new()));
+                std::mem::swap(&mut dummy_ident, &mut s.variable);
+
+                let new_set_expr = setbox_argument(dummy_ident, set_expr);
+
+                *expr = new_set_expr;
+            }
+            ExprKind::Require(r) => self.visit_require(r),
+            ExprKind::Let(l) => self.visit_let(l),
+        }
+    }
+
+    fn visit_define(&mut self, define: &mut Define) {
+        self.visit(&mut define.body);
+    }
+
+    fn visit_lambda_function(&mut self, lambda_function: &mut LambdaFunction) {
+        // Visit the body first, unwind the recursion on the way up
+        self.visit(&mut lambda_function.body);
+
+        let function_info = self
+            .analysis
+            .function_info
+            .get(&lambda_function.syntax_object_id)
+            .unwrap();
+
+        let mut mutable_variables = Vec::new();
+
+        // Which arguments do we need to wrap up
+        for var in &lambda_function.args {
+            if let Some(ident) = var.atom_identifier() {
+                if let Some(arg) = function_info.arguments().get(ident) {
+                    if arg.captured && arg.mutated {
+                        mutable_variables.push(var.clone());
+                    }
+                }
+            } else {
+                unreachable!()
+            }
+        }
+
+        if !mutable_variables.is_empty() {
+            let mut body = ExprKind::List(List::new(Vec::new()));
+
+            std::mem::swap(&mut lambda_function.body, &mut body);
+
+            let wrapped_lambda = LambdaFunction::new(
+                mutable_variables.clone(),
+                body,
+                lambda_function.location.clone(),
+            );
+
+            // Box the values!
+            let mut mutable_variables: Vec<_> =
+                mutable_variables.into_iter().map(box_argument).collect();
+
+            mutable_variables.insert(0, ExprKind::LambdaFunction(Box::new(wrapped_lambda)));
+            lambda_function.body = ExprKind::List(List::new(mutable_variables));
+        }
+    }
+}
+
 struct LiftLocallyDefinedFunctions<'a> {
     analysis: &'a Analysis,
     lifted_functions: Vec<ExprKind>,
@@ -2962,6 +3078,18 @@ impl<'a> SemanticAnalysis<'a> {
         }
     }
 
+    pub fn replace_mutable_captured_variables_with_boxes(&mut self) -> &mut Self {
+        let mut replacer = ReplaceSetOperationsWithBoxes {
+            analysis: &self.analysis,
+        };
+
+        for expr in self.exprs.iter_mut() {
+            replacer.visit(expr);
+        }
+
+        self
+    }
+
     /// Find all local pure functions, except those defined already at the top level and those defined with 'define',
     /// and replace them with a globally defined function. This means we're not going to be recreating
     /// the function _on every instance_ and instead can just grab them each time.

crates/steel-core/src/primitives/nums.rs

@@ -77,11 +77,26 @@ pub fn divide_primitive(args: &[SteelVal]) -> Result<SteelVal> {
         stop!(ArityMismatch => "/ requires at least one argument")
     }
 
+    if args.len() == 1 {
+        match &args[0] {
+            SteelVal::IntV(n) => return Ok(SteelVal::NumV((1 / n) as f64)),
+            SteelVal::NumV(n) => return Ok(SteelVal::NumV((1.0 / n) as f64)),
+            unexpected => {
+                stop!(TypeMismatch => "division expects a number, found: {:?}", unexpected)
+            }
+        }
+    }
+
+    let mut no_floats = true;
+
     let floats: Result<Vec<f64>> = args
         .iter()
         .map(|x| match x {
             SteelVal::IntV(n) => Ok(*n as f64),
-            SteelVal::NumV(n) => Ok(*n),
+            SteelVal::NumV(n) => {
+                no_floats = false;
+                Ok(*n)
+            }
             _ => stop!(TypeMismatch => "division expects a number"),
         })
         .collect();
@@ -91,7 +106,7 @@ pub fn divide_primitive(args: &[SteelVal]) -> Result<SteelVal> {
     if let Some(first) = floats.next() {
         let result = floats.fold(first, |acc, x| acc / x);
 
-        if result.fract() == 0.0 {
+        if no_floats && result.fract() == 0.0 {
             Ok(SteelVal::IntV(result as isize))
         } else {
             Ok(SteelVal::NumV(result))

crates/steel-core/src/steel_vm/vm.rs

@@ -5501,6 +5501,22 @@ fn read_alloc_handler(ctx: &mut VmCore<'_>) -> Result<()> {
         .borrow()[payload_size]
         .get();
 
+    // dbg!(payload_size);
+
+    // dbg!(ctx
+    //     .thread
+    //     .stack_frames
+    //     .last()
+    //     .unwrap()
+    //     .function
+    //     .heap_allocated()
+    //     .borrow()
+    //     .iter()
+    //     .map(|x| x.get())
+    //     .collect::<Vec<_>>());
+
+    dbg!(&value);
+
     ctx.thread.stack.push(value);
     ctx.ip += 1;
 

crates/steel-core/src/tests/success/letrec_simple_recursion.scm

@@ -1,10 +1,6 @@
 (define (test)
-    (let ((loop void))
-        (let ((loop-prime (lambda (x) 
-                            (if (= x 10000)
-                                x
-                                (loop (+ x 1))))))
-            (set! loop loop-prime))
+  (let ([loop void])
+    (let ([loop-prime (lambda (x) (if (= x 10000) x (loop (+ x 1))))]) (set! loop loop-prime))
     (loop 0)))
 
-(assert! (= (test) 10000))
+(assert! (= (test) 10000))

crates/steel-core/src/values/closed.rs

@@ -265,6 +265,10 @@ impl Heap {
             for heap_ref in function.heap_allocated.borrow().iter() {
                 context.mark_heap_reference(&heap_ref.strong_ptr())
             }
+
+            for value in function.captures() {
+                context.push_back(value.clone());
+            }
         }
 
         context.visit();