RustAst/src/ast/compiler/upvalues.rs

use std::collections::{HashMap, HashSet};
use crate::ast::nodes::{Node, UntypedKind, Symbol};
use crate::ast::types::Identity;

/// 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 { name, value } => {
                Self::visit(value, scopes, capture_map, current_lambda.clone());
                if let Some(current) = scopes.last_mut() {
                    current.insert(name.clone(), node.identity.clone());
                }
            }
            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::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(_) => {}
        }
    }
}