diff --git a/src/ast/compiler/lambda_collector.rs b/src/ast/compiler/lambda_collector.rs index 8e8fc4e..be1dbd7 100644 --- a/src/ast/compiler/lambda_collector.rs +++ b/src/ast/compiler/lambda_collector.rs @@ -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, +pub struct LambdaCollector<'a, T> { + registry: &'a mut HashMap>, } -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) { + pub fn collect(node: &BoundNode, registry: &'a mut HashMap>) { let mut collector = Self { registry }; collector.visit(node); } - fn visit(&mut self, node: &BoundNode) { + fn visit(&mut self, node: &BoundNode) { 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); } diff --git a/src/ast/compiler/optimizer.rs b/src/ast/compiler/optimizer.rs index 58c4b38..738fbaa 100644 --- a/src/ast/compiler/optimizer.rs +++ b/src/ast/compiler/optimizer.rs @@ -13,11 +13,12 @@ pub struct Optimizer { max_passes: usize, pub globals: Option>>>, pub global_purity: Option>>>, + pub lambda_registry: Option>>>, } 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>>) -> Self { @@ -30,6 +31,11 @@ impl Optimizer { self } + pub fn with_registry(mut self, registry: Rc>>) -> 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::() { + && sub.inlining_depth < 5 + && let Some(closure) = obj.as_any().downcast_ref::() + && (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::() { + && sub.inlining_depth < 5 + && let Some(closure) = obj.as_any().downcast_ref::() + && (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(¶ms_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 { 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 { 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) { @@ -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, used_globals: &mut HashSet) { + fn collect_usage(&self, node: &TypedNode, used_locals: &mut HashSet, used_globals: &mut HashSet, assigned_locals: &mut HashSet, assigned_globals: &mut HashSet) { match &node.kind { BoundKind::Constant(v) => { if let Value::Object(obj) = v && let Some(closure) = obj.as_any().downcast_ref::() { - 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, upvalues: HashMap, globals: HashMap, + ast_locals: HashMap, slot_mapping: HashMap, + inlining_stack: HashSet, + inlining_depth: usize, + assigned_locals: HashSet, + assigned_globals: HashSet, next_slot: u32, used_slots: HashSet, pure_slots: HashSet, @@ -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; diff --git a/src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_assignment_safety.snap b/src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_assignment_safety.snap index fc79742..35be8a0 100644 --- a/src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_assignment_safety.snap +++ b/src/ast/compiler/snapshots/myc__ast__compiler__optimizer__tests__opt_assignment_safety.snap @@ -1,6 +1,6 @@ --- source: src/ast/compiler/optimizer.rs -assertion_line: 720 +assertion_line: 844 expression: dump --- Lambda (Upvalues: 0) @@ -11,4 +11,4 @@ Lambda (Upvalues: 0) Set: Local(0) Constant: 20 - Constant: 20 + Get: x (Local(0)) diff --git a/src/ast/environment.rs b/src/ast/environment.rs index c8bfb8c..521a04c 100644 --- a/src/ast/environment.rs +++ b/src/ast/environment.rs @@ -24,6 +24,7 @@ pub struct Environment { pub global_purity: Rc>>, pub global_values: Rc>>, pub function_registry: Rc>>, + pub typed_function_registry: Rc>>, pub monomorph_cache: Rc>, 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)