RustAst/src/ast/compiler/optimizer/utils.rs

use crate::ast::nodes::{
    Address, AnalyzedNode, GlobalIdx, IdentifierBinding,
    NodeKind, VirtualId,
};
use crate::ast::types::{Identity, Value};
use crate::ast::vm::Closure;
use std::collections::HashSet;

// --- PathTracker ---
#[derive(Default)]
pub struct PathTracker {
    pub inlining_depth: usize,
    pub inlining_stack: HashSet<GlobalIdx>,
    pub identity_stack: HashSet<Identity>,
}

impl PathTracker {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn enter_lambda(&mut self, identity: &Identity) -> bool {
        if self.identity_stack.contains(identity) {
            return false;
        }
        self.identity_stack.insert(identity.clone());
        true
    }

    pub fn exit_lambda(&mut self, identity: &Identity) {
        self.identity_stack.remove(identity);
    }
}

// --- UsageInfo ---
#[derive(Default)]
pub struct UsageInfo {
    pub used: HashSet<Address<VirtualId>>,
    pub assigned: HashSet<Address<VirtualId>>,
    pub used_identities: HashSet<Identity>,
}

impl UsageInfo {
    pub fn is_assigned(&self, addr: &Address<VirtualId>) -> bool {
        self.assigned.contains(addr)
    }

    pub fn is_used(&self, addr: &Address<VirtualId>) -> bool {
        self.used.contains(addr)
    }

    pub fn collect(&mut self, node: &AnalyzedNode) {
        match &node.kind {
            NodeKind::Def { pattern, value, .. } => {
                self.collect_pattern(pattern);
                self.collect(value);
            }
            NodeKind::Constant(v) => {
                if let Value::Object(obj) = v
                    && let Some(closure) = obj.as_any().downcast_ref::<Closure>()
                {
                    self.used_identities
                        .insert(closure.function_node.identity.clone());
                    self.collect(&closure.function_node);
                }
            }
            NodeKind::Identifier { binding, .. } => {
                let addr = match binding {
                    IdentifierBinding::Reference(addr) => *addr,
                    IdentifierBinding::Declaration { addr, .. } => *addr,
                };
                self.used.insert(addr);
            }
            NodeKind::Assign { target, value, info, .. } => {
                if let Some(addr) = info.addr {
                    self.assigned.insert(addr);
                } else {
                    // Destructuring assign: collect assigned addresses from target pattern
                    self.collect_assigned_from_target(target);
                }
                self.collect(value);
            }
            NodeKind::GetField { rec, .. } => {
                self.collect(rec);
            }
            NodeKind::Lambda {
                params,
                body,
                info: lambda_info,
            } => {
                self.used_identities.insert(node.identity.clone());

                let mut inner_info = UsageInfo::default();
                inner_info.collect(params);
                inner_info.collect(body);

                // Propagate globals and identities
                for addr in &inner_info.used {
                    if let Address::Global(_) = addr {
                        self.used.insert(*addr);
                    }
                }
                for addr in &inner_info.assigned {
                    if let Address::Global(_) = addr {
                        self.assigned.insert(*addr);
                    }
                }
                self.used_identities.extend(inner_info.used_identities);

                // Map used upvalues to parent scope
                for addr in &inner_info.used {
                    if let Address::Upvalue(idx) = addr
                        && let Some(parent_addr) = lambda_info.upvalues.get(idx.0 as usize)
                    {
                        self.used.insert(*parent_addr);
                    }
                }
                // Map assigned upvalues to parent scope
                for addr in &inner_info.assigned {
                    if let Address::Upvalue(idx) = addr
                        && let Some(parent_addr) = lambda_info.upvalues.get(idx.0 as usize)
                    {
                        self.assigned.insert(*parent_addr);
                    }
                }
            }
            NodeKind::Block { exprs } => {
                for e in exprs {
                    self.collect(e);
                }
            }
            NodeKind::If {
                cond,
                then_br,
                else_br,
            } => {
                self.collect(cond);
                self.collect(then_br);
                if let Some(e) = else_br {
                    self.collect(e);
                }
            }
            NodeKind::Call { callee, args } => {
                self.collect(callee);
                self.collect(args);
            }
            NodeKind::Tuple { elements } => {
                for e in elements {
                    self.collect(e);
                }
            }
            NodeKind::Record { fields, .. } => {
                for (_, v) in fields {
                    self.collect(v);
                }
            }
            NodeKind::Expansion { expanded, .. } => {
                self.collect(expanded);
            }
            NodeKind::Again { args } => {
                self.collect(args);
            }
            NodeKind::Pipe { inputs, lambda, .. } => {
                for input in inputs {
                    self.collect(input);
                }
                self.collect(lambda);
            }
            NodeKind::Nop
            | NodeKind::FieldAccessor(_)
            | NodeKind::Extension(_)
            | NodeKind::Error => {}
            // Syntax-only variants that should not appear in AnalyzedPhase
            NodeKind::MacroDecl { .. }
            | NodeKind::Template(_)
            | NodeKind::Placeholder(_)
            | NodeKind::Splice(_) => {}
        }
    }

    pub fn collect_pattern(&mut self, node: &AnalyzedNode) {
        match &node.kind {
            NodeKind::Identifier {
                binding: IdentifierBinding::Declaration { .. },
                ..
            } => {}
            NodeKind::Def { .. } => {}
            NodeKind::Assign { info, .. } => {
                if let Some(addr) = info.addr {
                    self.assigned.insert(addr);
                }
            }
            NodeKind::Tuple { elements } => {
                for el in elements {
                    self.collect_pattern(el);
                }
            }
            _ => {}
        }
    }

    pub fn collect_assigned_from_target(&mut self, node: &AnalyzedNode) {
        match &node.kind {
            NodeKind::Identifier { binding, .. } => {
                let addr = match binding {
                    IdentifierBinding::Reference(addr)
                    | IdentifierBinding::Declaration { addr, .. } => *addr,
                };
                self.assigned.insert(addr);
            }
            NodeKind::Tuple { elements } => {
                for el in elements {
                    self.collect_assigned_from_target(el);
                }
            }
            _ => {}
        }
    }
}

/// Collects all addresses declared in a binding pattern (Identifier Declaration or nested Tuple).
///
/// Used by the optimizer to enumerate all slots bound by a destructuring `def` pattern,
/// enabling dead-def elimination for patterns that `extract_def_addr` cannot handle
/// (i.e. anything other than a plain `Identifier`).
pub fn collect_pattern_addrs(pattern: &AnalyzedNode, out: &mut Vec<Address<VirtualId>>) {
    match &pattern.kind {
        NodeKind::Identifier {
            binding: IdentifierBinding::Declaration { addr, .. },
            ..
        } => {
            out.push(*addr);
        }
        NodeKind::Def { pattern: inner, .. } => {
            collect_pattern_addrs(inner, out);
        }
        NodeKind::Tuple { elements } => {
            for el in elements {
                collect_pattern_addrs(el, out);
            }
        }
        _ => {}
    }
}