RustAst/src/ast/compiler/specializer.rs

use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;
use crate::ast::types::{StaticType, Value, Signature};
use crate::ast::compiler::bound_nodes::{BoundKind, Address, TypedNode, BoundNode};
use crate::ast::nodes::Node;

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

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

pub trait FunctionRegistry {
    fn resolve(&self, addr: Address) -> Option<Rc<BoundNode>>;
}

pub struct Specializer {
    cache: RefCell<HashMap<MonoCacheKey, (Value, StaticType)>>,
    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>,
    ) -> Self {
        Self {
            cache: RefCell::new(HashMap::new()),
            registry,
            compiler,
            _rtl_lookup: rtl_lookup,
        }
    }

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

    fn visit_node(&self, node: TypedNode) -> TypedNode {
        let (new_kind, new_ty) = match node.kind {
            BoundKind::Call { callee, args } => self.visit_call(*callee, args, node.ty.clone()),
            
            // Recursive traversal for other nodes
            BoundKind::If { cond, then_br, else_br } => {
                let cond = Box::new(self.visit_node(*cond));
                let then_br = Box::new(self.visit_node(*then_br));
                let else_br = else_br.map(|e| Box::new(self.visit_node(*e)));
                (BoundKind::If { cond, then_br, else_br }, node.ty)
            },
            BoundKind::Block { exprs } => {
                let exprs = exprs.into_iter().map(|e| self.visit_node(e)).collect();
                (BoundKind::Block { exprs }, node.ty)
            },
            BoundKind::Lambda { param_count, upvalues, body } => {
                // We do NOT specialize inside lambdas automatically unless called?
                // Actually, we should specialize the body as generic code.
                // But without known types for parameters, we can't do much deep specialization.
                // Delphi code: "Cannot specialize closures safely without more complex analysis".
                // We'll just visit the body essentially.
                // Wait, if we visit body, we might specialize calls inside it that don't depend on params.
                let body = Rc::new(self.visit_node((*body).clone())); 
                (BoundKind::Lambda { param_count, upvalues, body }, node.ty)
            },
            BoundKind::DefLocal { name, slot, value, captured_by } => {
                let value = Box::new(self.visit_node(*value));
                (BoundKind::DefLocal { name, slot, value, captured_by }, node.ty)
            },
            BoundKind::DefGlobal { name, global_index, value } => {
                let value = Box::new(self.visit_node(*value));
                (BoundKind::DefGlobal { name, global_index, value }, node.ty)
            },
            BoundKind::Set { addr, value } => {
                let value = Box::new(self.visit_node(*value));
                (BoundKind::Set { addr, value }, node.ty)
            },
            BoundKind::Tuple { elements } => {
                let elements = elements.into_iter().map(|e| self.visit_node(e)).collect();
                (BoundKind::Tuple { elements }, node.ty)
            },
            BoundKind::Map { entries } => {
                let entries = entries.into_iter().map(|(k, v)| (self.visit_node(k), self.visit_node(v))).collect();
                (BoundKind::Map { entries }, node.ty)
            },
            BoundKind::Expansion { original_call, bound_expanded } => {
                let bound_expanded = Box::new(self.visit_node(*bound_expanded));
                (BoundKind::Expansion { original_call, bound_expanded }, node.ty)
            },
            
            // Leaf nodes or uninteresting nodes
            k => (k, node.ty),
        };

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

    fn visit_call(&self, callee: TypedNode, args: Vec<TypedNode>, original_ty: StaticType) -> (BoundKind<StaticType>, StaticType) {
        // 1. Specialize children first
        let new_callee = self.visit_node(callee);
        let new_args: Vec<TypedNode> = args.into_iter().map(|a| self.visit_node(a)).collect();

        // 2. Check if this call is a candidate (Callee is Get(Address))
        let address = if let BoundKind::Get { addr, .. } = &new_callee.kind {
            *addr
        } else {
            // Not a direct call to a named function/variable
            return (BoundKind::Call { callee: Box::new(new_callee), args: new_args }, original_ty);
        };

        // 3. Check if all argument types are statically known
        let arg_types: Vec<StaticType> = new_args.iter().map(|a| a.ty.clone()).collect();
        if arg_types.iter().any(|t| matches!(t, StaticType::Any)) {
            // Cannot specialize with unknown types
            return (BoundKind::Call { callee: Box::new(new_callee), args: new_args }, original_ty);
        }

        // --- Optimization Candidate ---
        let key = MonoCacheKey { address, arg_types: arg_types.clone() };

        // 4. Check Cache
        if let Some((val, ret_ty)) = self.cache.borrow().get(&key) {
            // Cache Hit! Replace Callee with Constant(Function)
            let specialized_callee = Node {
                identity: new_callee.identity.clone(),
                kind: BoundKind::Constant(val.clone()),
                ty: StaticType::Function(Box::new(Signature {
                    params: arg_types,
                    ret: ret_ty.clone(),
                })),
            };
            return (BoundKind::Call { callee: Box::new(specialized_callee), args: new_args }, ret_ty.clone());
        }

        // 5. Check RTL (Host Functions) - TODO: Need Name lookup from Address?
        // Wait, Address::Global(idx) -> Name? We don't have the Name here easily unless we look up in global map.
        // But RTL functions are usually bound to Globals.
        // IF the Registry can give us the Name, we can look up RTL.
        // For now, let's assume we can resolve the function definition.
        
        // 6. Resolve Function Definition
        if let Some(func_node) = self.registry.as_ref().and_then(|r| r.resolve(address)) {
            // Check constraints (no closures with state)
            if let BoundKind::Lambda { upvalues, .. } = &func_node.kind {
                if !upvalues.is_empty() {
                        // Cannot specialize stateful closures trivially
                    return (BoundKind::Call { callee: Box::new(new_callee), args: new_args }, original_ty);
                }
            } else {
                // Not a lambda?
                return (BoundKind::Call { callee: Box::new(new_callee), args: new_args }, original_ty);
            }

            // 7. Compile Specialization (User Code)
            if let Some(compiler) = &self.compiler {
                match compiler(func_node, &arg_types) {
                    Ok((compiled_val, ret_ty)) => {
                        // Store in cache
                        self.cache.borrow_mut().insert(key, (compiled_val.clone(), ret_ty.clone()));

                        let specialized_callee = Node {
                            identity: new_callee.identity.clone(),
                            kind: BoundKind::Constant(compiled_val),
                            ty: StaticType::Function(Box::new(Signature {
                                params: arg_types,
                                ret: ret_ty.clone(),
                            })),
                        };
                        return (BoundKind::Call { callee: Box::new(specialized_callee), args: new_args }, ret_ty);
                    },
                    Err(_) => {
                        // Fallback on error
                    }
                }
            }
        }

        // Fallback: Dynamic Call
        (BoundKind::Call { callee: Box::new(new_callee), args: new_args }, original_ty)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::types::{Identity, NodeIdentity, SourceLocation, StaticType, Value, Signature};
    use crate::ast::compiler::bound_nodes::{BoundKind, Address, BoundNode, TypedNode};
    use crate::ast::nodes::Symbol;
    use std::rc::Rc;

    fn make_identity() -> Identity {
        Rc::new(NodeIdentity { location: SourceLocation { line: 0, col: 0 } })
    }

    fn make_typed_node(kind: BoundKind<StaticType>, ty: StaticType) -> TypedNode {
        crate::ast::nodes::Node {
            identity: make_identity(),
            kind,
            ty,
        }
    }

    // Mock Registry
    struct MockRegistry {
        functions: HashMap<Address, Rc<BoundNode>>,
    }
    
    impl MockRegistry {
        fn new() -> Self {
            Self { functions: HashMap::new() }
        }
        
        fn register(&mut self, addr: Address, node: BoundNode) {
            self.functions.insert(addr, Rc::new(node));
        }
    }
    
    impl FunctionRegistry for MockRegistry {
        fn resolve(&self, addr: Address) -> Option<Rc<BoundNode>> {
            self.functions.get(&addr).cloned()
        }
    }

    #[test]
    fn test_specialize_compiles_user_function() {
        // Setup Registry with a function definition
        let mut registry = MockRegistry::new();
        let addr = Address::Local(0);
        let name = Symbol::from("test_func");
        
        // Def: (fn [x] x) -- generic identity
        let func_node = BoundNode {
            identity: make_identity(),
            kind: BoundKind::Lambda { 
                param_count: 1, 
                upvalues: vec![], 
                body: Rc::new(BoundNode { identity: make_identity(), kind: BoundKind::Nop, ty: () }) 
            },
            ty: ()
        };
        registry.register(addr, func_node);

        // Setup Compiler Mock
        let compiler: CompileFunc = Rc::new(|_node: Rc<BoundNode>, _args: &[StaticType]| -> Result<(Value, StaticType), String> {
            // Return a specialized "compiled" value
            Ok((Value::Int(12345), StaticType::Int))
        });

        let spec = Specializer::new(Some(Rc::new(registry)), Some(compiler), None);
        
        // Call(Get(Local(0)), [Arg(Int)])
        let callee = make_typed_node(BoundKind::Get { addr, name: name.clone() }, StaticType::Any);
        let arg = make_typed_node(BoundKind::Constant(Value::Int(1)), StaticType::Int);
        
        let call_node = make_typed_node(
            BoundKind::Call { callee: Box::new(callee), args: vec![arg] },
            StaticType::Any
        );
        
        let result = spec.specialize(call_node);
        
        // Should be Call(Constant(12345), ...)
        if let BoundKind::Call { callee, .. } = result.kind {
            if let BoundKind::Constant(val) = callee.kind {
                match val {
                    Value::Int(12345) => (), 
                    _ => panic!("Expected compiled value 12345"),
                }
            } else {
                panic!("Expected Constant callee");
            }
        } else {
            panic!("Expected Call node");
        }
    }

    #[test]
    fn test_specialize_skips_unknown_types() {
        let spec = Specializer::new(None, None, None);
        let name0 = Symbol::from("f");
        let name1 = Symbol::from("x");
        
        // Call(Get(Local(0)), [Get(Local(1))]) where arg is Any
        let callee = make_typed_node(BoundKind::Get { addr: Address::Local(0), name: name0 }, StaticType::Function(Box::new(Signature { params: vec![StaticType::Any], ret: StaticType::Void })));
        let arg = make_typed_node(BoundKind::Get { addr: Address::Local(1), name: name1 }, StaticType::Any);
        
        let call_node = make_typed_node(
            BoundKind::Call { callee: Box::new(callee), args: vec![arg] },
            StaticType::Void
        );
        
        let result = spec.specialize(call_node);
        
        // Should remain a generic Call because arg type is Any
        if let BoundKind::Call { callee, .. } = result.kind {
            if let BoundKind::Get { .. } = callee.kind {
                // Correct: Still a Get, not a Constant(Function)
            } else {
                panic!("Expected generic Call to Get, got {:?}", callee.kind);
            }
        } else {
            panic!("Expected Call node");
        }
    }

    #[test]
    fn test_specialize_uses_cache() {
        // Setup cache with a pre-specialized function for (Int) -> Int
        let spec = Specializer::new(None, None, None);
        
        let addr = Address::Local(0);
        let name = Symbol::from("cached_func");
        let arg_types = vec![StaticType::Int];
        let key = MonoCacheKey { address: addr, arg_types: arg_types.clone() };
        
        // Mock a specialized function pointer
        let specialized_val = Value::Int(999); // Dummy value representing function
        let ret_ty = StaticType::Int;
        
        spec.cache.borrow_mut().insert(key, (specialized_val.clone(), ret_ty.clone()));
        
        // Create the call node: Call(Get(0), [Arg(Int)])
        let callee = make_typed_node(BoundKind::Get { addr, name }, StaticType::Any);
        let arg = make_typed_node(BoundKind::Constant(Value::Int(1)), StaticType::Int);
        
        let call_node = make_typed_node(
            BoundKind::Call { callee: Box::new(callee), args: vec![arg] },
            StaticType::Any
        );
        
        let result = spec.specialize(call_node);
        
        // Should now be Call(Constant(999), ...)
        if let BoundKind::Call { callee, .. } = result.kind {
            if let BoundKind::Constant(val) = callee.kind {
                match val {
                    Value::Int(999) => (), // Success
                    _ => panic!("Expected specialized value 999"),
                }
            } else {
                panic!("Expected Constant callee, got {:?}", callee.kind);
            }
        } else {
            panic!("Expected Call node");
        }
    }
}