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

use crate::ast::compiler::bound_nodes::{
    Address, AnalyzedNode, BoundKind, GlobalAnalyzedRegistry, GlobalIdx, UpvalueIdx, UpvalueSource,
};
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<GlobalAnalyzedRegistry>>>,
}

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<GlobalAnalyzedRegistry>>,
    ) -> Self {
        self.lambda_registry = Some(registry);
        self
    }

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

        let mut current = Rc::new(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 Rc::ptr_eq(&next, &current) {
                break;
            }
            current = next;
        }
        (*current).clone()
    }

    fn try_inline(
        &self,
        params: &AnalyzedNode,
        arg_nodes: &[Rc<AnalyzedNode>],
        body: &AnalyzedNode,
        sub: &mut SubstitutionMap,
        path: &mut PathTracker,
        base_sub: Option<SubstitutionMap>,
    ) -> Option<Rc<AnalyzedNode>> {
        let inliner = Inliner::new(&self.globals, &self.global_purity);
        let mut inner_sub = base_sub.unwrap_or_else(|| sub.new_for_inlining());
        
        if inliner.prepare_beta_reduction(params, arg_nodes, body, &mut inner_sub).is_some() {
            let res = self.visit_node(Rc::new(body.clone()), &mut inner_sub, path);
            
            sub.next_slot = inner_sub.next_slot;
            sub.used.extend(inner_sub.used.iter().cloned());
            sub.assigned.extend(inner_sub.assigned.iter().cloned());
            sub.captured_slots.extend(inner_sub.captured_slots.iter().cloned());
            
            Some(res)
        } else {
            None
        }
    }

    fn visit_node(
        &self,
        node_rc: Rc<AnalyzedNode>,
        sub: &mut SubstitutionMap,
        path: &mut PathTracker,
    ) -> Rc<AnalyzedNode> {
        let node = &*node_rc;
        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, name } => {
                let addr = *addr;
                if !sub.assigned.contains(&addr) {
                    if let Some(val) = sub.get_value(&addr)
                        && inliner.is_inlinable_value(val, addr)
                    {
                        return Rc::new(folder.make_constant_node(val.clone(), node));
                    }

                    if let Some(inlined_node) = sub.ast_substitutions.get(&addr) {
                        return inlined_node.clone();
                    }

                    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 Rc::new(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 { rec, field } => {
                let rec_opt = self.visit_node(rec.clone(), sub, path);

                if let BoundKind::Constant(Value::Record(layout, values)) = &rec_opt.kind
                    && let Some(idx) = layout.index_of(*field)
                {
                    return Rc::new(folder.make_constant_node(values[idx].clone(), node));
                }

                if Rc::ptr_eq(&rec_opt, rec) {
                    return node_rc;
                }

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

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

                if Rc::ptr_eq(&value_opt, value) {
                    return node_rc;
                }

                (
                    BoundKind::Set {
                        addr: sub.map_address(*addr),
                        value: value_opt,
                    },
                    node.ty.clone(),
                )
            }

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

                if let Address::Local(slot) = addr
                    && !captured_by.borrow().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);
                }

                if Rc::ptr_eq(&value_opt, value) {
                    return node_rc;
                }

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

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

                if self.enabled {
                    if let BoundKind::FieldAccessor(k) = &callee_opt.kind
                        && let BoundKind::Tuple { elements } = &args_opt.kind
                        && elements.len() == 1
                    {
                        let rec = elements[0].clone();
                        let metrics = node.ty.clone();
                        return Rc::new(Node {
                            identity: node.identity.clone(),
                            kind: BoundKind::GetField {
                                rec,
                                field: *k,
                            },
                            ty: metrics,
                        });
                    }

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

                    if let BoundKind::Lambda {
                        params,
                        body,
                        upvalues,
                        ..
                    } = &callee_opt.kind
                        && upvalues.is_empty()
                        && path.inlining_depth < 5
                        && !callee_opt.ty.is_recursive
                        && path.enter_lambda(&callee_opt.identity)
                    {
                        path.inlining_depth += 1;
                        let collapsed = self.try_inline(params, &arg_nodes, body, sub, path, None);
                        path.inlining_depth -= 1;
                        path.exit_lambda(&callee_opt.identity);

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

                    if let BoundKind::Get {
                        addr: Address::Global(idx),
                        ..
                    } = &callee_opt.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,
                                ..
                            } = &lambda_node.kind
                            && upvalues.is_empty()
                            && !lambda_node.ty.is_recursive
                        {
                            path.inlining_stack.insert(*idx);
                            path.inlining_depth += 1;
                            let collapsed = self.try_inline(params, &arg_nodes, body, sub, path, 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_opt.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 = sub.new_for_inlining();
                        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 collapsed = self.try_inline(
                            &closure.parameter_node,
                            &arg_nodes,
                            &closure.function_node,
                            sub,
                            path,
                            Some(closure_sub),
                        );
                        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_opt, &arg_nodes) {
                        return Rc::new(folded);
                    }
                }

                if Rc::ptr_eq(&callee_opt, callee) && Rc::ptr_eq(&args_opt, args) {
                    return node_rc;
                }

                (
                    BoundKind::Call {
                        callee: callee_opt,
                        args: args_opt,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Again { args } => {
                let args_opt = self.visit_node(args.clone(), sub, path);
                if Rc::ptr_eq(&args_opt, args) {
                    return node_rc;
                }
                (
                    BoundKind::Again {
                        args: args_opt,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::If {
                cond,
                then_br,
                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 Rc::new(folder.make_nop_node(node));
                    }
                }

                let then_br_opt = self.visit_node(then_br.clone(), sub, path);
                let else_br_opt = else_br
                    .as_ref()
                    .map(|e| self.visit_node(e.clone(), sub, path));

                let unchanged = Rc::ptr_eq(&cond_opt, cond) && Rc::ptr_eq(&then_br_opt, then_br) && match (&else_br_opt, else_br) {
                    (Some(a), Some(b)) => Rc::ptr_eq(a, b),
                    (None, None) => true,
                    _ => false,
                };

                if unchanged {
                    return node_rc;
                }

                (
                    BoundKind::If {
                        cond: cond_opt,
                        then_br: then_br_opt,
                        else_br: else_br_opt,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Pipe {
                inputs,
                lambda,
                out_type,
            } => {
                let mut o_inputs = Vec::with_capacity(inputs.len());
                let mut inputs_changed = false;
                for input in inputs {
                    let opt = self.visit_node(input.clone(), sub, path);
                    if !Rc::ptr_eq(&opt, input) {
                        inputs_changed = true;
                    }
                    o_inputs.push(opt);
                }
                let o_lambda = self.visit_node(lambda.clone(), sub, path);

                if !inputs_changed && Rc::ptr_eq(&o_lambda, lambda) {
                    return node_rc;
                }

                (
                    BoundKind::Pipe {
                        inputs: o_inputs,
                        lambda: o_lambda,
                        out_type: out_type.clone(),
                    },
                    node.ty.clone(),
                )
            }
            BoundKind::Block { statements, result } => {
                let mut info = UsageInfo::default();
                for s in statements {
                    info.collect(s);
                }
                info.collect(result);

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

                let mut new_statements = Vec::with_capacity(statements.len());
                let mut block_changed = false;

                for s in statements {
                    if self.enabled {
                        let removable = match &s.kind {
                            BoundKind::Define { addr, value, captured_by, .. } => {
                                !info.is_used(addr)
                                    && (if let Address::Local(slot) = addr {
                                        !sub.captured_slots.contains(slot)
                                    } else {
                                        true
                                    })
                                    && captured_by.borrow().is_empty()
                                    && (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
                            }
                            _ => s.ty.purity >= Purity::SideEffectFree,
                        };

                        if removable {
                            block_changed = true;
                            continue;
                        }
                    }

                    let opt = self.visit_node(s.clone(), sub, path);
                    if self.enabled && matches!(opt.kind, BoundKind::Nop) {
                        block_changed = true;
                        continue;
                    }
                    if !Rc::ptr_eq(&opt, s) {
                        block_changed = true;
                    }
                    new_statements.push(opt);
                }

                let new_result = self.visit_node(result.clone(), sub, path);
                if !Rc::ptr_eq(&new_result, result) {
                    block_changed = true;
                }

                if !block_changed && Rc::ptr_eq(&new_result, result) {
                    return node_rc;
                }

                if self.enabled && new_statements.is_empty() {
                    return new_result;
                }

                (BoundKind::Block { statements: new_statements, result: new_result }, node.ty.clone())
                }

            BoundKind::Lambda {
                params,
                upvalues,
                body,
                positional_count,
                max_slots,
            } => {
                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;
                let mut upvalues_changed = false;

                for (old_idx, source) in upvalues.iter().enumerate() {
                    let capture_addr = match source {
                        UpvalueSource::Local(s) => Address::Local(*s),
                        UpvalueSource::Upvalue(i) => Address::Upvalue(*i),
                    };

                    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);
                        upvalues_changed = true;
                    } else {
                        mapping.push(Some(new_upvalues.len() as u32));
                        let mapped_addr = sub.map_address(capture_addr);
                        let mapped_source = match mapped_addr {
                            Address::Local(s) => UpvalueSource::Local(s),
                            Address::Upvalue(i) => UpvalueSource::Upvalue(i),
                            Address::Global(_) => panic!("Cannot map upvalue to global")
                        };
                        if mapped_source != *source {
                            upvalues_changed = true;
                        }
                        new_upvalues.push(mapped_source);
                    }
                }

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

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

                if !upvalues_changed && Rc::ptr_eq(&params_opt, params) && Rc::ptr_eq(&reindexed_body, body) {
                    return node_rc;
                }

                (
                    BoundKind::Lambda {
                        params: params_opt,
                        upvalues: new_upvalues,
                        body: reindexed_body,
                        positional_count: *positional_count,
                        max_slots: *max_slots,
                    },
                    node.ty.clone(),
                )
            }

            BoundKind::Destructure { pattern, value } => {
                let val_opt = self.visit_node(value.clone(), sub, path);
                let pat_opt = 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);
                }

                if Rc::ptr_eq(&val_opt, value) && Rc::ptr_eq(&pat_opt, pattern) {
                    return node_rc;
                }

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

            BoundKind::Tuple { elements } => {
                let mut new_elements = Vec::with_capacity(elements.len());
                let mut changed = false;
                for e in elements {
                    let opt = self.visit_node(e.clone(), sub, path);
                    if !Rc::ptr_eq(&opt, e) {
                        changed = true;
                    }
                    new_elements.push(opt);
                }
                if !changed {
                    return node_rc;
                }
                (BoundKind::Tuple { elements: new_elements }, node.ty.clone())
            }
            BoundKind::Record { layout, values } => {
                let mut mapped_values = Vec::with_capacity(values.len());
                let mut changed = false;
                for v in values {
                    let opt = self.visit_node(v.clone(), sub, path);
                    if !Rc::ptr_eq(&opt, v) {
                        changed = true;
                    }
                    mapped_values.push(opt);
                }

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

                if !changed {
                    return node_rc;
                }

                (
                    BoundKind::Record {
                        layout: layout.clone(),
                        values: mapped_values,
                    },
                    node.ty.clone(),
                )
            }
            BoundKind::Expansion {
                original_call,
                bound_expanded,
            } => {
                path.inlining_depth += 1;
                let expanded_opt = self.visit_node(bound_expanded.clone(), sub, path);
                path.inlining_depth -= 1;
                
                if Rc::ptr_eq(&expanded_opt, bound_expanded) {
                    return node_rc;
                }

                (
                    BoundKind::Expansion {
                        original_call: original_call.clone(),
                        bound_expanded: expanded_opt,
                    },
                    node.ty.clone(),
                )
            }
            k => (k.clone(), node.ty.clone()),
        };

        Rc::new(Node {
            identity: node.identity.clone(),
            kind: new_kind,
            ty: metrics,
        })
    }

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