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

use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, GlobalIdx, UpvalueIdx};
use crate::ast::nodes::Node;
use crate::ast::types::{Purity, Value};
use crate::ast::vm::Closure;
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

use super::folder::Folder;
use super::inliner::Inliner;
use super::substitution_map::SubstitutionMap;
use super::utils::{PathTracker, UsageInfo};

pub struct Optimizer {
    pub enabled: bool,
    max_passes: usize,
    pub globals: Option<Rc<RefCell<Vec<Value>>>>,
    pub global_purity: Option<Rc<RefCell<HashMap<GlobalIdx, Purity>>>>,
    pub lambda_registry: Option<Rc<RefCell<HashMap<GlobalIdx, AnalyzedNode>>>>,
}

impl Optimizer {
    pub fn new(enabled: bool) -> Self {
        Self {
            enabled,
            max_passes: 5,
            globals: None,
            global_purity: None,
            lambda_registry: None,
        }
    }

    pub fn with_globals(mut self, globals: Rc<RefCell<Vec<Value>>>) -> Self {
        self.globals = Some(globals);
        self
    }

    pub fn with_purity(mut self, purity: Rc<RefCell<HashMap<GlobalIdx, Purity>>>) -> Self {
        self.global_purity = Some(purity);
        self
    }

    pub fn with_registry(
        mut self,
        registry: Rc<RefCell<HashMap<GlobalIdx, AnalyzedNode>>>,
    ) -> Self {
        self.lambda_registry = Some(registry);
        self
    }

    pub fn optimize(&self, node: AnalyzedNode) -> AnalyzedNode {
        if !self.enabled {
            return node;
        }

        let mut current = node;
        for _ in 0..self.max_passes {
            let mut sub = SubstitutionMap::new();
            let mut path = PathTracker::new();
            let next = self.visit_node(current.clone(), &mut sub, &mut path);
            if next == current {
                break;
            }
            current = next;
        }
        current
    }

    fn visit_node(
        &self,
        node: AnalyzedNode,
        sub: &mut SubstitutionMap,
        path: &mut PathTracker,
    ) -> AnalyzedNode {
        let folder = Folder::new(&self.globals);
        let inliner = Inliner::new(&self.globals, &self.global_purity);
        let (new_kind, metrics) = match node.kind {
            BoundKind::Get { addr, ref name } => {
                if !sub.assigned.contains(&addr) {
                    // 1. Try inlining from current value substitution map (locals/globals/upvalues)
                    if let Some(val) = sub.get_value(&addr)
                        && inliner.is_inlinable_value(val, addr)
                    {
                        return folder.make_constant_node(val.clone(), &node);
                    }

                    // 2. Try inlining from AST substitution map (pure expressions)
                    if let Some(inlined_node) = sub.ast_substitutions.get(&addr) {
                        return inlined_node.clone();
                    }

                    // 3. Fallback for Globals: check the actual VM environment
                    if let Address::Global(idx) = addr
                        && let Some(globals_rc) = &self.globals
                    {
                        let globals = globals_rc.borrow();
                        if let Some(val) = globals.get(idx.0 as usize)
                            && inliner.is_inlinable_value(val, addr)
                        {
                            return folder.make_constant_node(val.clone(), &node);
                        }
                    }
                }

                sub.used.insert(addr);
                (
                    BoundKind::Get {
                        addr: sub.map_address(addr),
                        name: name.clone(),
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::FieldAccessor(k) => (BoundKind::FieldAccessor(k), node.ty.clone()),

            BoundKind::GetField { ref rec, field } => {
                let rec_opt = self.visit_node((**rec).clone(), sub, path);

                // Constant folding for Field Access
                if let BoundKind::Constant(Value::Record(layout, values)) = &rec_opt.kind
                    && let Some(idx) = layout.index_of(field)
                {
                    return folder.make_constant_node(values[idx].clone(), &node);
                }

                (
                    BoundKind::GetField {
                        rec: Box::new(rec_opt),
                        field,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Set { addr, value } => {
                let value = Box::new(self.visit_node(*value, sub, path));
                if let BoundKind::Constant(val) = &value.kind {
                    sub.add_value(addr, val.clone());
                } else {
                    sub.remove_value(&addr);
                }
                (
                    BoundKind::Set {
                        addr: sub.map_address(addr),
                        value,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Define {
                ref name,
                addr,
                kind,
                ref value,
                ref captured_by,
            } => {
                let value_opt = Box::new(self.visit_node((**value).clone(), sub, path));

                if let Address::Local(slot) = addr
                    && !captured_by.is_empty()
                {
                    sub.captured_slots.insert(slot);
                }

                if let BoundKind::Constant(val) = &value_opt.kind {
                    sub.add_value(addr, val.clone());
                } else {
                    sub.remove_value(&addr);
                }

                if let Address::Global(global_index) = addr
                    && value_opt.ty.purity > Purity::Impure
                    && let Some(purity_rc) = &self.global_purity
                {
                    purity_rc
                        .borrow_mut()
                        .insert(global_index, value_opt.ty.purity);
                }

                (
                    BoundKind::Define {
                        name: name.clone(),
                        addr: sub.map_address(addr),
                        kind,
                        value: value_opt,
                        captured_by: captured_by.clone(),
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Call { callee, args } => {
                let callee = self.visit_node(*callee, sub, path);
                let args = self.visit_node(*args, sub, path);

                if self.enabled {
                    // Optimized Field Access Transformation
                    if let BoundKind::FieldAccessor(k) = &callee.kind
                        && let BoundKind::Tuple { elements } = &args.kind
                        && elements.len() == 1
                    {
                        let rec = elements[0].clone();
                        let metrics = node.ty.clone();
                        return Node {
                            identity: node.identity,
                            kind: BoundKind::GetField {
                                rec: Box::new(rec),
                                field: *k,
                            },
                            ty: metrics,
                        };
                    }

                    let mut arg_nodes = Vec::new();
                    self.flatten_tuple(args.clone(), &mut arg_nodes);

                    if let BoundKind::Lambda {
                        params,
                        body,
                        upvalues,
                        positional_count,
                    } = &callee.kind
                        && upvalues.is_empty()
                        && positional_count.is_some()
                        && path.inlining_depth < 5
                        && !callee.ty.is_recursive
                        && path.enter_lambda(&callee.identity)
                    {
                        path.inlining_depth += 1;
                        let mut inner_sub = sub.new_inner();
                        let collapsed = if inliner
                            .prepare_beta_reduction(params, &arg_nodes, body, &mut inner_sub)
                            .is_some()
                        {
                            Some(self.visit_node((**body).clone(), &mut inner_sub, path))
                        } else {
                            None
                        };
                        path.inlining_depth -= 1;
                        path.exit_lambda(&callee.identity);

                        if let Some(res) = collapsed {
                            return res;
                        }
                    }

                    if let BoundKind::Get {
                        addr: Address::Global(idx),
                        ..
                    } = &callee.kind
                        && let Some(registry_rc) = &self.lambda_registry
                        && path.inlining_depth < 5
                        && !path.inlining_stack.contains(idx)
                    {
                        let registry = registry_rc.borrow();
                        if let Some(lambda_node) = registry.get(idx)
                            && let BoundKind::Lambda {
                                params,
                                body,
                                upvalues,
                                positional_count,
                            } = &lambda_node.kind
                            && upvalues.is_empty()
                            && positional_count.is_some()
                            && !lambda_node.ty.is_recursive
                        {
                            let mut inner_sub = sub.new_inner();
                            path.inlining_stack.insert(*idx);
                            path.inlining_depth += 1;
                            let collapsed = if inliner
                                .prepare_beta_reduction(params, &arg_nodes, body, &mut inner_sub)
                                .is_some()
                            {
                                Some(self.visit_node((**body).clone(), &mut inner_sub, path))
                            } else {
                                None
                            };
                            path.inlining_depth -= 1;
                            path.inlining_stack.remove(idx);

                            if let Some(res) = collapsed {
                                return res;
                            }
                        }
                    }

                    if let BoundKind::Constant(Value::Object(ref obj)) = callee.kind
                        && path.inlining_depth < 5
                        && let Some(closure) = obj.as_any().downcast_ref::<Closure>()
                        && (closure.upvalues.is_empty()
                            || closure.function_node.ty.purity >= Purity::SideEffectFree)
                        && !closure.function_node.ty.is_recursive
                        && path.enter_lambda(&closure.function_node.identity)
                    {
                        let mut closure_sub = SubstitutionMap::new();
                        for (i, cell) in closure.upvalues.iter().enumerate() {
                            closure_sub.add_value(
                                Address::Upvalue(UpvalueIdx(i as u32)),
                                cell.borrow().clone(),
                            );
                        }

                        path.inlining_depth += 1;
                        let inlined_body = self.visit_node(
                            (*closure.function_node).clone(),
                            &mut closure_sub,
                            path,
                        );

                        let collapsed = if inliner
                            .prepare_beta_reduction(
                                &closure.parameter_node,
                                &arg_nodes,
                                &inlined_body,
                                &mut closure_sub,
                            )
                            .is_some()
                        {
                            Some(self.visit_node(inlined_body, &mut closure_sub, path))
                        } else {
                            None
                        };
                        path.inlining_depth -= 1;
                        path.exit_lambda(&closure.function_node.identity);

                        if let Some(res) = collapsed {
                            return res;
                        }
                    }

                    if let Some(folded) = folder.try_fold_pure(&callee, &arg_nodes) {
                        return folded;
                    }
                }

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

            BoundKind::Again { args } => {
                let args = self.visit_node(*args, sub, path);
                (
                    BoundKind::Again {
                        args: Box::new(args),
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::If {
                ref cond,
                ref then_br,
                ref else_br,
            } => {
                let cond_opt = self.visit_node((**cond).clone(), sub, path);
                if self.enabled
                    && let BoundKind::Constant(ref val) = cond_opt.kind
                {
                    if val.is_truthy() {
                        return self.visit_node((**then_br).clone(), sub, path);
                    } else if let Some(else_node) = else_br {
                        return self.visit_node((**else_node).clone(), sub, path);
                    } else {
                        return folder.make_nop_node(&node);
                    }
                }

                let then_br = Box::new(self.visit_node((**then_br).clone(), sub, path));
                let else_br = else_br
                    .as_ref()
                    .map(|e| Box::new(self.visit_node((**e).clone(), sub, path)));
                (
                    BoundKind::If {
                        cond: Box::new(cond_opt),
                        then_br,
                        else_br,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Pipe {
                inputs,
                lambda,
                out_type,
            } => {
                let mut o_inputs = Vec::with_capacity(inputs.len());
                for input in inputs {
                    o_inputs.push(self.visit_node(input, sub, path));
                }
                let o_lambda = Box::new(self.visit_node(*lambda, sub, path));
                (
                    BoundKind::Pipe {
                        inputs: o_inputs,
                        lambda: o_lambda,
                        out_type: out_type.clone(),
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Block { ref exprs } => {
                let mut info = UsageInfo::default();
                if !exprs.is_empty() {
                    for e in exprs {
                        info.collect(e);
                    }
                }

                sub.assigned.extend(info.assigned.iter().cloned());

                let mut new_exprs = Vec::with_capacity(exprs.len());
                let last_idx = exprs.len().saturating_sub(1);

                for (i, e) in exprs.iter().enumerate() {
                    let is_last = i == last_idx;
                    if self.enabled && !is_last {
                        let removable = match &e.kind {
                            BoundKind::Define { addr, value, .. } => {
                                !info.is_used(addr)
                                    && (if let Address::Local(slot) = addr {
                                        !sub.captured_slots.contains(slot)
                                    } else {
                                        true
                                    })
                                    && (value.ty.purity >= Purity::SideEffectFree
                                        || matches!(value.kind, BoundKind::Lambda { .. }))
                            }
                            BoundKind::Set { addr, value, .. } => {
                                !info.is_used(addr)
                                    && (if let Address::Local(slot) = addr {
                                        !sub.captured_slots.contains(slot)
                                    } else {
                                        true
                                    })
                                    && value.ty.purity >= Purity::SideEffectFree
                            }
                            _ => e.ty.purity >= Purity::SideEffectFree,
                        };

                        if removable {
                            continue;
                        }
                    }

                    let opt = self.visit_node(e.clone(), sub, path);
                    if self.enabled && matches!(opt.kind, BoundKind::Nop) && !is_last {
                        continue;
                    }
                    new_exprs.push(opt);
                }

                if self.enabled {
                    if new_exprs.is_empty() {
                        return folder.make_nop_node(&node);
                    } else if new_exprs.len() == 1 {
                        return new_exprs.pop().unwrap();
                    }
                }

                (BoundKind::Block { exprs: new_exprs }, node.ty.clone())
            }

            BoundKind::Lambda { .. } => {
                let (params, original_upvalues, body, positional_count) =
                    if let BoundKind::Lambda {
                        params,
                        upvalues,
                        body,
                        positional_count,
                    } = &node.kind
                    {
                        (params, upvalues, body, positional_count)
                    } else {
                        unreachable!()
                    };

                let mut info = UsageInfo::default();
                info.collect(&node);

                let mut new_upvalues = Vec::new();
                let mut mapping = Vec::new();
                let mut next_inner_subs = sub.new_inner();
                next_inner_subs.assigned = info.assigned;

                for (old_idx, capture_addr) in original_upvalues.iter().enumerate() {
                    let mut inlined_val = None;
                    if !sub.assigned.contains(capture_addr)
                        && let Some(val) = sub.get_value(capture_addr)
                    {
                        inlined_val = Some(val.clone());
                    }

                    if let Address::Local(slot) = capture_addr {
                        sub.captured_slots.insert(*slot);
                    }

                    if let Some(val) = inlined_val {
                        next_inner_subs
                            .add_value(Address::Upvalue(UpvalueIdx(old_idx as u32)), val);
                        mapping.push(None);
                    } else {
                        mapping.push(Some(new_upvalues.len() as u32));
                        new_upvalues.push(sub.map_address(*capture_addr));
                    }
                }

                let params_node =
                    self.visit_node(params.as_ref().clone(), &mut next_inner_subs, path);
                inliner.collect_parameter_slots(&params_node, &mut next_inner_subs);

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

                (
                    BoundKind::Lambda {
                        params: Rc::new(params_node),
                        upvalues: new_upvalues,
                        body: Rc::new(reindexed_body),
                        positional_count: *positional_count,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Destructure {
                ref pattern,
                ref value,
            } => {
                let val_opt = Box::new(self.visit_node((**value).clone(), sub, path));
                let pat_opt = Box::new(self.visit_node((**pattern).clone(), sub, path));

                let mut info = UsageInfo::default();
                info.collect_pattern(&pat_opt);
                for addr in info.assigned {
                    sub.remove_value(&addr);
                }

                (
                    BoundKind::Destructure {
                        pattern: pat_opt,
                        value: val_opt,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Tuple { elements } => {
                let elements = elements
                    .into_iter()
                    .map(|e| self.visit_node(e, sub, path))
                    .collect();
                (BoundKind::Tuple { elements }, node.ty.clone())
            }
            BoundKind::Record {
                ref layout,
                ref values,
            } => {
                let mapped_values: Vec<_> = values
                    .iter()
                    .map(|v| self.visit_node(v.clone(), sub, path))
                    .collect();

                if self.enabled
                    && let Some(folded) = folder.try_fold_record(layout, &mapped_values, &node)
                {
                    return folded;
                }

                (
                    BoundKind::Record {
                        layout: layout.clone(),
                        values: mapped_values,
                    },
                    node.ty.clone(),
                )
            }
            BoundKind::Expansion {
                ref original_call,
                ref bound_expanded,
            } => {
                path.inlining_depth += 1;
                let bound_expanded =
                    Box::new(self.visit_node((**bound_expanded).clone(), sub, path));
                path.inlining_depth -= 1;
                (
                    BoundKind::Expansion {
                        original_call: original_call.clone(),
                        bound_expanded,
                    },
                    node.ty.clone(),
                )
            }
            k => (k, node.ty.clone()),
        };

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

    fn flatten_tuple(&self, node: AnalyzedNode, into: &mut Vec<AnalyzedNode>) {
        match node.kind {
            BoundKind::Tuple { elements } => {
                for el in elements {
                    self.flatten_tuple(el, into);
                }
            }
            BoundKind::Nop => {}
            _ => into.push(node),
        }
    }
}