Brummel/RustAst
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Add typed lambda registry to Optimizer

Browse Source 
The Optimizer now accepts a typed lambda registry, allowing it to
perform more aggressive inlining and beta-reduction on globally defined
functions. This is a significant step towards optimizing recursive and
globally defined lambdas more effectively.
This commit is contained in:
Michael Schimmel
2026-02-22 02:31:55 +01:00
parent 8e8d85c06c
commit 2123f1d279
4 changed files with 197 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/lambda_collector.rs
+7 -7
View File
@@ -3,18 +3,18 @@ use crate::ast::compiler::bound_nodes::{BoundKind, Address, BoundNode};
/// A pass that collects all global function definitions (lambdas) into a registry.
/// This allows the Specializer to retrieve the original AST of a function for monomorphization.
pub struct LambdaCollector<'a> {
registry: &'a mut HashMap<u32, BoundNode>,
pub struct LambdaCollector<'a, T> {
registry: &'a mut HashMap<u32, BoundNode<T>>,
}
impl<'a> LambdaCollector<'a> {
impl<'a, T: Clone> LambdaCollector<'a, T> {
/// Performs a full traversal of the AST and populates the provided registry.
pub fn collect(node: &BoundNode, registry: &'a mut HashMap<u32, BoundNode>) {
pub fn collect(node: &BoundNode<T>, registry: &'a mut HashMap<u32, BoundNode<T>>) {
let mut collector = Self { registry };
collector.visit(node);
}
fn visit(&mut self, node: &BoundNode) {
fn visit(&mut self, node: &BoundNode<T>) {
match &node.kind {
BoundKind::Block { exprs } => {
for expr in exprs {
@@ -25,7 +25,7 @@ impl<'a> LambdaCollector<'a> {
BoundKind::DefGlobal { global_index, value, .. } => {
// Register the full Lambda node as the template.
if let BoundKind::Lambda { .. } = &value.kind {
self.registry.insert(*global_index, (**value).clone());
self.registry.insert(*global_index, (*value).as_ref().clone());
}
self.visit(value);
}
@@ -35,7 +35,7 @@ impl<'a> LambdaCollector<'a> {
if let Address::Global(global_index) = addr
&& let BoundKind::Lambda { .. } = &value.kind
{
self.registry.insert(*global_index, (**value).clone());
self.registry.insert(*global_index, (*value).as_ref().clone());
}
self.visit(value);
}
src/ast/compiler/optimizer.rs
+181 -67
View File
@@ -13,11 +13,12 @@ pub struct Optimizer {
max_passes: usize,
pub globals: Option<Rc<RefCell<Vec<Value>>>>,
pub global_purity: Option<Rc<RefCell<HashMap<u32, bool>>>>,
pub lambda_registry: Option<Rc<RefCell<HashMap<u32, TypedNode>>>>,
}
impl Optimizer {
pub fn new(enabled: bool) -> Self {
Self { enabled, max_passes: 5, globals: None, global_purity: None }
Self { enabled, max_passes: 5, globals: None, global_purity: None, lambda_registry: None }
}
pub fn with_globals(mut self, globals: Rc<RefCell<Vec<Value>>>) -> Self {
@@ -30,6 +31,11 @@ impl Optimizer {
self
}
pub fn with_registry(mut self, registry: Rc<RefCell<HashMap<u32, TypedNode>>>) -> Self {
self.lambda_registry = Some(registry);
self
}
pub fn optimize(&self, node: TypedNode) -> TypedNode {
if !self.enabled {
return node;
@@ -52,12 +58,17 @@ impl Optimizer {
BoundKind::Get { addr, name } => {
let new_addr = match addr {
Address::Local(slot) => {
if let Some(val) = sub.locals.get(&slot) {
return Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
if !sub.assigned_locals.contains(&slot) {
if let Some(inlined_node) = sub.ast_locals.get(&slot) {
return inlined_node.clone();
}
if let Some(val) = sub.locals.get(&slot) {
return Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
}
}
sub.used_slots.insert(slot);
Address::Local(sub.map_slot(slot))
@@ -73,24 +84,26 @@ impl Optimizer {
Address::Upvalue(idx)
}
Address::Global(idx) => {
if let Some(val) = sub.globals.get(&idx) {
return Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
}
if let Some(globals_rc) = &self.globals {
let globals = globals_rc.borrow();
if let Some(val) = globals.get(idx as usize)
&& self.is_inlinable_value(val)
{
if !sub.assigned_globals.contains(&idx) {
if let Some(val) = sub.globals.get(&idx) {
return Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
}
if let Some(globals_rc) = &self.globals {
let globals = globals_rc.borrow();
if let Some(val) = globals.get(idx as usize)
&& self.is_inlinable_value(val)
{
return Node {
identity: node.identity,
ty: val.static_type(),
kind: BoundKind::Constant(val.clone()),
};
}
}
}
sub.used_globals.insert(idx);
Address::Global(idx)
@@ -128,28 +141,64 @@ impl Optimizer {
(BoundKind::Set { addr: new_addr, value }, node.ty)
},
BoundKind::Call { callee, args } => {
let callee = self.visit_node(*callee, sub);
let args = self.visit_node(*args, sub);
if self.enabled {
// Case 1: Beta-Reduction for Lambda Literals
if let BoundKind::Lambda { params, body, .. } = &callee.kind
&& let Some(collapsed) = self.try_beta_reduce(params, &args, (**body).clone()) {
return self.visit_node(collapsed, sub);
}
// Case 2: Cracking and Inlining for Constant Closures
BoundKind::Call { callee, args } => {
let callee = self.visit_node(*callee, sub);
let args = self.visit_node(*args, sub);
if self.enabled {
// Case 1: Beta-Reduction for Lambda Literals
if let BoundKind::Lambda { params, body, upvalues, positional_count } = &callee.kind
&& upvalues.is_empty()
&& positional_count.is_some()
&& let Some(collapsed) = self.try_beta_reduce(params, &args, (**body).clone()) {
return collapsed;
}
// Case 1.5: Beta-Reduction for Global Registry functions (from Registry)
if let BoundKind::Get { addr: Address::Global(idx), .. } = &callee.kind
&& let Some(registry_rc) = &self.lambda_registry
&& sub.inlining_depth < 5
&& !sub.inlining_stack.contains(idx) {
let registry = registry_rc.borrow();
if let Some(lambda_node) = registry.get(idx)
&& let BoundKind::Lambda { params, body, upvalues, positional_count } = &lambda_node.kind
&& upvalues.is_empty()
&& positional_count.is_some()
{
// RECURSION CHECK: Don't inline if the function calls itself
let mut used_in_body = HashSet::new();
let mut assigned_in_body = HashSet::new();
let mut used_globals_in_body = HashSet::new();
let mut assigned_globals_in_body = HashSet::new();
self.collect_usage(body, &mut used_in_body, &mut used_globals_in_body, &mut assigned_in_body, &mut assigned_globals_in_body);
if !used_globals_in_body.contains(idx) {
let mut inner_sub = SubstitutionMap::new();
inner_sub.inlining_stack = sub.inlining_stack.clone();
inner_sub.inlining_stack.insert(*idx);
inner_sub.inlining_depth = sub.inlining_depth + 1;
if let Some(collapsed) = self.try_beta_reduce_with_sub(params.as_ref(), &args, body.as_ref().clone(), &mut inner_sub) {
return collapsed;
}
}
}
}
// Case 2: Cracking and Inlining for Constant Closures
if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
&& sub.inlining_depth < 5
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>()
&& (closure.upvalues.is_empty() || self.is_pure(&closure.function_node))
{
let mut closure_sub = SubstitutionMap::new();
closure_sub.inlining_depth = sub.inlining_depth + 1;
for (i, cell) in closure.upvalues.iter().enumerate() {
closure_sub.add_upvalue(i as u32, cell.borrow().clone());
}
let inlined_body = self.visit_node((*closure.function_node).clone(), &mut closure_sub);
if let Some(collapsed) = self.try_beta_reduce(&closure.parameter_node, &args, inlined_body) {
return self.visit_node(collapsed, sub);
if let Some(collapsed) = self.try_beta_reduce_with_sub(&closure.parameter_node, &args, inlined_body, &mut closure_sub) {
return collapsed;
}
}
@@ -160,8 +209,12 @@ impl Optimizer {
if self.enabled
&& let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
&& sub.inlining_depth < 5
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>()
&& (closure.upvalues.is_empty() || self.is_pure(&closure.function_node))
{
let mut closure_sub = SubstitutionMap::new();
closure_sub.inlining_depth = sub.inlining_depth + 1;
for (i, cell) in closure.upvalues.iter().enumerate() {
closure_sub.add_upvalue(i as u32, cell.borrow().clone());
}
@@ -208,9 +261,12 @@ impl Optimizer {
BoundKind::Block { exprs } => {
let mut used_in_block = HashSet::new();
let mut used_globals_in_block = HashSet::new();
let mut assigned_locals_in_block = HashSet::new();
let mut assigned_globals_in_block = HashSet::new();
if !exprs.is_empty() {
for e in &exprs {
self.collect_usage(e, &mut used_in_block, &mut used_globals_in_block);
self.collect_usage(e, &mut used_in_block, &mut used_globals_in_block, &mut assigned_locals_in_block, &mut assigned_globals_in_block);
}
if let Some(last) = exprs.last() {
match &last.kind {
@@ -221,6 +277,11 @@ impl Optimizer {
}
}
// IMPORTANT: Propagate block assignments to the outer map so we don't
// incorrectly inline these variables in the rest of the script.
for slot in &assigned_locals_in_block { sub.assigned_locals.insert(*slot); }
for idx in &assigned_globals_in_block { sub.assigned_globals.insert(*idx); }
let mut new_exprs = Vec::with_capacity(exprs.len());
let last_idx = exprs.len().saturating_sub(1);
@@ -233,7 +294,7 @@ impl Optimizer {
!used_in_block.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(value)
}
BoundKind::DefGlobal { global_index, value, .. } => {
!used_globals_in_block.contains(global_index) && self.is_pure(value)
!used_globals_in_block.contains(global_index) && (self.is_pure(value) || matches!(value.kind, BoundKind::Lambda { .. }))
}
BoundKind::Set { addr: Address::Local(slot), value, .. } => {
!used_in_block.contains(slot) && !sub.captured_slots.contains(slot) && self.is_pure(value)
@@ -266,10 +327,27 @@ impl Optimizer {
},
BoundKind::Lambda { params, upvalues: original_upvalues, body, positional_count } => {
BoundKind::Lambda { .. } => {
let (params, original_upvalues, body, positional_count) = if let BoundKind::Lambda { params, upvalues, body, positional_count } = &node.kind {
(params, upvalues, body, positional_count)
} else { unreachable!() };
let mut used_in_lambda = HashSet::new();
let mut used_globals_in_lambda = HashSet::new();
let mut assigned_locals_in_lambda = HashSet::new();
let mut assigned_globals_in_lambda = HashSet::new();
self.collect_usage(&node, &mut used_in_lambda, &mut used_globals_in_lambda, &mut assigned_locals_in_lambda, &mut assigned_globals_in_lambda);
// Track captures from outer scope that are assigned in this lambda
for idx in &assigned_globals_in_lambda { sub.assigned_globals.insert(*idx); }
let mut new_upvalues = Vec::new();
let mut mapping = Vec::new();
let mut next_inner_subs = SubstitutionMap::new();
next_inner_subs.inlining_depth = sub.inlining_depth;
next_inner_subs.assigned_locals = assigned_locals_in_lambda;
next_inner_subs.assigned_globals = assigned_globals_in_lambda;
for (old_idx, capture_addr) in original_upvalues.iter().enumerate() {
let mut inlined_val = None;
@@ -299,13 +377,13 @@ impl Optimizer {
}
// 1. Visit parameters to determine their new slots
let params_node = self.visit_node(params.as_ref().clone(), &mut SubstitutionMap::new());
let params_node = self.visit_node(params.as_ref().clone(), &mut next_inner_subs);
// 2. IMPORTANT: Pre-register parameter slots in the inner substitution map
// so they are mapped 1:1 and don't get reassigned to higher slots in the body.
self.collect_parameter_slots(&params_node, &mut next_inner_subs);
let body_node = self.visit_node((*body).clone(), &mut next_inner_subs);
let body_node = self.visit_node(body.as_ref().clone(), &mut next_inner_subs);
let reindexed_body = if new_upvalues.len() != original_upvalues.len() {
sub.reindex_upvalues(body_node, &mapping)
} else {
@@ -316,7 +394,7 @@ impl Optimizer {
params: Rc::new(params_node),
upvalues: new_upvalues,
body: Rc::new(reindexed_body),
positional_count
positional_count: *positional_count
}, node.ty)
},
@@ -362,7 +440,9 @@ impl Optimizer {
(BoundKind::Record { fields }, node.ty)
},
BoundKind::Expansion { original_call, bound_expanded } => {
sub.inlining_depth += 1;
let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub));
sub.inlining_depth -= 1;
(BoundKind::Expansion { original_call, bound_expanded }, node.ty)
},
k => (k, node.ty),
@@ -401,7 +481,8 @@ impl Optimizer {
fn is_pure(&self, node: &TypedNode) -> bool {
match &node.kind {
BoundKind::Constant(_) | BoundKind::Parameter { .. } | BoundKind::Nop | BoundKind::Lambda { .. } => true,
BoundKind::Constant(_) | BoundKind::Parameter { .. } | BoundKind::Nop => true,
BoundKind::Lambda { body, .. } => self.is_pure(body),
BoundKind::Get { addr, .. } => {
match addr {
Address::Global(idx) => {
@@ -421,7 +502,8 @@ impl Optimizer {
}
BoundKind::Block { exprs } => exprs.iter().all(|e| self.is_pure(e)),
BoundKind::Expansion { bound_expanded, .. } => self.is_pure(bound_expanded),
BoundKind::DefLocal { value, .. } | BoundKind::Set { addr: Address::Local(_), value } => self.is_pure(value),
BoundKind::DefLocal { value, .. } => self.is_pure(value),
BoundKind::Set { .. } => false,
BoundKind::Call { callee, args } => {
let callee_pure = match &callee.kind {
BoundKind::Get { addr: Address::Global(idx), .. } => {
@@ -481,16 +563,22 @@ impl Optimizer {
fn try_beta_reduce(&self, params: &TypedNode, args: &TypedNode, body: TypedNode) -> Option<TypedNode> {
let mut sub = SubstitutionMap::new();
self.try_beta_reduce_with_sub(params, args, body, &mut sub)
}
fn try_beta_reduce_with_sub(&self, params: &TypedNode, args: &TypedNode, body: TypedNode, sub: &mut SubstitutionMap) -> Option<TypedNode> {
let mut arg_vals = Vec::new();
self.flatten_typed_tuple(args, &mut arg_vals);
let mut slot_index = 0;
self.map_params_to_args(params, &arg_vals, &mut slot_index, &mut sub);
self.map_params_to_args(params, &arg_vals, &mut slot_index, sub);
if sub.locals.len() < arg_vals.len() { return None; }
if sub.locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() { return None; }
// Safety: Only proceed if we mapped EXACTLY the number of arguments provided.
// And if we actually have parameters to bind (or if it's a 0-arg call).
if slot_index != arg_vals.len() { return None; }
if sub.locals.is_empty() && sub.ast_locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() { return None; }
Some(self.visit_node(body, &mut sub))
Some(self.visit_node(body, sub))
}
fn flatten_typed_tuple(&self, node: &TypedNode, into: &mut Vec<TypedNode>) {
@@ -504,10 +592,17 @@ impl Optimizer {
fn map_params_to_args(&self, pattern: &TypedNode, args: &[TypedNode], offset: &mut usize, sub: &mut SubstitutionMap) {
match &pattern.kind {
BoundKind::Parameter { slot, .. } => {
if let Some(arg) = args.get(*offset)
&& let BoundKind::Constant(val) = &arg.kind {
if let Some(arg) = args.get(*offset) {
if let BoundKind::Constant(val) = &arg.kind {
sub.add_local(*slot, val.clone());
} else if let BoundKind::Lambda { upvalues, .. } = &arg.kind {
if upvalues.is_empty() {
sub.add_ast_local(*slot, arg.clone());
}
} else if let BoundKind::Get { addr: Address::Global(_), .. } = &arg.kind {
sub.add_ast_local(*slot, arg.clone());
}
}
// Parameters in beta-reduction are also identity mapped if not inlined
sub.slot_mapping.insert(*slot, *slot);
sub.next_slot = sub.next_slot.max(*slot + 1);
@@ -520,28 +615,33 @@ impl Optimizer {
}
}
fn collect_usage(&self, node: &TypedNode, used_locals: &mut HashSet<u32>, used_globals: &mut HashSet<u32>) {
fn collect_usage(&self, node: &TypedNode, used_locals: &mut HashSet<u32>, used_globals: &mut HashSet<u32>, assigned_locals: &mut HashSet<u32>, assigned_globals: &mut HashSet<u32>) {
match &node.kind {
BoundKind::Constant(v) => {
if let Value::Object(obj) = v
&& let Some(closure) = obj.as_any().downcast_ref::<Closure>()
{
self.collect_usage(&closure.function_node, used_locals, used_globals);
self.collect_usage(&closure.function_node, used_locals, used_globals, assigned_locals, assigned_globals);
}
}
BoundKind::Get { addr, .. } | BoundKind::Set { addr, .. } => {
BoundKind::Get { addr, .. } => {
match addr {
Address::Local(slot) => { used_locals.insert(*slot); }
Address::Global(idx) => { used_globals.insert(*idx); }
_ => {}
}
if let BoundKind::Set { value, .. } = &node.kind {
self.collect_usage(value, used_locals, used_globals);
}
BoundKind::Set { addr, value } => {
match addr {
Address::Local(slot) => { assigned_locals.insert(*slot); }
Address::Global(idx) => { assigned_globals.insert(*idx); }
_ => {}
}
self.collect_usage(value, used_locals, used_globals, assigned_locals, assigned_globals);
}
BoundKind::Lambda { params, body, upvalues, .. } => {
self.collect_usage(params, used_locals, used_globals);
self.collect_usage(body, used_locals, used_globals);
self.collect_usage(params, used_locals, used_globals, assigned_locals, assigned_globals);
self.collect_usage(body, used_locals, used_globals, assigned_locals, assigned_globals);
for addr in upvalues {
match addr {
Address::Local(slot) => { used_locals.insert(*slot); }
@@ -551,31 +651,31 @@ impl Optimizer {
}
}
BoundKind::Block { exprs } => {
for e in exprs { self.collect_usage(e, used_locals, used_globals); }
for e in exprs { self.collect_usage(e, used_locals, used_globals, assigned_locals, assigned_globals); }
}
BoundKind::If { cond, then_br, else_br } => {
self.collect_usage(cond, used_locals, used_globals);
self.collect_usage(then_br, used_locals, used_globals);
if let Some(e) = else_br { self.collect_usage(e, used_locals, used_globals); }
self.collect_usage(cond, used_locals, used_globals, assigned_locals, assigned_globals);
self.collect_usage(then_br, used_locals, used_globals, assigned_locals, assigned_globals);
if let Some(e) = else_br { self.collect_usage(e, used_locals, used_globals, assigned_locals, assigned_globals); }
}
BoundKind::Call { callee, args } => {
self.collect_usage(callee, used_locals, used_globals);
self.collect_usage(args, used_locals, used_globals);
self.collect_usage(callee, used_locals, used_globals, assigned_locals, assigned_globals);
self.collect_usage(args, used_locals, used_globals, assigned_locals, assigned_globals);
}
BoundKind::Tuple { elements } => {
for e in elements { self.collect_usage(e, used_locals, used_globals); }
for e in elements { self.collect_usage(e, used_locals, used_globals, assigned_locals, assigned_globals); }
}
BoundKind::Record { fields } => {
for (k, v) in fields {
self.collect_usage(k, used_locals, used_globals);
self.collect_usage(v, used_locals, used_globals);
self.collect_usage(k, used_locals, used_globals, assigned_locals, assigned_globals);
self.collect_usage(v, used_locals, used_globals, assigned_locals, assigned_globals);
}
}
BoundKind::DefLocal { value, .. } | BoundKind::DefGlobal { value, .. } => {
self.collect_usage(value, used_locals, used_globals);
self.collect_usage(value, used_locals, used_globals, assigned_locals, assigned_globals);
}
BoundKind::Expansion { bound_expanded, .. } => {
self.collect_usage(bound_expanded, used_locals, used_globals);
self.collect_usage(bound_expanded, used_locals, used_globals, assigned_locals, assigned_globals);
}
_ => {}
}
@@ -587,7 +687,12 @@ struct SubstitutionMap {
locals: HashMap<u32, Value>,
upvalues: HashMap<u32, Value>,
globals: HashMap<u32, Value>,
ast_locals: HashMap<u32, TypedNode>,
slot_mapping: HashMap<u32, u32>,
inlining_stack: HashSet<u32>,
inlining_depth: usize,
assigned_locals: HashSet<u32>,
assigned_globals: HashSet<u32>,
next_slot: u32,
used_slots: HashSet<u32>,
pure_slots: HashSet<u32>,
@@ -601,7 +706,12 @@ impl SubstitutionMap {
locals: HashMap::new(),
upvalues: HashMap::new(),
globals: HashMap::new(),
ast_locals: HashMap::new(),
slot_mapping: HashMap::new(),
inlining_stack: HashSet::new(),
inlining_depth: 0,
assigned_locals: HashSet::new(),
assigned_globals: HashSet::new(),
next_slot: 0,
used_slots: HashSet::new(),
pure_slots: HashSet::new(),
@@ -610,6 +720,10 @@ impl SubstitutionMap {
}
}
fn add_ast_local(&mut self, slot: u32, node: TypedNode) {
self.ast_locals.insert(slot, node);
}
fn map_slot(&mut self, old_slot: u32) -> u32 {
if let Some(&new_slot) = self.slot_mapping.get(&old_slot) {
return new_slot;
src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_assignment_safety.snap
+2 -2
View File
@@ -1,6 +1,6 @@
---
source: src/ast/compiler/optimizer.rs
assertion_line: 720
assertion_line: 844
expression: dump
---
Lambda (Upvalues: 0) <Metadata: Metadata { ty: Function(Signature { params: Tuple([]), ret: Int }), is_tail: true }>
@@ -11,4 +11,4 @@ Lambda (Upvalues: 0) <Metadata: Metadata { ty: Function(Signature { params: Tupl
Constant: 10 <Metadata: Metadata { ty: Int, is_tail: false }>
Set: Local(0) <Metadata: Metadata { ty: Int, is_tail: false }>
Constant: 20 <Metadata: Metadata { ty: Int, is_tail: false }>
Constant: 20 <Metadata: Metadata { ty: Int, is_tail: true }>
Get: x (Local(0)) <Metadata: Metadata { ty: Int, is_tail: true }>
src/ast/environment.rs
+7 -1
View File
@@ -24,6 +24,7 @@ pub struct Environment {
pub global_purity: Rc<RefCell<HashMap<u32, bool>>>,
pub global_values: Rc<RefCell<Vec<Value>>>,
pub function_registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
pub typed_function_registry: Rc<RefCell<HashMap<u32, TypedNode>>>,
pub monomorph_cache: Rc<RefCell<MonoCache>>,
pub debug_mode: bool,
pub optimization: bool,
@@ -89,6 +90,7 @@ impl Environment {
global_purity: Rc::new(RefCell::new(HashMap::new())),
global_values: Rc::new(RefCell::new(Vec::new())),
function_registry: Rc::new(RefCell::new(HashMap::new())),
typed_function_registry: Rc::new(RefCell::new(HashMap::new())),
monomorph_cache: Rc::new(RefCell::new(HashMap::new())),
debug_mode: false,
optimization: true,
@@ -180,6 +182,9 @@ impl Environment {
let checker = TypeChecker::new(self.global_types.clone());
let typed_ast = checker.check(bound_ast, &[])?;
// 7. Collect Typed Lambdas
LambdaCollector::collect(&typed_ast, &mut self.typed_function_registry.borrow_mut());
Ok(typed_ast)
}
@@ -191,7 +196,8 @@ impl Environment {
// 2. Optimize (Level 1: Cracking, Level 2: Collapsing)
let optimizer = Optimizer::new(self.optimization)
.with_globals(self.global_values.clone())
.with_purity(self.global_purity.clone());
.with_purity(self.global_purity.clone())
.with_registry(self.typed_function_registry.clone());
let optimized = optimizer.optimize(specialized);
// 3. TCO (Always performed, converts to ExecNode)