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

use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, LocalSlot, UpvalueIdx, UpvalueSource};
use crate::ast::nodes::Node;
use crate::ast::types::Value;
use std::collections::{HashMap, HashSet};
use std::rc::Rc;

#[derive(Default)]
pub struct SubstitutionMap {
    pub values: HashMap<Address, Value>,
    pub ast_substitutions: HashMap<Address, Rc<AnalyzedNode>>,
    pub slot_mapping: HashMap<LocalSlot, LocalSlot>,
    pub assigned: HashSet<Address>,
    pub next_slot: u32,
    pub used: HashSet<Address>,
    pub captured_slots: HashSet<LocalSlot>,
}

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

    /// Creates a new SubstitutionMap for an inner scope (like an inlined Lambda).
    /// Safely inherits only Global substitutions, because Local and Upvalue
    /// addresses are relative to the specific function frame and would overlap.
    pub fn new_inner(&self) -> Self {
        let mut inner = Self::new();
        for (k, v) in &self.values {
            if matches!(k, Address::Global(_)) {
                inner.values.insert(*k, v.clone());
            }
        }
        for (k, v) in &self.ast_substitutions {
            if matches!(k, Address::Global(_)) {
                inner.ast_substitutions.insert(*k, v.clone());
            }
        }
        inner
    }

    pub fn new_for_inlining(&self) -> Self {
        let mut inner = self.new_inner();
        inner.next_slot = self.next_slot;
        inner
    }

    pub fn add_ast_substitution(&mut self, addr: Address, node: AnalyzedNode) {
        self.ast_substitutions.insert(addr, Rc::new(node));
    }

    pub fn map_slot(&mut self, old_slot: LocalSlot) -> LocalSlot {
        if let Some(&new_slot) = self.slot_mapping.get(&old_slot) {
            return new_slot;
        }
        let new_slot = LocalSlot(self.next_slot);
        self.slot_mapping.insert(old_slot, new_slot);
        self.next_slot += 1;
        new_slot
    }

    pub fn map_address(&mut self, addr: Address) -> Address {
        match addr {
            Address::Local(slot) => Address::Local(self.map_slot(slot)),
            other => other,
        }
    }

    pub fn add_value(&mut self, addr: Address, val: Value) {
        self.values.insert(addr, val);
    }

    pub fn get_value(&self, addr: &Address) -> Option<&Value> {
        self.values.get(addr)
    }

    pub fn remove_value(&mut self, addr: &Address) {
        self.values.remove(addr);
    }

    fn reindex_addr(&self, addr: Address, mapping: &[Option<u32>]) -> Address {
        if let Address::Upvalue(idx) = addr
            && let Some(res) = mapping.get(idx.0 as usize)
            && let Some(new_idx) = res
        {
            Address::Upvalue(UpvalueIdx(*new_idx))
        } else {
            addr
        }
    }

    fn reindex_source(&self, source: UpvalueSource, mapping: &[Option<u32>]) -> UpvalueSource {
        match source {
            UpvalueSource::Upvalue(idx) => {
                if let Some(res) = mapping.get(idx.0 as usize)
                    && let Some(new_idx) = res
                {
                    UpvalueSource::Upvalue(UpvalueIdx(*new_idx))
                } else {
                    source
                }
            }
            _ => source
        }
    }

    pub fn reindex_upvalues(&self, node_rc: Rc<AnalyzedNode>, mapping: &[Option<u32>]) -> Rc<AnalyzedNode> {
        let node = &*node_rc;
        let (new_kind, metrics) = match &node.kind {
            BoundKind::Get { addr, name } => (
                BoundKind::Get {
                    addr: self.reindex_addr(*addr, mapping),
                    name: name.clone(),
                },
                node.ty.clone(),
            ),
            BoundKind::Lambda {
                params,
                upvalues,
                body,
                positional_count,
                max_slots,
            } => {
                let mut next_upvalues = Vec::new();
                for source in upvalues {
                    next_upvalues.push(self.reindex_source(*source, mapping));
                }
                (
                    BoundKind::Lambda {
                        params: params.clone(),
                        upvalues: next_upvalues,
                        body: body.clone(),
                        positional_count: *positional_count,
                        max_slots: *max_slots,
                    },
                    node.ty.clone(),
                )
            }
            BoundKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let cond = self.reindex_upvalues(cond.clone(), mapping);
                let then_br = self.reindex_upvalues(then_br.clone(), mapping);
                let else_br = else_br.as_ref().map(|e| self.reindex_upvalues(e.clone(), mapping));
                (
                    BoundKind::If {
                        cond,
                        then_br,
                        else_br,
                    },
                    node.ty.clone(),
                )
            }
            BoundKind::Block { statements, result } => {
                let t_statements = statements
                    .iter()
                    .map(|s| self.reindex_upvalues(s.clone(), mapping))
                    .collect();
                let t_result = self.reindex_upvalues(result.clone(), mapping);
                (BoundKind::Block { statements: t_statements, result: t_result }, node.ty.clone())
            }
            BoundKind::Call { callee, args } => {
                let callee = self.reindex_upvalues(callee.clone(), mapping);
                let args = self.reindex_upvalues(args.clone(), mapping);
                (BoundKind::Call { callee, args }, node.ty.clone())
            }
            BoundKind::Define {
                name,
                addr,
                kind,
                value,
                identity,
                captured_by,
            } => {
                let value = self.reindex_upvalues(value.clone(), mapping);
                (
                    BoundKind::Define {
                        name: name.clone(),
                        addr: *addr,
                        kind: *kind,
                        value,
                        identity: identity.clone(),
                        captured_by: captured_by.clone(),
                    },
                    node.ty.clone(),
                )
            }
            BoundKind::Set { addr, value } => {
                let value = self.reindex_upvalues(value.clone(), mapping);
                (
                    BoundKind::Set {
                        addr: self.reindex_addr(*addr, mapping),
                        value,
                    },
                    node.ty.clone(),
                )
            }
            BoundKind::Tuple { elements } => {
                let elements = elements
                    .iter()
                    .map(|e| self.reindex_upvalues(e.clone(), mapping))
                    .collect();
                (BoundKind::Tuple { elements }, node.ty.clone())
            }
            BoundKind::Record { layout, values } => {
                let values = values
                    .iter()
                    .map(|v| self.reindex_upvalues(v.clone(), mapping))
                    .collect();
                (BoundKind::Record { layout: layout.clone(), values }, node.ty.clone())
            }
            BoundKind::Expansion { original_call, bound_expanded } => {
                let bound_expanded = self.reindex_upvalues(bound_expanded.clone(), mapping);
                (
                    BoundKind::Expansion {
                        original_call: original_call.clone(),
                        bound_expanded,
                    },
                    node.ty.clone(),
                )
            }
            k => (k.clone(), node.ty.clone()),
        };
        Rc::new(Node {
            identity: node.identity.clone(),
            kind: new_kind,
            ty: metrics,
        })
    }
}