Handle macro expansion in lambda collection

Recursively traverse through `BoundKind::Expansion` nodes when
collecting lambdas to ensure that macros correctly expanded into lambdas
are registered. This also updates assignments to global variables to
correctly track lambdas that have been assigned after macro expansion.

Added a new integration test to verify that macro-wrapped function calls
are correctly inlined and optimized, and that unused definitions are
removed by dead code elimination.

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
+                if let Address::Global(global_index) = addr {
-                    && let BoundKind::Lambda { .. } = &value.kind
+                    let mut current = value;
-                {
+                    while let BoundKind::Expansion { bound_expanded, .. } = &current.kind {
-                    self.registry
+                        current = bound_expanded;
-                        .insert(*global_index, (*value).as_ref().clone());
+                    }
+
+                    if let BoundKind::Lambda { .. } = &current.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
+                if let Address::Global(global_index) = addr {
-                    && let BoundKind::Lambda { .. } = &value.kind
+                    let mut current = value;
-                {
+                    while let BoundKind::Expansion { bound_expanded, .. } = &current.kind {
-                    self.registry
+                        current = bound_expanded;
-                        .insert(*global_index, (*value).as_ref().clone());
+                    }
+
+                    if let BoundKind::Lambda { .. } = &current.kind {
+                        self.registry
+                            .insert(*global_index, (*current).as_ref().clone());
+                    }
                 }
                 self.visit(value);
             }

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");
+    }
 }