RustAst/src/ast/compiler/optimizer.rs

use std::rc::Rc;
use crate::ast::compiler::bound_nodes::{BoundKind, TypedNode, Address};
use crate::ast::types::{Value, StaticType};
use crate::ast::vm::Closure;
use crate::ast::nodes::Node;
use std::collections::HashMap;

/// The Optimizer performs Phase 2 (Cracking) and Phase 2.5 (Aggressive Collapsing).
pub struct Optimizer {
    pub level: u32,
    max_passes: usize,
}

impl Optimizer {
    pub fn new(level: u32) -> Self {
        Self { level, max_passes: 5 }
    }

    pub fn optimize(&self, node: TypedNode) -> TypedNode {
        if self.level == 0 {
            return node;
        }
        
        let mut current = node;
        for _ in 0..self.max_passes {
            let mut sub = SubstitutionMap::new();
            let next = self.visit_node(current.clone(), &mut sub);
            if next == current {
                break;
            }
            current = next;
        }
        current
    }

    fn visit_node(&self, node: TypedNode, sub: &mut SubstitutionMap) -> TypedNode {
        let (new_kind, new_ty) = match node.kind {
            BoundKind::Get { addr, name } => {
                match addr {
                    Address::Local(slot) => {
                        if let Some(val) = sub.locals.get(&slot) {
                            return Node {
                                identity: node.identity,
                                ty: val.static_type(),
                                kind: BoundKind::Constant(val.clone()),
                            };
                        }
                    }
                    Address::Upvalue(idx) => {
                        if let Some(val) = sub.upvalues.get(&idx) {
                            return Node {
                                identity: node.identity,
                                ty: val.static_type(),
                                kind: BoundKind::Constant(val.clone()),
                            };
                        }
                    }
                    _ => {}
                }
                (BoundKind::Get { addr, name }, node.ty)
            }

            BoundKind::Call { callee, args } => {
                let callee = self.visit_node(*callee, sub);
                let args = self.visit_node(*args, sub);
                
                if self.level >= 2 {
                    // Case 1: Beta-Reduction for Lambda Literals
                    if let BoundKind::Lambda { params, body, .. } = &callee.kind
                        && let Some(collapsed) = self.try_beta_reduce(params, &args, (**body).clone()) {
                            return self.visit_node(collapsed, sub);
                        }
                    
                    // Case 2: Cracking and Inlining for Constant Closures
                    if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
                        && let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
                            let mut closure_sub = SubstitutionMap::new();
                            for (i, cell) in closure.upvalues.iter().enumerate() {
                                closure_sub.add_upvalue(i as u32, cell.borrow().clone());
                            }
                            let inlined_body = self.visit_node((*closure.function_node).clone(), &mut closure_sub);
                            
                            if let Some(collapsed) = self.try_beta_reduce(&closure.parameter_node, &args, inlined_body) {
                                return self.visit_node(collapsed, sub);
                            }
                        }

                    if let Some(folded) = self.try_fold_intrinsic(&callee, &args) {
                        return folded;
                    }
                }

                if self.level >= 1
                    && let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
                        && let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
                            let mut closure_sub = SubstitutionMap::new();
                            for (i, cell) in closure.upvalues.iter().enumerate() {
                                closure_sub.add_upvalue(i as u32, cell.borrow().clone());
                            }
                            let inlined_body = self.visit_node((*closure.function_node).clone(), &mut closure_sub);
                            
                            let cracked_lambda = Node {
                                identity: callee.identity.clone(),
                                ty: callee.ty.clone(),
                                kind: BoundKind::Lambda {
                                    params: closure.parameter_node.clone(),
                                    upvalues: vec![],
                                    body: Rc::new(inlined_body),
                                    positional_count: closure.positional_count,
                                },
                            };
                            return Node {
                                identity: node.identity,
                                kind: BoundKind::Call { callee: Box::new(cracked_lambda), args: Box::new(args) },
                                ty: node.ty,
                            };
                        }

                (BoundKind::Call { callee: Box::new(callee), args: Box::new(args) }, node.ty)
            },

            BoundKind::If { cond, then_br, else_br } => {
                let cond = self.visit_node(*cond, sub);
                if self.level >= 2
                    && let BoundKind::Constant(ref val) = cond.kind {
                        if val.is_truthy() {
                            return self.visit_node(*then_br, sub);
                        } else if let Some(else_node) = else_br {
                            return self.visit_node(*else_node, sub);
                        } else {
                            return Node { identity: node.identity, kind: BoundKind::Nop, ty: StaticType::Void };
                        }
                    }
                
                let then_br = Box::new(self.visit_node(*then_br, sub));
                let else_br = else_br.map(|e| Box::new(self.visit_node(*e, sub)));
                (BoundKind::If { cond: Box::new(cond), then_br, else_br }, node.ty)
            },

            BoundKind::Block { exprs } => {
                let mut new_exprs = Vec::with_capacity(exprs.len());
                for e in exprs {
                    let opt = self.visit_node(e, sub);
                    if self.level >= 2 && matches!(opt.kind, BoundKind::Nop) {
                        continue;
                    }
                    new_exprs.push(opt);
                }
                
                if self.level >= 2 {
                    if new_exprs.is_empty() {
                        return Node { identity: node.identity, kind: BoundKind::Nop, ty: StaticType::Void };
                    } else if new_exprs.len() == 1 {
                        return new_exprs.pop().unwrap();
                    }
                }
                
                let ty = if new_exprs.is_empty() { StaticType::Void } else { new_exprs.last().unwrap().ty.clone() };
                (BoundKind::Block { exprs: new_exprs }, ty)
            },

            BoundKind::Lambda { params, upvalues: original_upvalues, body, positional_count } => {
                let mut new_upvalues = Vec::new();
                let mut mapping = Vec::new();
                let mut next_inner_subs = SubstitutionMap::new();

                for (old_idx, capture_addr) in original_upvalues.iter().enumerate() {
                    let mut inlined_val = None;
                    match capture_addr {
                        Address::Local(slot) => {
                            if let Some(val) = sub.locals.get(slot) { inlined_val = Some(val.clone()); }
                        }
                        Address::Upvalue(idx) => {
                            if let Some(val) = sub.upvalues.get(idx) { inlined_val = Some(val.clone()); }
                        }
                        _ => {}
                    }

                    if let Some(val) = inlined_val {
                        next_inner_subs.add_upvalue(old_idx as u32, val);
                        mapping.push(None);
                    } else {
                        mapping.push(Some(new_upvalues.len() as u32));
                        new_upvalues.push(*capture_addr);
                    }
                }

                let params = Rc::new(self.visit_node(params.as_ref().clone(), &mut SubstitutionMap::new()));
                let body_node = self.visit_node((*body).clone(), &mut next_inner_subs);
                
                let reindexed_body = if new_upvalues.len() != original_upvalues.len() {
                    sub.reindex_upvalues(body_node, &mapping)
                } else {
                    body_node
                };

                (BoundKind::Lambda { 
                    params, 
                    upvalues: new_upvalues, 
                    body: Rc::new(reindexed_body), 
                    positional_count 
                }, node.ty)
            },
            
            BoundKind::DefLocal { name, slot, value, captured_by } => {
                let value = Box::new(self.visit_node(*value, sub));
                if let BoundKind::Constant(val) = &value.kind {
                    sub.add_local(slot, val.clone());
                } else {
                    sub.locals.remove(&slot);
                }
                (BoundKind::DefLocal { name, slot, value, captured_by }, node.ty)
            },
            BoundKind::DefGlobal { name, global_index, value } => {
                let value = Box::new(self.visit_node(*value, sub));
                (BoundKind::DefGlobal { name, global_index, value }, node.ty)
            },
            BoundKind::Set { addr, value } => {
                let value = Box::new(self.visit_node(*value, sub));
                if let Address::Local(slot) = addr {
                    if let BoundKind::Constant(val) = &value.kind {
                        sub.add_local(slot, val.clone());
                    } else {
                        sub.locals.remove(&slot);
                    }
                }
                (BoundKind::Set { addr, value }, node.ty)
            },
            BoundKind::Tuple { elements } => {
                let elements = elements.into_iter().map(|e| self.visit_node(e, sub)).collect();
                (BoundKind::Tuple { elements }, node.ty)
            },
            BoundKind::Record { fields } => {
                let fields = fields.into_iter().map(|(k, v)| (self.visit_node(k, sub), self.visit_node(v, sub))).collect();
                (BoundKind::Record { fields }, node.ty)
            },
            BoundKind::Expansion { original_call, bound_expanded } => {
                let bound_expanded = Box::new(self.visit_node(*bound_expanded, sub));
                (BoundKind::Expansion { original_call, bound_expanded }, node.ty)
            },
            k => (k, node.ty),
        };

        Node { identity: node.identity, kind: new_kind, ty: new_ty }
    }

    fn try_beta_reduce(&self, params: &TypedNode, args: &TypedNode, body: TypedNode) -> Option<TypedNode> {
        if self.contains_def_local(&body) {
            return None;
        }

        let mut sub = SubstitutionMap::new();
        let mut arg_vals = Vec::new();
        self.flatten_typed_tuple(args, &mut arg_vals);
        
        let mut slot_index = 0;
        self.map_params_to_args(params, &arg_vals, &mut slot_index, &mut sub);

        if sub.locals.len() < arg_vals.len() {
            return None;
        }

        if sub.locals.is_empty() && sub.upvalues.is_empty() && !arg_vals.is_empty() {
            return None;
        }

        Some(self.visit_node(body, &mut sub))
    }

    fn flatten_typed_tuple(&self, node: &TypedNode, into: &mut Vec<TypedNode>) {
        if let BoundKind::Tuple { elements } = &node.kind {
            for el in elements {
                self.flatten_typed_tuple(el, into);
            }
        } else if !matches!(node.kind, BoundKind::Nop) {
            into.push(node.clone());
        }
    }

    fn map_params_to_args(&self, pattern: &TypedNode, args: &[TypedNode], offset: &mut usize, sub: &mut SubstitutionMap) {
        match &pattern.kind {
            BoundKind::Parameter { slot, .. } => {
                if let Some(arg) = args.get(*offset)
                    && let BoundKind::Constant(val) = &arg.kind {
                        sub.add_local(*slot, val.clone());
                    }
                *offset += 1;
            }
            BoundKind::Tuple { elements } => {
                for el in elements {
                    self.map_params_to_args(el, args, offset, sub);
                }
            }
            _ => {}
        }
    }

    fn try_fold_intrinsic(&self, callee: &TypedNode, args: &TypedNode) -> Option<TypedNode> {
        if let BoundKind::Get { name, .. } = &callee.kind
            && let BoundKind::Tuple { elements } = &args.kind
            && elements.len() == 2
        {
            let val_a = self.get_const_int(&elements[0]);
            let val_b = self.get_const_int(&elements[1]);
            
            if let (Some(a), Some(b)) = (val_a, val_b) {
                let res = match &*name.name {
                    "+" => Some(Value::Int(a + b)),
                    "-" => Some(Value::Int(a - b)),
                    "*" => Some(Value::Int(a * b)),
                    "/" if b != 0 => Some(Value::Int(a / b)),
                    _ => None
                };
                if let Some(val) = res {
                    return Some(Node {
                        identity: callee.identity.clone(),
                        ty: StaticType::Int,
                        kind: BoundKind::Constant(val),
                    });
                }
            }
        }
        None
    }

    fn get_const_int(&self, node: &TypedNode) -> Option<i64> {
        match &node.kind {
            BoundKind::Constant(Value::Int(i)) => Some(*i),
            BoundKind::Tuple { elements } if elements.len() == 1 => self.get_const_int(&elements[0]),
            _ => None
        }
    }

    fn contains_def_local(&self, node: &TypedNode) -> bool {
        match &node.kind {
            BoundKind::DefLocal { .. } => true,
            BoundKind::If { cond, then_br, else_br } => {
                self.contains_def_local(cond) || self.contains_def_local(then_br) || else_br.as_ref().is_some_and(|e| self.contains_def_local(e))
            }
            BoundKind::Block { exprs } => exprs.iter().any(|e| self.contains_def_local(e)),
            BoundKind::Call { callee, args } => self.contains_def_local(callee) || self.contains_def_local(args),
            BoundKind::Lambda { .. } => false, 
            BoundKind::Tuple { elements } => elements.iter().any(|e| self.contains_def_local(e)),
            BoundKind::Record { fields } => fields.iter().any(|(k, v)| self.contains_def_local(k) || self.contains_def_local(v)),
            BoundKind::Set { value, .. } => self.contains_def_local(value),
            _ => false,
        }
    }
}

/// Helper for transitively inlining values and cleaning up capture lists.
struct SubstitutionMap {
    /// Mapping Address::Local(slot) -> Value
    locals: HashMap<u32, Value>,
    /// Mapping Address::Upvalue(idx) -> Value
    upvalues: HashMap<u32, Value>,
}

impl SubstitutionMap {
    fn new() -> Self {
        Self { locals: HashMap::new(), upvalues: HashMap::new() }
    }

    fn add_local(&mut self, slot: u32, val: Value) {
        self.locals.insert(slot, val);
    }

    fn add_upvalue(&mut self, idx: u32, val: Value) {
        self.upvalues.insert(idx, val);
    }

    /// Re-maps Get(Upvalue(old_idx)) to Get(Upvalue(new_idx)) based on a mapping table.
    fn reindex_upvalues(&self, node: TypedNode, mapping: &[Option<u32>]) -> TypedNode {
        let (new_kind, new_ty) = match node.kind {
            BoundKind::Get { addr: Address::Upvalue(idx), name } => {
                if let Some(res) = mapping.get(idx as usize) {
                    match res {
                        Some(new_idx) => (BoundKind::Get { addr: Address::Upvalue(*new_idx), name }, node.ty),
                        None => (BoundKind::Get { addr: Address::Upvalue(idx), name }, node.ty), // Should have been inlined
                    }
                } else {
                    (BoundKind::Get { addr: Address::Upvalue(idx), name }, node.ty)
                }
            },
            
            BoundKind::Lambda { params, upvalues, body, positional_count } => {
                // IMPORTANT: If this nested lambda captures an upvalue from our current scope,
                // we MUST re-index it in its own capture list!
                let mut next_upvalues = Vec::new();
                for addr in upvalues {
                    if let Address::Upvalue(idx) = addr 
                        && let Some(res) = mapping.get(idx as usize) {
                            if let Some(new_idx) = res {
                                next_upvalues.push(Address::Upvalue(*new_idx));
                            }
                            continue; // Inlined or re-indexed
                    }
                    next_upvalues.push(addr);
                }

                // Note: We don't recurse into the body with the SAME mapping, 
                // because nested Get(Upvalue) nodes refer to THIS lambda's capture list.
                (BoundKind::Lambda { params, upvalues: next_upvalues, body, positional_count }, node.ty)
            },

            BoundKind::If { cond, then_br, else_br } => {

                let cond = Box::new(self.reindex_upvalues(*cond, mapping));
                let then_br = Box::new(self.reindex_upvalues(*then_br, mapping));
                let else_br = else_br.map(|e| Box::new(self.reindex_upvalues(*e, mapping)));
                (BoundKind::If { cond, then_br, else_br }, node.ty)
            },
            BoundKind::Block { exprs } => {
                let exprs = exprs.into_iter().map(|e| self.reindex_upvalues(e, mapping)).collect();
                (BoundKind::Block { exprs }, node.ty)
            },
            BoundKind::Call { callee, args } => {
                let callee = Box::new(self.reindex_upvalues(*callee, mapping));
                let args = Box::new(self.reindex_upvalues(*args, mapping));
                (BoundKind::Call { callee, args }, node.ty)
            },
            BoundKind::DefLocal { name, slot, value, captured_by } => {
                let value = Box::new(self.reindex_upvalues(*value, mapping));
                (BoundKind::DefLocal { name, slot, value, captured_by }, node.ty)
            },
            BoundKind::Set { addr, value } => {
                let value = Box::new(self.reindex_upvalues(*value, mapping));
                (BoundKind::Set { addr, value }, node.ty)
            },
            BoundKind::Tuple { elements } => {
                let elements = elements.into_iter().map(|e| self.reindex_upvalues(e, mapping)).collect();
                (BoundKind::Tuple { elements }, node.ty)
            },
            BoundKind::Record { fields } => {
                let fields = fields.into_iter().map(|(k, v)| (self.reindex_upvalues(k, mapping), self.reindex_upvalues(v, mapping))).collect();
                (BoundKind::Record { fields }, node.ty)
            },
            k => (k, node.ty),
        };

        Node { identity: node.identity, kind: new_kind, ty: new_ty }
    }
}