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

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

#[derive(Default)]
pub struct SubstitutionMap {
    pub values: HashMap<Address, Value>,
    pub ast_substitutions: HashMap<Address, 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 add_ast_substitution(&mut self, addr: Address, node: AnalyzedNode) {
        self.ast_substitutions.insert(addr, 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
        }
    }

    pub fn reindex_upvalues(&self, node: AnalyzedNode, mapping: &[Option<u32>]) -> AnalyzedNode {
        let (new_kind, metrics) = match node.kind {
            BoundKind::Get { addr, name } => (
                BoundKind::Get {
                    addr: self.reindex_addr(addr, mapping),
                    name,
                },
                node.ty.clone(),
            ),
            BoundKind::Lambda {
                params,
                upvalues,
                body,
                positional_count,
            } => {
                let mut next_upvalues = Vec::new();
                for addr in upvalues {
                    next_upvalues.push(self.reindex_addr(addr, mapping));
                }
                (
                    BoundKind::Lambda {
                        params,
                        upvalues: next_upvalues,
                        body,
                        positional_count,
                    },
                    node.ty.clone(),
                )
            }
            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.clone(),
                )
            }
            BoundKind::Block { exprs } => {
                let exprs = exprs
                    .into_iter()
                    .map(|e| self.reindex_upvalues(e, mapping))
                    .collect();
                (BoundKind::Block { exprs }, node.ty.clone())
            }
            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.clone())
            }
            BoundKind::Define {
                name,
                addr,
                kind,
                value,
                captured_by,
            } => {
                let value = Box::new(self.reindex_upvalues(*value, mapping));
                (
                    BoundKind::Define {
                        name,
                        addr,
                        kind,
                        value,
                        captured_by,
                    },
                    node.ty.clone(),
                )
            }
            BoundKind::Set { addr, value } => {
                let value = Box::new(self.reindex_upvalues(*value, mapping));
                (BoundKind::Set { addr, value }, node.ty.clone())
            }
            BoundKind::Tuple { elements } => {
                let elements = elements
                    .into_iter()
                    .map(|e| self.reindex_upvalues(e, mapping))
                    .collect();
                (BoundKind::Tuple { elements }, node.ty.clone())
            }
            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.clone())
            }
            k => (k, node.ty.clone()),
        };
        Node {
            identity: node.identity.clone(),
            kind: new_kind,
            ty: metrics,
        }
    }
}