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.
This commit is contained in:
@@ -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<Node<BoundKind, StaticType>>,
|
||||
},
|
||||
|
||||
TailCall {
|
||||
callee: Box<Node<BoundKind, StaticType>>,
|
||||
args: Vec<Node<BoundKind, StaticType>>,
|
||||
},
|
||||
|
||||
Block {
|
||||
exprs: Vec<Node<BoundKind, StaticType>>,
|
||||
},
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
pub mod binder;
|
||||
pub mod bound_nodes;
|
||||
pub mod tco;
|
||||
|
||||
pub use binder::*;
|
||||
pub use bound_nodes::*;
|
||||
pub use tco::*;
|
||||
|
||||
@@ -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<BoundKind, StaticType>) -> Node<BoundKind, StaticType> {
|
||||
// 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<BoundKind, StaticType>, is_tail_position: bool) -> Node<BoundKind, StaticType> {
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user