From 11f6115fc17ec7b4bbcb710f98d7bf52cce827d9 Mon Sep 17 00:00:00 2001
From: Michael Schimmel <schimmel@mycelium.de>
Date: Wed, 25 Feb 2026 13:47:50 +0100
Subject: [PATCH] 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 | 30 +++++++++++++++++---------
 src/integration_test.rs              | 32 ++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 10 deletions(-)

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, .. } = &current.kind {
+                        current = bound_expanded;
+                    }
+
+                    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
-                    && 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, .. } = &current.kind {
+                        current = bound_expanded;
+                    }
+
+                    if let BoundKind::Lambda { .. } = &current.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");
+    }
 }