Refactor: Remove explicit upvalue analysis pass

The `UpvalueAnalyzer` pass is no longer necessary as the binder now directly tracks captures. This commit removes the `UpvalueAnalyzer` struct and associated logic from `src/ast/compiler/upvalues.rs`. The `Binder::bind_root` function now returns the `captures` map, which is then processed by a new `CapturePass` in `src/ast/environment.rs` before type checking. This consolidates capture logic within the binder and its subsequent processing steps.
2026-02-25 12:15:59 +01:00
parent 82daf03522
commit 0371c21523
5 changed files with 92 additions and 204 deletions
@@ -8,6 +8,7 @@ use std::rc::Rc;
 #[derive(Debug, Clone)]
 struct LocalInfo {
    slot: u32,
+    identity: Identity,
    // Note: Binder doesn't strictly need the type anymore,
    // but it might be useful for built-ins during resolution.
    // For now we keep it as Any or Unknown.
@@ -28,7 +29,7 @@ impl CompilerScope {
        }
    }

-    fn define(&mut self, sym: &Symbol) -> Result<u32, String> {
+    fn define(&mut self, sym: &Symbol, identity: Identity) -> Result<u32, String> {
        if self.locals.contains_key(sym) {
            return Err(format!(
                "Variable '{}' is already defined in this scope level.",
@@ -40,6 +41,7 @@ impl CompilerScope {
            sym.clone(),
            LocalInfo {
                slot,
+                identity,
                _ty: StaticType::Any,
            },
        );
@@ -59,14 +61,16 @@ enum ScopeKind {
 }

 struct FunctionCompiler {
+    identity: Identity,
    scope: CompilerScope,
    upvalues: Vec<Address>,
    kind: ScopeKind,
 }

 impl FunctionCompiler {
-    fn new(kind: ScopeKind) -> Self {
+    fn new(kind: ScopeKind, identity: Identity) -> Self {
        Self {
+            identity,
            scope: CompilerScope::new(),
            upvalues: Vec::new(),
            kind,
@@ -76,7 +80,8 @@ impl FunctionCompiler {
    fn define_variable(
        &mut self,
        name: &Symbol,
-        globals: &Rc<RefCell<HashMap<Symbol, u32>>>,
+        identity: Identity,
+        globals: &Rc<RefCell<HashMap<Symbol, (u32, Identity)>>>,
    ) -> Result<Address, String> {
        match self.kind {
            ScopeKind::Root => {
@@ -88,11 +93,11 @@ impl FunctionCompiler {
                    ));
                }
                let idx = globals_map.len() as u32;
-                globals_map.insert(name.clone(), idx);
+                globals_map.insert(name.clone(), (idx, identity));
                Ok(Address::Global(idx))
            }
            ScopeKind::Local => {
-                let slot = self.scope.define(name)?;
+                let slot = self.scope.define(name, identity)?;
                Ok(Address::Local(slot))
            }
        }
@@ -110,48 +115,53 @@ impl FunctionCompiler {

 pub struct Binder {
    functions: Vec<FunctionCompiler>,
-    // Globals mapping: Symbol -> Index
-    globals: Rc<RefCell<HashMap<Symbol, u32>>>,
+    // Globals mapping: Symbol -> (Index, DefinitionIdentity)
+    globals: Rc<RefCell<HashMap<Symbol, (u32, Identity)>>>,
    // Map of Declaration Identity -> List of Lambda Identities that capture it
-    capture_map: HashMap<Identity, Vec<Identity>>,
+    capture_map: HashMap<Identity, std::collections::HashSet<Identity>>,
 }

 impl Binder {
-    pub fn new(globals: Rc<RefCell<HashMap<Symbol, u32>>>) -> Self {
-        Self::with_boxed(globals, HashMap::new())
-    }
-
-    fn with_boxed(
-        globals: Rc<RefCell<HashMap<Symbol, u32>>>,
-        captures: HashMap<Identity, Vec<Identity>>,
-    ) -> Self {
+    pub fn new(globals: Rc<RefCell<HashMap<Symbol, (u32, Identity)>>>) -> Self {
        let mut binder = Self {
            functions: Vec::new(),
            globals,
-            capture_map: captures,
+            capture_map: HashMap::new(),
        };
-        binder
-            .functions
-            .push(FunctionCompiler::new(ScopeKind::Root));
+        binder.functions.push(FunctionCompiler::new(
+            ScopeKind::Root,
+            Rc::new(crate::ast::types::NodeIdentity {
+                location: crate::ast::types::SourceLocation { line: 0, col: 0 },
+            }),
+        ));
        binder
    }

    pub fn bind_root(
-        globals: Rc<RefCell<HashMap<Symbol, u32>>>,
+        globals: Rc<RefCell<HashMap<Symbol, (u32, Identity)>>>,
        node: &Node<UntypedKind>,
-    ) -> Result<BoundNode, String> {
-        let captures = crate::ast::compiler::upvalues::UpvalueAnalyzer::analyze(node);
-        let mut binder = Self::with_boxed(globals, captures);
-        binder.bind(node)
+    ) -> Result<(BoundNode, HashMap<Identity, Vec<Identity>>), String> {
+        let mut binder = Self::new(globals);
+        let bound = binder.bind(node)?;
+
+        // Convert HashSet to sorted Vec
+        let final_captures = binder
+            .capture_map
+            .into_iter()
+            .map(|(k, v)| (k, v.into_iter().collect()))
+            .collect();
+
+        Ok((bound, final_captures))
    }

    fn declare_variable(
        &mut self,
        name: &Symbol,
+        identity: Identity,
        _kind: crate::ast::compiler::bound_nodes::DeclarationKind,
    ) -> Result<Address, String> {
        let current_fn = self.functions.last_mut().unwrap();
-        current_fn.define_variable(name, &self.globals)
+        current_fn.define_variable(name, identity, &self.globals)
    }

    pub fn bind(&mut self, node: &Node<UntypedKind>) -> Result<BoundNode, String> {
@@ -203,14 +213,10 @@ impl Binder {
                if let UntypedKind::Parameter(ref name) = target.kind {
                    let addr = self.declare_variable(
                        name,
+                        node.identity.clone(), // Identity of the Def node
                        crate::ast::compiler::bound_nodes::DeclarationKind::Variable,
                    )?;
                    let val_node = self.bind(value)?;
-                    let captured_by = self
-                        .capture_map
-                        .get(&target.identity)
-                        .cloned()
-                        .unwrap_or_default();

                    Ok(self.make_node(
                        node.identity.clone(),
@@ -219,7 +225,7 @@ impl Binder {
                            addr,
                            kind: crate::ast::compiler::bound_nodes::DeclarationKind::Variable,
                            value: Box::new(val_node),
-                            captured_by,
+                            captured_by: Vec::new(), // Will be filled by post-pass
                        },
                    ))
                } else {
@@ -265,7 +271,10 @@ impl Binder {

            UntypedKind::Lambda { params, body } => {
                let identity = node.identity.clone();
-                self.functions.push(FunctionCompiler::new(ScopeKind::Local));
+                self.functions.push(FunctionCompiler::new(
+                    ScopeKind::Local,
+                    identity.clone(),
+                ));

                // 1. Bind the parameter pattern/tuple
                let params_bound = self.bind_pattern(params, DeclarationKind::Parameter)?;
@@ -409,7 +418,13 @@ impl Binder {
            if let Some(info) = self.functions[i].scope.resolve(sym) {
                let mut addr = Address::Local(info.slot);

+                // Record the capture for each lambda level in between
                for k in (i + 1)..=current_fn_idx {
+                    let lambda_id = self.functions[k].identity.clone();
+                    self.capture_map
+                        .entry(info.identity.clone())
+                        .or_default()
+                        .insert(lambda_id);
                    addr = Address::Upvalue(self.functions[k].add_upvalue(addr));
                }
                return Ok(addr);
@@ -418,7 +433,7 @@ impl Binder {

        // 3. Try Global
        let globals = self.globals.borrow();
-        if let Some(idx) = globals.get(sym) {
+        if let Some((idx, _)) = globals.get(sym) {
            return Ok(Address::Global(*idx));
        }

@@ -428,7 +443,7 @@ impl Binder {
                name: sym.name.clone(),
                context: None,
            };
-            if let Some(idx) = globals.get(&fallback_sym) {
+            if let Some((idx, _)) = globals.get(&fallback_sym) {
                return Ok(Address::Global(*idx));
            }
        }
@@ -443,12 +458,7 @@ impl Binder {
    ) -> Result<BoundNode, String> {
        match &node.kind {
            UntypedKind::Parameter(sym) => {
-                let addr = self.declare_variable(sym, kind)?;
-                let captured_by = self
-                    .capture_map
-                    .get(&node.identity)
-                    .cloned()
-                    .unwrap_or_default();
+                let addr = self.declare_variable(sym, node.identity.clone(), kind)?;

                Ok(self.make_node(
                    node.identity.clone(),
@@ -457,7 +467,7 @@ impl Binder {
                        addr,
                        kind,
                        value: Box::new(self.make_node(node.identity.clone(), BoundKind::Nop)),
-                        captured_by,
+                        captured_by: Vec::new(), // Filled by post-pass
                    },
                ))
            }
@@ -530,19 +540,16 @@ mod tests {
        let untyped = parser.parse_expression().unwrap();

        let globals = Rc::new(RefCell::new(HashMap::new()));
-        let bound = Binder::bind_root(globals, &untyped).unwrap();
+        let (bound, captures) = Binder::bind_root(globals, &untyped).unwrap();

        // Structure: Lambda -> Block -> [ Define(x), Define(f), Get(x) ]
        if let BoundKind::Lambda { body, .. } = &bound.kind {
            if let BoundKind::Block { exprs } = &body.kind {
                let x_decl = &exprs[0];
-                if let BoundKind::Define {
-                    captured_by, addr, ..
-                } = &x_decl.kind
-                {
+                if let BoundKind::Define { addr, .. } = &x_decl.kind {
                    assert!(matches!(addr, Address::Local(_)));
                    assert!(
-                        !captured_by.is_empty(),
+                        captures.contains_key(&x_decl.identity),
                        "Variable 'x' should have capturers because it is used in lambda 'f'"
                    );
                } else {
@@ -566,18 +573,15 @@ mod tests {
        let untyped = parser.parse_expression().unwrap();

        let globals = Rc::new(RefCell::new(HashMap::new()));
-        let bound = Binder::bind_root(globals, &untyped).unwrap();
+        let (bound, captures) = Binder::bind_root(globals, &untyped).unwrap();

        if let BoundKind::Lambda { body, .. } = &bound.kind {
            if let BoundKind::Block { exprs } = &body.kind {
                let x_decl = &exprs[0];
-                if let BoundKind::Define {
-                    captured_by, addr, ..
-                } = &x_decl.kind
-                {
+                if let BoundKind::Define { addr, .. } = &x_decl.kind {
                    assert!(matches!(addr, Address::Local(_)));
                    assert!(
-                        captured_by.is_empty(),
+                        !captures.contains_key(&x_decl.identity),
                        "Variable 'x' should NOT have any capturers"
                    );
                } else {
@@ -713,7 +713,15 @@ mod tests {
        let expanded = expander.expand(untyped).unwrap();

        let mut global_names = HashMap::new();
-        global_names.insert(Symbol::from("*"), 0);
+        global_names.insert(
+            Symbol::from("*"),
+            (
+                0,
+                Rc::new(crate::ast::types::NodeIdentity {
+                    location: crate::ast::types::SourceLocation { line: 0, col: 0 },
+                }),
+            ),
+        );
        let globals = Rc::new(RefCell::new(global_names));

        let result = Binder::bind_root(globals, &expanded);
@@ -740,8 +748,8 @@ mod tests {
        let expanded = expander.expand(untyped).unwrap();

        let globals = Rc::new(RefCell::new(HashMap::new()));
-        let bound = Binder::bind_root(globals, &expanded);
-        assert!(bound.is_ok(), "Hygiene failed: {:?}", bound.err());
+        let result = Binder::bind_root(globals, &expanded);
+        assert!(result.is_ok(), "Hygiene failed: {:?}", result.err());
    }

    #[test]
@@ -760,11 +768,11 @@ mod tests {
        let expanded = expander.expand(untyped).unwrap();

        let globals = Rc::new(RefCell::new(HashMap::new()));
-        let bound = Binder::bind_root(globals, &expanded);
+        let result = Binder::bind_root(globals, &expanded);
        assert!(
-            bound.is_ok(),
+            result.is_ok(),
            "Explicit Hygiene with Backticks failed: {:?}",
-            bound.err()
+            result.err()
        );
    }
 }
@@ -1,6 +1,7 @@
 pub mod analyzer;
 pub mod binder;
 pub mod bound_nodes;
+pub mod captures;
 pub mod dumper;
 pub mod lambda_collector;
 pub mod macros;
@@ -8,14 +9,13 @@ pub mod optimizer;
 pub mod specializer;
 pub mod tco;
 pub mod type_checker;
-pub mod upvalues;

 pub use binder::*;
 pub use bound_nodes::*;
+pub use captures::*;
 pub use dumper::*;
 pub use macros::*;
 pub use optimizer::*;
 pub use specializer::*;
 pub use tco::*;
 pub use type_checker::*;
-pub use upvalues::*;
@@ -1,135 +0,0 @@
-use crate::ast::nodes::{Node, Symbol, UntypedKind};
-use crate::ast::types::Identity;
-use std::collections::{HashMap, HashSet};
-
-/// Analyzes the AST to find which lambdas capture which variable declarations.
-/// Returns a map: Declaration Identity -> List of Lambda Identities that capture it.
-pub struct UpvalueAnalyzer;
-
-impl UpvalueAnalyzer {
-    pub fn analyze(root: &Node<UntypedKind>) -> HashMap<Identity, Vec<Identity>> {
-        let mut capture_map: HashMap<Identity, HashSet<Identity>> = HashMap::new();
-        let mut scopes = vec![HashMap::new()]; // Root scope
-
-        Self::visit(root, &mut scopes, &mut capture_map, None);
-
-        // Convert HashSet back to sorted Vec for the final result
-        capture_map
-            .into_iter()
-            .map(|(decl, lambdas)| (decl, lambdas.into_iter().collect()))
-            .collect()
-    }
-
-    fn visit_params(node: &Node<UntypedKind>, scope: &mut HashMap<Symbol, Identity>) {
-        match &node.kind {
-            UntypedKind::Parameter(sym) => {
-                scope.insert(sym.clone(), node.identity.clone());
-            }
-            UntypedKind::Tuple { elements } => {
-                for el in elements {
-                    Self::visit_params(el, scope);
-                }
-            }
-            _ => {} // Ignore other nodes in parameter patterns for now
-        }
-    }
-
-    fn visit(
-        node: &Node<UntypedKind>,
-        scopes: &mut Vec<HashMap<Symbol, Identity>>,
-        capture_map: &mut HashMap<Identity, HashSet<Identity>>,
-        current_lambda: Option<Identity>,
-    ) {
-        match &node.kind {
-            UntypedKind::Identifier(sym) => {
-                // Resolve name in scope stack (from inner to outer)
-                for (depth, scope) in scopes.iter().rev().enumerate() {
-                    if let Some(decl_id) = scope.get(sym) {
-                        if depth > 0 {
-                            // Captured from an outer scope!
-                            if let Some(lambda_id) = &current_lambda {
-                                capture_map
-                                    .entry(decl_id.clone())
-                                    .or_default()
-                                    .insert(lambda_id.clone());
-                            }
-                        }
-                        break;
-                    }
-                }
-            }
-            UntypedKind::Def { target, value } => {
-                Self::visit(value, scopes, capture_map, current_lambda.clone());
-                if let Some(current) = scopes.last_mut() {
-                    Self::visit_params(target, current);
-                }
-            }
-            UntypedKind::Lambda { params, body } => {
-                let mut new_scope = HashMap::new();
-                Self::visit_params(params, &mut new_scope);
-
-                scopes.push(new_scope);
-                // The current node is the lambda causing captures in its body
-                Self::visit(body, scopes, capture_map, Some(node.identity.clone()));
-                scopes.pop();
-            }
-            UntypedKind::MacroDecl { params, body, .. } => {
-                let mut new_scope = HashMap::new();
-                Self::visit_params(params, &mut new_scope);
-                scopes.push(new_scope);
-                Self::visit(body, scopes, capture_map, current_lambda);
-                scopes.pop();
-            }
-            UntypedKind::If {
-                cond,
-                then_br,
-                else_br,
-            } => {
-                Self::visit(cond, scopes, capture_map, current_lambda.clone());
-                Self::visit(then_br, scopes, capture_map, current_lambda.clone());
-                if let Some(e) = else_br {
-                    Self::visit(e, scopes, capture_map, current_lambda.clone());
-                }
-            }
-            UntypedKind::Assign { target, value } => {
-                Self::visit(target, scopes, capture_map, current_lambda.clone());
-                Self::visit(value, scopes, capture_map, current_lambda.clone());
-            }
-            UntypedKind::Call { callee, args } => {
-                Self::visit(callee, scopes, capture_map, current_lambda.clone());
-                Self::visit(args, scopes, capture_map, current_lambda.clone());
-            }
-            UntypedKind::Again { args } => {
-                Self::visit(args, scopes, capture_map, current_lambda.clone());
-            }
-            UntypedKind::Block { exprs } => {
-                for expr in exprs {
-                    Self::visit(expr, scopes, capture_map, current_lambda.clone());
-                }
-            }
-            UntypedKind::Tuple { elements } => {
-                for el in elements {
-                    Self::visit(el, scopes, capture_map, current_lambda.clone());
-                }
-            }
-            UntypedKind::Record { fields } => {
-                for (k, v) in fields {
-                    Self::visit(k, scopes, capture_map, current_lambda.clone());
-                    Self::visit(v, scopes, capture_map, current_lambda.clone());
-                }
-            }
-            UntypedKind::Expansion { call: _, expanded } => {
-                Self::visit(expanded, scopes, capture_map, current_lambda);
-            }
-            UntypedKind::Template(body)
-            | UntypedKind::Placeholder(body)
-            | UntypedKind::Splice(body) => {
-                Self::visit(body, scopes, capture_map, current_lambda);
-            }
-            UntypedKind::Nop
-            | UntypedKind::Constant(_)
-            | UntypedKind::Extension(_)
-            | UntypedKind::Parameter(_) => {}
-        }
-    }
-}
@@ -17,10 +17,12 @@ use crate::ast::compiler::specializer::{FunctionRegistry, MonoCache, Specializer
 use crate::ast::compiler::tco::{ExecNode, TCO};
 use crate::ast::rtl;
 use crate::ast::rtl::intrinsics;
-use crate::ast::types::{Object, Purity, StaticType, Value};
+use crate::ast::types::{
+    Identity, NodeIdentity, Object, Purity, SourceLocation, StaticType, Value,
+};

 pub struct Environment {
-    pub global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
+    pub global_names: Rc<RefCell<HashMap<Symbol, (u32, Identity)>>>,
    pub global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
    pub global_purity: Rc<RefCell<HashMap<u32, Purity>>>,
    pub global_values: Rc<RefCell<Vec<Value>>>,
@@ -55,7 +57,7 @@ impl FunctionRegistry for EnvFunctionRegistry {
 }

 struct RuntimeMacroEvaluator {
-    global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
+    global_names: Rc<RefCell<HashMap<Symbol, (u32, Identity)>>>,
    global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
    global_values: Rc<RefCell<Vec<Value>>>,
 }
@@ -72,7 +74,8 @@ impl MacroEvaluator for RuntimeMacroEvaluator {
            return Ok(Value::Object(Rc::new(arg_node.clone()) as Rc<dyn Object>));
        }

-        let bound_ast = Binder::bind_root(self.global_names.clone(), node)?;
+        let (bound_ast, captures) = Binder::bind_root(self.global_names.clone(), node)?;
+        let bound_ast = crate::ast::compiler::captures::CapturePass::apply(bound_ast, &captures);
        let checker = TypeChecker::new(self.global_types.clone());
        let typed_ast = checker.check(bound_ast, &[])?;
        let exec_ast = TCO::optimize(Analyzer::analyze(&typed_ast, &HashMap::new())); // Minimal analysis for macro eval
@@ -131,7 +134,10 @@ impl Environment {
        let mut purity = self.global_purity.borrow_mut();

        let idx = values.len() as u32;
-        names.insert(Symbol::from(name), idx);
+        let identity = Rc::new(NodeIdentity {
+            location: SourceLocation { line: 0, col: 0 },
+        });
+        names.insert(Symbol::from(name), (idx, identity));
        types.insert(idx, ty);
        purity.insert(idx, func.purity);
        values.push(Value::Function(func));
@@ -161,7 +167,10 @@ impl Environment {
        let mut purity = self.global_purity.borrow_mut();

        let idx = values.len() as u32;
-        names.insert(Symbol::from(name), idx);
+        let identity = Rc::new(NodeIdentity {
+            location: SourceLocation { line: 0, col: 0 },
+        });
+        names.insert(Symbol::from(name), (idx, identity));
        types.insert(idx, ty);
        purity.insert(idx, Purity::Pure);
        values.push(val);
@@ -186,7 +195,9 @@ impl Environment {
        }

        let expanded_ast = self.get_expander().expand(untyped_ast)?;
-        let bound_ast = Binder::bind_root(self.global_names.clone(), &expanded_ast)?;
+        let (bound_ast, captures) = Binder::bind_root(self.global_names.clone(), &expanded_ast)?;
+        let bound_ast = crate::ast::compiler::captures::CapturePass::apply(bound_ast, &captures);
+
        LambdaCollector::collect(&bound_ast, &mut self.function_registry.borrow_mut());

        let checker = TypeChecker::new(self.global_types.clone());