From 98d33449129b9280603912e03fa52d019bb6bffa Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Tue, 17 Feb 2026 23:10:55 +0100 Subject: [PATCH] feat: Add tail call optimization Introduce a TCO pass to the compiler and modify the VM to handle tail calls. This allows for deep recursion without stack overflow. --- examples/tco_test.myc | 8 ++ src/ast/compiler/bound_nodes.rs | 7 +- src/ast/compiler/mod.rs | 2 + src/ast/compiler/tco.rs | 128 ++++++++++++++++++++++++++++++++ src/ast/environment.rs | 9 ++- src/ast/types.rs | 3 + src/ast/vm.rs | 75 +++++++++++++------ 7 files changed, 205 insertions(+), 27 deletions(-) create mode 100644 examples/tco_test.myc create mode 100644 src/ast/compiler/tco.rs diff --git a/examples/tco_test.myc b/examples/tco_test.myc new file mode 100644 index 0000000..3872969 --- /dev/null +++ b/examples/tco_test.myc @@ -0,0 +1,8 @@ +(do + (def count_down (fn [n] + (if (< n 1) + "done" + (count_down (- n 1))))) + + (count_down 1000000) +) diff --git a/src/ast/compiler/bound_nodes.rs b/src/ast/compiler/bound_nodes.rs index f06014f..f66dd75 100644 --- a/src/ast/compiler/bound_nodes.rs +++ b/src/ast/compiler/bound_nodes.rs @@ -9,7 +9,7 @@ pub enum Address { Global(u32), // Index in the global environment vector } -#[derive(Debug)] +#[derive(Debug, Clone)] pub enum BoundKind { Nop, Constant(Value), @@ -47,6 +47,11 @@ pub enum BoundKind { args: Vec>, }, + TailCall { + callee: Box>, + args: Vec>, + }, + Block { exprs: Vec>, }, diff --git a/src/ast/compiler/mod.rs b/src/ast/compiler/mod.rs index 95b4160..6df5bb3 100644 --- a/src/ast/compiler/mod.rs +++ b/src/ast/compiler/mod.rs @@ -1,5 +1,7 @@ pub mod binder; pub mod bound_nodes; +pub mod tco; pub use binder::*; pub use bound_nodes::*; +pub use tco::*; diff --git a/src/ast/compiler/tco.rs b/src/ast/compiler/tco.rs new file mode 100644 index 0000000..9c60723 --- /dev/null +++ b/src/ast/compiler/tco.rs @@ -0,0 +1,128 @@ +use std::rc::Rc; +use crate::ast::nodes::Node; +use crate::ast::compiler::bound_nodes::BoundKind; +use crate::ast::types::StaticType; + +pub struct TCO; + +impl TCO { + pub fn optimize(node: Node) -> Node { + // Top-level starts NOT in tail position usually, unless the script result is returned immediately + // which implies tail position for the script body if viewed as a function. + // Let's assume standard execution context (not tail call) for the script root. + Self::transform(node, false) + } + + fn transform(node: Node, is_tail_position: bool) -> Node { + match node.kind { + BoundKind::Call { callee, args } => { + let new_callee = Box::new(Self::transform(*callee, false)); + let new_args: Vec<_> = args.into_iter().map(|arg| Self::transform(arg, false)).collect(); + + if is_tail_position { + Node { + kind: BoundKind::TailCall { + callee: new_callee, + args: new_args, + }, + ..node + } + } else { + Node { + kind: BoundKind::Call { + callee: new_callee, + args: new_args, + }, + ..node + } + } + }, + + BoundKind::If { cond, then_br, else_br } => { + let new_cond = Box::new(Self::transform(*cond, false)); + let new_then = Box::new(Self::transform(*then_br, is_tail_position)); + let new_else = else_br.map(|e| Box::new(Self::transform(*e, is_tail_position))); + + Node { + kind: BoundKind::If { + cond: new_cond, + then_br: new_then, + else_br: new_else, + }, + ..node + } + }, + + BoundKind::Block { exprs } => { + if exprs.is_empty() { + return Node { + kind: BoundKind::Block { exprs }, + ..node + }; + } + + let last_idx = exprs.len() - 1; + let mut new_exprs = Vec::with_capacity(exprs.len()); + + for (i, expr) in exprs.into_iter().enumerate() { + let is_last = i == last_idx; + // Only the last expression inherits the tail context + new_exprs.push(Self::transform(expr, is_tail_position && is_last)); + } + + Node { + kind: BoundKind::Block { exprs: new_exprs }, + ..node + } + }, + + BoundKind::Lambda { param_count, upvalues, body } => { + // The body of a lambda is implicitly in tail position when the lambda is called. + // However, we process the lambda definition here. + // The body node inside needs to be transformed assuming it IS in tail position relative to the function. + let new_body = Rc::new(Self::transform((*body).clone(), true)); + + Node { + kind: BoundKind::Lambda { + param_count, + upvalues, + body: new_body, + }, + ..node + } + }, + + BoundKind::Set { addr, value } => { + Node { + kind: BoundKind::Set { addr, value: Box::new(Self::transform(*value, false)) }, + ..node + } + }, + BoundKind::DefGlobal { global_index, value } => { + Node { + kind: BoundKind::DefGlobal { global_index, value: Box::new(Self::transform(*value, false)) }, + ..node + } + }, + BoundKind::Map { entries } => { + let new_entries = entries.into_iter().map(|(k, v)| { + (Self::transform(k, false), Self::transform(v, false)) + }).collect(); + Node { + kind: BoundKind::Map { entries: new_entries }, + ..node + } + }, + BoundKind::Tuple { elements } => { + let new_elements = elements.into_iter().map(|e| Self::transform(e, false)).collect(); + Node { + kind: BoundKind::Tuple { elements: new_elements }, + ..node + } + }, + + // Atoms and variables don't change + _ => node, + } + } +} diff --git a/src/ast/environment.rs b/src/ast/environment.rs index f2d8ce9..a74f623 100644 --- a/src/ast/environment.rs +++ b/src/ast/environment.rs @@ -6,6 +6,8 @@ use crate::ast::parser::Parser; use crate::ast::compiler::binder::Binder; use crate::ast::vm::VM; +use crate::ast::compiler::tco::TCO; + pub struct Environment { pub global_names: Rc>>, pub global_values: Rc>>, @@ -115,9 +117,12 @@ impl Environment { // 3. Bind & Type Check let mut binder = Binder::new(self.global_names.clone()); - let bound_ast = binder.bind(&untyped_ast)?; + let mut bound_ast = binder.bind(&untyped_ast)?; - // 4. Execute + // 4. Optimize + bound_ast = TCO::optimize(bound_ast); + + // 5. Execute let mut vm = VM::new(self.global_values.clone()); vm.run(&bound_ast) } diff --git a/src/ast/types.rs b/src/ast/types.rs index 340b1b6..a4c44ec 100644 --- a/src/ast/types.rs +++ b/src/ast/types.rs @@ -70,6 +70,7 @@ pub enum Value { Function(Rc) -> Value>), Object(Rc), // For compiled Closures and other opaque types Cell(Rc>), // Boxed value for captures + TailCallRequest(Rc, Vec), // Internal: For TCO } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -113,6 +114,7 @@ impl Value { Value::Function { .. } => StaticType::Any, // Dynamic function Value::Object(o) => StaticType::Object(o.type_name()), Value::Cell(c) => c.borrow().static_type(), + Value::TailCallRequest(_, _) => panic!("Internal VM state leaked: TailCallRequest"), } } } @@ -147,6 +149,7 @@ impl fmt::Display for Value { Value::Function(_) => write!(f, ""), Value::Object(o) => write!(f, "<{}>", o.type_name()), Value::Cell(c) => write!(f, "{}", c.borrow()), + Value::TailCallRequest(_, _) => panic!("Internal VM state leaked: TailCallRequest"), } } } diff --git a/src/ast/vm.rs b/src/ast/vm.rs index 2c446e4..ce2224b 100644 --- a/src/ast/vm.rs +++ b/src/ast/vm.rs @@ -114,40 +114,67 @@ impl VM { Ok(Value::Object(Rc::new(closure))) }, + + BoundKind::TailCall { callee, args } => { + let func_val = self.eval(callee)?; + let mut arg_vals = Vec::with_capacity(args.len()); + for arg in args { + arg_vals.push(self.eval(arg)?); + } + + match func_val { + Value::Object(obj) => Ok(Value::TailCallRequest(obj, arg_vals)), + // Native functions can't be TCO'd this way (they return value directly), + // so we just execute them and return Value. + Value::Function(f) => Ok(f(arg_vals)), + _ => Err(format!("Tail call target is not a function: {}", func_val)), + } + }, BoundKind::Call { callee, args } => { - let func_val = self.eval(callee)?; + let mut func_val = self.eval(callee)?; let mut arg_vals = Vec::with_capacity(args.len()); for arg in args { arg_vals.push(self.eval(arg)?); } - match func_val { - Value::Function(f) => Ok(f(arg_vals)), - Value::Object(obj) => { - if let Some(closure) = obj.as_any().downcast_ref::() { - let old_stack_top = self.stack.len(); - let closure_rc = Rc::new(closure.clone()); - - self.stack.extend(arg_vals); - - self.frames.push(CallFrame { - stack_base: old_stack_top, - closure: Some(closure_rc.clone()), - }); + // Trampoline Loop + loop { + match func_val { + Value::Function(f) => return Ok(f(arg_vals)), + Value::Object(obj) => { + if let Some(closure) = obj.as_any().downcast_ref::() { + let old_stack_top = self.stack.len(); + let closure_rc = Rc::new(closure.clone()); + + self.stack.extend(arg_vals); // Moves arg_vals into stack + + self.frames.push(CallFrame { + stack_base: old_stack_top, + closure: Some(closure_rc.clone()), + }); - let result = self.eval(&closure.function_node); + let result = self.eval(&closure.function_node); - self.frames.pop(); - self.stack.truncate(old_stack_top); - - result - } else { - Err(format!("Object is not a closure: {}", obj.type_name())) - } - }, - _ => Err(format!("Attempt to call non-function: {}", func_val)), + self.frames.pop(); + self.stack.truncate(old_stack_top); + + match result { + Ok(Value::TailCallRequest(next_obj, next_args)) => { + func_val = Value::Object(next_obj); + arg_vals = next_args; + continue; + }, + Ok(val) => return Ok(val), + Err(e) => return Err(e), + } + } else { + return Err(format!("Object is not a closure: {}", obj.type_name())); + } + }, + _ => return Err(format!("Attempt to call non-function: {}", func_val)), + } } },