RustAst/src/ast/compiler/specializer.rs

use crate::ast::nodes::{
    Address, AnalyzedNode, AnalyzedPhase, IdentifierBinding, Node, NodeKind, NodeMetrics, VirtualId,
};
use crate::ast::types::{Purity, Signature, StaticType, Value};
use std::cell::RefCell;
use std::collections::HashMap;
use std::rc::Rc;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct MonoCacheKey {
    pub address: Address<VirtualId>,
    pub arg_types: Vec<StaticType>,
}

pub type CompileFunc = Rc<dyn Fn(Rc<Node<AnalyzedPhase>>, &[StaticType]) -> Result<(Value, StaticType), String>>;
pub type RtlLookupFunc = Rc<dyn Fn(&str, &[StaticType]) -> Option<(Value, StaticType)>>;

pub trait FunctionRegistry {
    fn resolve(&self, addr: Address<VirtualId>) -> Option<Rc<Node<AnalyzedPhase>>>;
    fn resolve_analyzed(&self, _addr: Address<VirtualId>) -> Option<Rc<AnalyzedNode>> {
        None
    }
}

pub type MonoCache = HashMap<MonoCacheKey, (Value, StaticType)>;

pub struct Specializer {
    pub cache: Rc<RefCell<MonoCache>>,
    registry: Option<Rc<dyn FunctionRegistry>>,
    compiler: Option<CompileFunc>,
    rtl_lookup: Option<RtlLookupFunc>,
}

impl Specializer {
    pub fn new(
        registry: Option<Rc<dyn FunctionRegistry>>,
        compiler: Option<CompileFunc>,
        rtl_lookup: Option<RtlLookupFunc>,
        cache: Option<Rc<RefCell<MonoCache>>>,
    ) -> Self {
        Self {
            cache: cache.unwrap_or_else(|| Rc::new(RefCell::new(HashMap::new()))),
            registry,
            compiler,
            rtl_lookup,
        }
    }

    pub fn specialize(&self, node: AnalyzedNode) -> AnalyzedNode {
        self.visit_node(node)
    }

    fn visit_node(&self, node: AnalyzedNode) -> AnalyzedNode {
        let (new_kind, metrics) = match node.kind {
            NodeKind::Call { callee, args } => {
                let (new_callee, new_args, _ret_ty) =
                    self.specialize_call_logic(callee, args, node.ty.original.ty.clone());

                let new_metrics = node.ty.clone();
                (
                    NodeKind::Call {
                        callee: Rc::new(new_callee),
                        args: Rc::new(new_args),
                    },
                    new_metrics,
                )
            }

            NodeKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let cond = Rc::new(self.visit_node(cond.as_ref().clone()));
                let then_br = Rc::new(self.visit_node(then_br.as_ref().clone()));
                let else_br = else_br.map(|e| Rc::new(self.visit_node(e.as_ref().clone())));
                (
                    NodeKind::If {
                        cond,
                        then_br,
                        else_br,
                    },
                    node.ty.clone(),
                )
            }
            NodeKind::Block { exprs } => {
                let exprs = exprs.into_iter().map(|e| Rc::new(self.visit_node(e.as_ref().clone()))).collect();
                (NodeKind::Block { exprs }, node.ty.clone())
            }
            NodeKind::Lambda {
                params,
                body,
                info,
            } => {
                let params = Rc::new(self.visit_node(params.as_ref().clone()));
                let body = Rc::new(self.visit_node(body.as_ref().clone()));
                (
                    NodeKind::Lambda {
                        params,
                        body,
                        info,
                    },
                    node.ty.clone(),
                )
            }
            NodeKind::Def {
                pattern,
                value,
                info,
            } => {
                let value = Rc::new(self.visit_node(value.as_ref().clone()));
                (
                    NodeKind::Def {
                        pattern,
                        value,
                        info,
                    },
                    node.ty.clone(),
                )
            }
            NodeKind::Assign { target, value, info } => {
                let value = Rc::new(self.visit_node(value.as_ref().clone()));
                (NodeKind::Assign { target, value, info }, node.ty.clone())
            }
            NodeKind::Tuple { elements } => {
                let elements = elements.into_iter().map(|e| Rc::new(self.visit_node(e.as_ref().clone()))).collect();
                (NodeKind::Tuple { elements }, node.ty.clone())
            }
            NodeKind::Record { fields, layout } => {
                let fields = fields.into_iter().map(|(k, v)| (k, Rc::new(self.visit_node(v.as_ref().clone())))).collect();
                (NodeKind::Record { fields, layout }, node.ty.clone())
            }
            NodeKind::Expansion {
                original_call,
                expanded,
            } => {
                let expanded = Rc::new(self.visit_node(expanded.as_ref().clone()));
                (
                    NodeKind::Expansion {
                        original_call,
                        expanded,
                    },
                    node.ty.clone(),
                )
            }
            k => (k, node.ty.clone()),
        };

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

    fn specialize_call_logic(
        &self,
        callee: Rc<AnalyzedNode>,
        args: Rc<AnalyzedNode>,
        original_ty: StaticType,
    ) -> (AnalyzedNode, AnalyzedNode, StaticType) {
        let new_callee = self.visit_node(callee.as_ref().clone());
        let new_args = self.visit_node(args.as_ref().clone());

        let address = if let NodeKind::Identifier {
            binding: IdentifierBinding::Reference(addr),
            ..
        } = &new_callee.kind
        {
            *addr
        } else {
            return (new_callee, new_args, original_ty);
        };

        let arg_types: Vec<StaticType> =
            if let StaticType::Tuple(elements) = &new_args.ty.original.ty {
                elements.clone()
            } else {
                vec![new_args.ty.original.ty.clone()]
            };

        if arg_types.iter().any(|t| matches!(t, StaticType::Any)) {
            return (new_callee, new_args, original_ty);
        }

        let key = MonoCacheKey {
            address,
            arg_types: arg_types.clone(),
        };

        if let Some((val, ret_ty)) = self.cache.borrow().get(&key) {
            let specialized_callee = self.make_constant_node(
                val.clone(),
                StaticType::Function(Box::new(Signature {
                    params: StaticType::Tuple(arg_types),
                    ret: ret_ty.clone(),
                })),
                &new_callee,
            );
            return (specialized_callee, new_args, ret_ty.clone());
        }

        if let Some(rtl_lookup) = &self.rtl_lookup
            && let NodeKind::Identifier { symbol, .. } = &new_callee.kind
            && let Some((val, ret_ty)) = rtl_lookup(&symbol.name, &arg_types)
        {
            self.cache
                .borrow_mut()
                .insert(key.clone(), (val.clone(), ret_ty.clone()));
            let specialized_callee = self.make_constant_node(
                val.clone(),
                StaticType::Function(Box::new(Signature {
                    params: StaticType::Tuple(arg_types),
                    ret: ret_ty.clone(),
                })),
                &new_callee,
            );
            return (specialized_callee, new_args, ret_ty);
        }

        if let Some(registry) = &self.registry
            && let Some(func_node) = registry.resolve_analyzed(address)
            && func_node.ty.is_recursive
        {
            return (new_callee, new_args, original_ty);
        }

        if let Some(compiler) = &self.compiler
            && let Some(func_node) = self.registry.as_ref().and_then(|r| r.resolve(address))
            && let Ok((compiled_val, ret_ty)) = compiler(func_node, &arg_types)
        {
            self.cache
                .borrow_mut()
                .insert(key, (compiled_val.clone(), ret_ty.clone()));

            // Only replace the callee if the compiled value is actually a function/object.
            // If it's a scalar (like 30 from folding), we DON'T fold here.
            // We keep the Call but update the callee to the specialized version if it's an object.
            if let Value::Object(_) | Value::Function(_) = &compiled_val {
                let specialized_callee = self.make_constant_node(
                    compiled_val,
                    StaticType::Function(Box::new(Signature {
                        params: StaticType::Tuple(arg_types),
                        ret: ret_ty.clone(),
                    })),
                    &new_callee,
                );
                return (specialized_callee, new_args, ret_ty);
            }
        }

        (new_callee, new_args, original_ty)
    }

    fn make_constant_node(
        &self,
        val: Value,
        ty: StaticType,
        template: &AnalyzedNode,
    ) -> AnalyzedNode {
        let typed_original = Rc::new(Node {
            identity: template.identity.clone(),
            kind: NodeKind::Constant(val.clone()),
            ty: ty.clone(),
        });
        Node {
            identity: template.identity.clone(),
            kind: NodeKind::Constant(val),
            ty: NodeMetrics {
                original: typed_original,
                purity: Purity::Pure,
                is_recursive: false,
            },
        }
    }
}