From 6c8e2df45cd09a5b123a3f62cf5c297bc58febfb Mon Sep 17 00:00:00 2001
From: Brummel <mould@mycelium.de>
Date: Sun, 29 Mar 2026 18:50:05 +0200
Subject: [PATCH] Refactor: Improve type propagation in optimizer

Refactor the `updated_metrics` function to `refreshed_metrics`. This
function now clones the `NodeMetrics` if the new type is identical to
the original, avoiding unnecessary allocations.

Also, adjust the `Block` and `Lambda` node kinds to correctly propagate
concrete types after inlining, especially when dealing with polymorphic
functions. This ensures that the return types of blocks and lambdas
reflect the actual computed types rather than stale type variables.

This change improves type inference accuracy and prepares for more
advanced optimizations by ensuring type information is correctly
propagated through the AST.
---
 src/ast/compiler/optimizer/engine.rs | 33 ++++++++++++++++++----------
 tests/optimizer.rs                   | 28 +++++++++++++++++++++++
 2 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/src/ast/compiler/optimizer/engine.rs b/src/ast/compiler/optimizer/engine.rs
index 145b4f4..f81ce08 100644
--- a/src/ast/compiler/optimizer/engine.rs
+++ b/src/ast/compiler/optimizer/engine.rs
@@ -2,7 +2,7 @@ use crate::ast::nodes::{
     Address, AnalyzedNode, AssignBinding, GlobalAnalyzedRegistry,
     IdentifierBinding, LambdaBinding, Node, NodeKind, NodeMetrics, UpvalueIdx, VirtualId,
 };
-use crate::ast::types::{Purity, StaticType, Value};
+use crate::ast::types::{Purity, Signature, StaticType, Value};
 use crate::ast::vm::GlobalStore;
 
 use std::cell::RefCell;
@@ -39,9 +39,12 @@ fn compute_tuple_type(elements: &[Rc<AnalyzedNode>]) -> StaticType {
     }
 }
 
-/// Build a new `NodeMetrics` with an updated `StaticType`, reusing the original
-/// node's identity, kind, and comments. Only call when the type actually changed.
-fn updated_metrics(old: &NodeMetrics, new_ty: StaticType) -> NodeMetrics {
+/// Return `NodeMetrics` with an updated `StaticType` if it differs from the
+/// original, otherwise clone cheaply. Avoids unnecessary `Rc` allocations.
+fn refreshed_metrics(old: &NodeMetrics, new_ty: StaticType) -> NodeMetrics {
+    if new_ty == old.original.ty {
+        return old.clone();
+    }
     let orig = &old.original;
     NodeMetrics {
         original: Rc::new(Node {
@@ -614,7 +617,9 @@ impl Optimizer {
                     return node_rc;
                 }
 
-                (NodeKind::Block { exprs: new_exprs }, node.ty.clone())
+                let last_ty = new_exprs.last().map(|e| e.ty.original.ty.clone())
+                    .unwrap_or(StaticType::Void);
+                (NodeKind::Block { exprs: new_exprs }, refreshed_metrics(&node.ty, last_ty))
             }
 
             NodeKind::Lambda {
@@ -679,6 +684,15 @@ impl Optimizer {
                     return node_rc;
                 }
 
+                let metrics = if let StaticType::Function(sig) = &node.ty.original.ty {
+                    refreshed_metrics(&node.ty, StaticType::Function(Box::new(Signature {
+                        params: sig.params.clone(),
+                        ret: reindexed_body.ty.original.ty.clone(),
+                    })))
+                } else {
+                    node.ty.clone()
+                };
+
                 (
                     NodeKind::Lambda {
                         params: params_opt,
@@ -688,7 +702,7 @@ impl Optimizer {
                             positional_count: lambda_info.positional_count,
                         },
                     },
-                    node.ty.clone(),
+                    metrics,
                 )
             }
 
@@ -705,12 +719,7 @@ impl Optimizer {
                 if !changed {
                     return node_rc;
                 }
-                let recomputed = compute_tuple_type(&new_elements);
-                let metrics = if recomputed != node.ty.original.ty {
-                    updated_metrics(&node.ty, recomputed)
-                } else {
-                    node.ty.clone()
-                };
+                let metrics = refreshed_metrics(&node.ty, compute_tuple_type(&new_elements));
                 (NodeKind::Tuple { elements: new_elements }, metrics)
             }
             NodeKind::Record { fields, layout } => {
diff --git a/tests/optimizer.rs b/tests/optimizer.rs
index c65b72e..cd46707 100644
--- a/tests/optimizer.rs
+++ b/tests/optimizer.rs
@@ -249,3 +249,31 @@ fn test_inlined_tuple_has_concrete_type_heterogeneous() {
         "Inlined tuple must not contain unresolved TypeVars. Dump:\n{dump}"
     );
 }
+
+#[test]
+fn test_block_and_lambda_propagate_concrete_type_after_inlining() {
+    // When a polymorphic function is inlined and the result tuple gets a concrete type,
+    // the enclosing Block and outer Lambda return type must propagate that concrete type
+    // upward — not retain stale TypeVars from before inlining.
+    let env = Environment::new();
+    let source = r#"(do
+        (def last (fn [s] (s 0)))
+        (def f (fn [n] n))
+        (def r (series 5))
+        (push r 4)
+        [(last f) (last r)]
+    )"#;
+    let dump = env.dump_ast_compact(source).unwrap();
+    // The outer Lambda and Block must carry concrete types, not stale TypeVars.
+    // Inner polymorphic lambdas (like `last`) legitimately retain TypeVars.
+    let first_line = dump.lines().next().unwrap();
+    assert!(
+        !first_line.contains('?'),
+        "Outer Lambda return type must be concrete. First line:\n{first_line}"
+    );
+    let block_line = dump.lines().find(|l| l.contains("Block")).unwrap();
+    assert!(
+        !block_line.contains('?'),
+        "Block type must be concrete. Line:\n{block_line}"
+    );
+}