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(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<BoundNode>;
}

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: 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 } => {
                let (new_callee, new_args, ret_ty) = self.specialize_call_logic(*callee, *args, node.ty.clone());
                (BoundKind::Call { callee: Box::new(new_callee), args: Box::new(new_args) }, ret_ty)
            },
            
            BoundKind::TailCall { callee, args } => {
                let (new_callee, new_args, ret_ty) = self.specialize_call_logic(*callee, *args, node.ty.clone());
                (BoundKind::TailCall { callee: Box::new(new_callee), args: Box::new(new_args) }, ret_ty)
            },

            // 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 { params, upvalues, body, positional_count } => {
                let params = Rc::new(self.visit_node(params.as_ref().clone()));
                let body = Rc::new(self.visit_node((*body).clone())); 
                (BoundKind::Lambda { params, upvalues, body, positional_count }, 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 specialize_call_logic(&self, callee: TypedNode, args: TypedNode, original_ty: StaticType) -> (TypedNode, TypedNode, StaticType) {
        // 1. Specialize children first
        let new_callee = self.visit_node(callee);
        let new_args = self.visit_node(args);

        // 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 (new_callee, new_args, original_ty);
        };

        // 3. Check if all argument types are statically known
        let arg_types: Vec<StaticType> = if let StaticType::Tuple(elements) = &new_args.ty {
            elements.clone()
        } else {
            vec![new_args.ty.clone()]
        };

        if arg_types.iter().any(|t| matches!(t, StaticType::Any)) {
            // Cannot specialize with unknown types
            return (new_callee, 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: StaticType::Tuple(arg_types),
                    ret: ret_ty.clone(),
                })),
            };
            return (specialized_callee, new_args, ret_ty.clone());
        }

        // 5. Check RTL (Host Functions)
        if let Some(rtl_lookup) = &self.rtl_lookup 
            && let BoundKind::Get { name, .. } = &new_callee.kind 
            && let Some((val, ret_ty)) = rtl_lookup(&name.name, &arg_types) 
        {
             // Cache Hit (RTL)
             self.cache.borrow_mut().insert(key.clone(), (val.clone(), ret_ty.clone()));
             
             let specialized_callee = Node {
                identity: new_callee.identity.clone(),
                kind: BoundKind::Constant(val.clone()),
                ty: StaticType::Function(Box::new(Signature {
                    params: StaticType::Tuple(arg_types),
                    ret: ret_ty.clone(),
                })),
            };
            return (specialized_callee, new_args, ret_ty);
        }
        
        // 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 (new_callee, new_args, original_ty);
                }
            } else {
                // Not a lambda?
                return (new_callee, 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)) => {
                        let res_val: Value = compiled_val;
                        let res_ty: StaticType = ret_ty;
                        
                        // Store in cache
                        self.cache.borrow_mut().insert(key, (res_val.clone(), res_ty.clone()));

                        // PERFORMANCE: Flatten the argument tuple to match the specialized signature.
                        // Since we are specializing, we can convert [[1 2] 3] into a flat [1 2 3] Tuple node.
                        let flat_elements = self.flatten_tuple(new_args.clone());
                        let flat_types = flat_elements.iter().map(|e| e.ty.clone()).collect();
                        let flattened_args = Node {
                            identity: new_args.identity.clone(),
                            kind: BoundKind::Tuple { elements: flat_elements },
                            ty: StaticType::Tuple(flat_types),
                        };

                        let specialized_callee = Node {
                            identity: new_callee.identity.clone(),
                            kind: BoundKind::Constant(res_val),
                            ty: StaticType::Function(Box::new(Signature {
                                params: flattened_args.ty.clone(),
                                ret: res_ty.clone(),
                            })),
                        };
                        return (specialized_callee, flattened_args, res_ty);
                    },
                    Err(_) => {
                        // Fallback on error
                    }
                }
            }
        }

        // Fallback: Dynamic Call
        (new_callee, new_args, original_ty)
    }

    fn flatten_tuple(&self, node: TypedNode) -> Vec<TypedNode> {
        match node.kind {
            BoundKind::Tuple { elements } => {
                let mut flat = Vec::new();
                for el in elements {
                    flat.extend(self.flatten_tuple(el));
                }
                flat
            }
            _ => vec![node],
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::types::{Identity, NodeIdentity, SourceLocation, StaticType, Value, Signature};
    use crate::ast::compiler::bound_nodes::{BoundKind, Address, 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, BoundNode>,
    }
    
    impl MockRegistry {
        fn new() -> Self {
            Self { functions: HashMap::new() }
        }
        
        fn register(&mut self, addr: Address, node: BoundNode) {
            self.functions.insert(addr, node);
        }
    }
    
    impl FunctionRegistry for MockRegistry {
        fn resolve(&self, addr: Address) -> Option<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 { 
                params: Rc::new(BoundNode {
                    identity: make_identity(),
                    kind: BoundKind::Tuple { 
                        elements: vec![
                            BoundNode {
                                identity: make_identity(),
                                kind: BoundKind::Parameter { name: name.clone(), slot: 0 },
                                ty: ()
                            }
                        ]
                    },
                    ty: ()
                }),
                upvalues: vec![], 
                body: Rc::new(BoundNode { identity: make_identity(), kind: BoundKind::Nop, ty: () }),
                positional_count: Some(1),
            },
            ty: ()
        };
        registry.register(addr, func_node);

        // Setup Compiler Mock
        let compiler: CompileFunc = Rc::new(|_node: 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, None);
        
        // Call(Get(Local(0)), Tuple([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 args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Int]));
        
        let call_node = make_typed_node(
            BoundKind::Call { callee: Box::new(callee), args: Box::new(args_tuple) },
            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, None);
        let name0 = Symbol::from("f");
        let name1 = Symbol::from("x");
        
        // Call(Get(Local(0)), Tuple([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: StaticType::Tuple(vec![StaticType::Any]), ret: StaticType::Void })));
        let arg = make_typed_node(BoundKind::Get { addr: Address::Local(1), name: name1 }, StaticType::Any);
        let args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Any]));
        
        let call_node = make_typed_node(
            BoundKind::Call { callee: Box::new(callee), args: Box::new(args_tuple) },
            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, 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), Tuple([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 args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Int]));
        
        let call_node = make_typed_node(
            BoundKind::Call { callee: Box::new(callee), args: Box::new(args_tuple) },
            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");
        }
    }

    #[test]
    fn test_specialize_uses_rtl_lookup() {
        // Setup RTL Lookup Mock
        let rtl_lookup: RtlLookupFunc = Rc::new(|name, _args| {
            if name == "rtl_func" {
                Some((Value::Int(888), StaticType::Int))
            } else {
                None
            }
        });
        
        let spec = Specializer::new(None, None, Some(rtl_lookup), None);
        
        let addr = Address::Global(10);
        let name = Symbol::from("rtl_func");
        
        let callee = make_typed_node(BoundKind::Get { addr, name }, StaticType::Any);
        let arg = make_typed_node(BoundKind::Constant(Value::Int(1)), StaticType::Int);
        let args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Int]));
        
        let call_node = make_typed_node(
            BoundKind::Call { callee: Box::new(callee), args: Box::new(args_tuple) },
            StaticType::Any
        );
        
        let result = spec.specialize(call_node);
        
        if let BoundKind::Call { callee, .. } = result.kind {
            if let BoundKind::Constant(val) = callee.kind {
                match val {
                    Value::Int(888) => (), 
                    _ => panic!("Expected RTL value 888"),
                }
            } else {
                panic!("Expected Constant callee from RTL");
            }
        } else {
            panic!("Expected Call node");
        }
    }

    #[test]
    fn test_specialize_preserves_tail_call() {
        let spec = Specializer::new(None, None, None, None);
        
        let addr = Address::Local(0);
        let name = Symbol::from("tail_func");
        let arg_types = vec![StaticType::Int];
        let key = MonoCacheKey { address: addr, arg_types: arg_types.clone() };
        
        let specialized_val = Value::Int(777);
        let ret_ty = StaticType::Int;
        
        spec.cache.borrow_mut().insert(key, (specialized_val.clone(), ret_ty.clone()));
        
        let callee = make_typed_node(BoundKind::Get { addr, name }, StaticType::Any);
        let arg = make_typed_node(BoundKind::Constant(Value::Int(1)), StaticType::Int);
        let args_tuple = make_typed_node(BoundKind::Tuple { elements: vec![arg] }, StaticType::Tuple(vec![StaticType::Int]));
        
        // Use TailCall here
        let call_node = make_typed_node(
            BoundKind::TailCall { callee: Box::new(callee), args: Box::new(args_tuple) },
            StaticType::Any
        );
        
        let result = spec.specialize(call_node);
        
        if let BoundKind::TailCall { callee, .. } = result.kind {
             if let BoundKind::Constant(val) = callee.kind {
                match val {
                    Value::Int(777) => (), 
                    _ => panic!("Expected specialized value 777"),
                }
            } else {
                panic!("Expected Constant callee");
            }
        } else {
            panic!("Expected TailCall node, got {:?}", result.kind);
        }
    }
}