RustAst/src/ast/compiler/lambda_collector.rs

use crate::ast::nodes::{
    Address, BoundLike, GlobalIdx, IdentifierBinding, Node, NodeKind,
};
use std::collections::HashMap;
use std::rc::Rc;

/// A pass that collects all global function definitions (lambdas) into a registry.
/// This allows the Specializer to retrieve the original AST of a function for monomorphization.
pub struct LambdaCollector<'a, P: BoundLike> {
    registry: &'a mut HashMap<GlobalIdx, Rc<Node<P>>>,
}

impl<'a, P> LambdaCollector<'a, P>
where
    P: BoundLike,
{
    /// Performs a full traversal of the AST and populates the provided registry.
    pub fn collect(node: &Node<P>, registry: &'a mut HashMap<GlobalIdx, Rc<Node<P>>>) {
        let mut collector = Self { registry };
        collector.visit(node);
    }

    fn visit(&mut self, node: &Node<P>) {
        match &node.kind {
            NodeKind::Block { exprs } => {
                for expr in exprs {
                    self.visit(expr);
                }
            }

            NodeKind::Def { pattern, value, .. } => {
                // Register global function definitions (lambdas)
                if let NodeKind::Identifier {
                    binding: IdentifierBinding::Declaration {
                        addr: Address::Global(global_index),
                        ..
                    },
                    ..
                } = &pattern.kind
                {
                    let mut current = value;
                    while let NodeKind::Expansion { expanded, .. } = &current.kind {
                        current = expanded;
                    }

                    if let NodeKind::Lambda { .. } = &current.kind {
                        self.registry
                            .insert(*global_index, (*current).clone());
                    }
                }
                self.visit(value);
            }

            NodeKind::Assign { value, info, .. } => {
                // Also track assignments to globals if they hold lambdas.
                if let Some(Address::Global(global_index)) = &info.addr {
                    let mut current = value;
                    while let NodeKind::Expansion { expanded, .. } = &current.kind {
                        current = expanded;
                    }

                    if let NodeKind::Lambda { .. } = &current.kind {
                        self.registry
                            .insert(*global_index, (*current).clone());
                    }
                }
                self.visit(value);
            }

            NodeKind::If {
                cond,
                then_br,
                else_br,
            } => {
                self.visit(cond);
                self.visit(then_br);
                if let Some(e) = else_br {
                    self.visit(e);
                }
            }

            NodeKind::Lambda { params, body, .. } => {
                self.visit(params);
                self.visit(body);
            }

            NodeKind::Call { callee, args } => {
                self.visit(callee);
                self.visit(args);
            }

            NodeKind::Tuple { elements } => {
                for el in elements {
                    self.visit(el);
                }
            }

            NodeKind::Record { fields, .. } => {
                for (_, v) in fields {
                    self.visit(v);
                }
            }

            NodeKind::Expansion { expanded, .. } => {
                self.visit(expanded);
            }

            _ => {} // Leaf nodes
        }
    }
}