diff --git a/src/ast/compiler/lambda_collector.rs b/src/ast/compiler/lambda_collector.rs index 9a65bda..53cff95 100644 --- a/src/ast/compiler/lambda_collector.rs +++ b/src/ast/compiler/lambda_collector.rs @@ -24,22 +24,32 @@ impl<'a, T: Clone> LambdaCollector<'a, T> { BoundKind::Define { addr, value, .. } => { // Register global function definitions (lambdas) - if let Address::Global(global_index) = addr - && let BoundKind::Lambda { .. } = &value.kind - { - self.registry - .insert(*global_index, (*value).as_ref().clone()); + if let Address::Global(global_index) = addr { + let mut current = value; + while let BoundKind::Expansion { bound_expanded, .. } = ¤t.kind { + current = bound_expanded; + } + + if let BoundKind::Lambda { .. } = ¤t.kind { + self.registry + .insert(*global_index, (*current).as_ref().clone()); + } } self.visit(value); } BoundKind::Set { addr, value } => { // Also track assignments to globals if they hold lambdas. - if let Address::Global(global_index) = addr - && let BoundKind::Lambda { .. } = &value.kind - { - self.registry - .insert(*global_index, (*value).as_ref().clone()); + if let Address::Global(global_index) = addr { + let mut current = value; + while let BoundKind::Expansion { bound_expanded, .. } = ¤t.kind { + current = bound_expanded; + } + + if let BoundKind::Lambda { .. } = ¤t.kind { + self.registry + .insert(*global_index, (*current).as_ref().clone()); + } } self.visit(value); } diff --git a/src/integration_test.rs b/src/integration_test.rs index 41149aa..ffec9ad 100644 --- a/src/integration_test.rs +++ b/src/integration_test.rs @@ -401,4 +401,36 @@ mod tests { let res = env.run_script(source); assert_eq!(format!("{}", res.unwrap()), "3"); } + + #[test] + fn test_macro_inlining_identity_collision() { + let source = r#" + (do + (macro wrap [f] `(fn [x] (~f x))) + + (def add1 (fn [x] (+ x 1))) + (def add2 (fn [x] (+ x 2))) + + (def w1 (wrap add1)) + (def w2 (wrap add2)) + + (w1 (w2 10))) + "#; + + // 1. Verify the result is correct + let env_run = Environment::new(); + let res = env_run.run_script(source).expect("Failed to run script"); + assert_eq!(format!("{}", res), "13"); + + // 2. Verify that it was actually folded into a constant by the optimizer + let env_dump = Environment::new(); + let dump = env_dump.dump_ast(source).expect("Failed to dump AST"); + assert!( + dump.contains("Constant: 13"), + "Macro-wrapped calls should be fully folded to 13. Dump:\n{}", + dump + ); + // The definitions add1, add2, w1, w2 should be gone after dead code elimination + assert!(!dump.contains("Define Variable"), "Definitions should be removed by DCE"); + } }