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

use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, VirtualId};
use crate::ast::types::{Purity, Value};
use crate::ast::vm::Closure;
use std::cell::RefCell;
use std::collections::HashSet;
use std::rc::Rc;

use super::substitution_map::SubstitutionMap;
use super::utils::UsageInfo;

pub struct Inliner<'a> {
    pub globals: &'a Option<Rc<RefCell<Vec<Value>>>>,
    pub root_purity: &'a Option<Rc<RefCell<Vec<Purity>>>>,
}

impl<'a> Inliner<'a> {
    pub fn new(
        globals: &'a Option<Rc<RefCell<Vec<Value>>>>,
        root_purity: &'a Option<Rc<RefCell<Vec<Purity>>>>,
    ) -> Self {
        Self {
            globals,
            root_purity,
        }
    }

    pub fn is_inlinable_value(&self, val: &Value, addr: Address<VirtualId>) -> bool {
        let type_ok = match val {
            Value::Int(_)
            | Value::Float(_)
            | Value::Bool(_)
            | Value::Text(_)
            | Value::Keyword(_)
            | Value::Record(_, _)
            | Value::DateTime(_) => true,
            Value::Object(obj) => {
                if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
                    closure.upvalues.is_empty() && !closure.function_node.ty.is_recursive
                } else {
                    false
                }
            }
            _ => false,
        };

        if !type_ok {
            return false;
        }

        if let Address::Global(idx) = addr {
            if let Some(purity_rc) = &self.root_purity {
                let purity = purity_rc
                    .borrow()
                    .get(idx.0 as usize)
                    .cloned()
                    .unwrap_or(Purity::Impure);
                return purity >= Purity::Pure;
            }
            return false;
        }

        true
    }

    pub fn prepare_beta_reduction(
        &self,
        params: &AnalyzedNode,
        arg_vals: &[Rc<AnalyzedNode>],
        body: &AnalyzedNode,
        sub: &mut SubstitutionMap,
    ) -> Option<()> {
        let mut body_usage = UsageInfo::default();
        body_usage.collect(body);

        let mut slot_index = 0;
        self.map_params_to_args(params, arg_vals, &mut slot_index, sub, &body_usage);

        if slot_index != arg_vals.len() {
            return None;
        }

        let mut param_slots = HashSet::new();
        self.collect_parameter_slots_set(params, &mut param_slots);

        for slot in param_slots {
            let addr = Address::Local(slot);
            if (body_usage.is_used(&addr) || body_usage.is_assigned(&addr))
                && sub.get_value(&addr).is_none()
                && !sub.ast_substitutions.contains_key(&addr)
            {
                return None;
            }
        }

        if sub.values.is_empty() && sub.ast_substitutions.is_empty() && !arg_vals.is_empty() {
            return None;
        }

        Some(())
    }

    pub fn map_params_to_args(
        &self,
        pattern: &AnalyzedNode,
        args: &[Rc<AnalyzedNode>],
        offset: &mut usize,
        sub: &mut SubstitutionMap,
        body_usage: &UsageInfo,
    ) {
        match &pattern.kind {
            BoundKind::Define { addr, .. } => {
                if let Some(arg) = args.get(*offset)
                    && !body_usage.is_assigned(addr)
                {
                    let mut core_arg = arg.as_ref();
                    while let BoundKind::Expansion { bound_expanded, .. } = &core_arg.kind {
                        core_arg = bound_expanded.as_ref();
                    }

                    if let BoundKind::Constant(val) = &core_arg.kind {
                        sub.add_value(*addr, val.clone());
                    } else if let BoundKind::Lambda { upvalues, .. } = &core_arg.kind
                        && upvalues.is_empty()
                    {
                        sub.add_ast_substitution(*addr, core_arg.clone());
                    } else if let BoundKind::Get {
                        addr: Address::Global(_),
                        ..
                    } = &core_arg.kind
                    {
                        sub.add_ast_substitution(*addr, core_arg.clone());
                    }
                }

                if let Address::Local(slot) = addr {
                    sub.map_slot(*slot);
                }
                *offset += 1;
            }
            BoundKind::Tuple { elements } => {
                for el in elements {
                    self.map_params_to_args(el, args, offset, sub, body_usage);
                }
            }
            _ => {}
        }
    }

    pub fn collect_parameter_slots_set(&self, node: &AnalyzedNode, slots: &mut HashSet<VirtualId>) {
        match &node.kind {
            BoundKind::Define {
                addr: Address::Local(slot),
                ..
            } => {
                slots.insert(*slot);
            }
            BoundKind::Tuple { elements } => {
                for el in elements {
                    self.collect_parameter_slots_set(el, slots);
                }
            }
            _ => {}
        }
    }
}