diff --git a/src/ast/compiler/specializer.rs b/src/ast/compiler/specializer.rs index 2440a29..ce48485 100644 --- a/src/ast/compiler/specializer.rs +++ b/src/ast/compiler/specializer.rs @@ -1,4 +1,3 @@ - use std::collections::HashMap; use std::rc::Rc; use std::cell::RefCell; @@ -23,7 +22,7 @@ pub struct Specializer { cache: RefCell>, registry: Option>, compiler: Option, - _rtl_lookup: Option, + rtl_lookup: Option, } impl Specializer { @@ -36,7 +35,7 @@ impl Specializer { cache: RefCell::new(HashMap::new()), registry, compiler, - _rtl_lookup: rtl_lookup, + rtl_lookup, } } @@ -46,8 +45,16 @@ impl Specializer { 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()), + 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: 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: new_args }, ret_ty) + }, + // Recursive traversal for other nodes BoundKind::If { cond, then_br, else_br } => { let cond = Box::new(self.visit_node(*cond)); @@ -60,12 +67,6 @@ impl Specializer { (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) }, @@ -105,7 +106,7 @@ impl Specializer { } } - fn visit_call(&self, callee: TypedNode, args: Vec, original_ty: StaticType) -> (BoundKind, StaticType) { + fn specialize_call_logic(&self, callee: TypedNode, args: Vec, original_ty: StaticType) -> (TypedNode, Vec, StaticType) { // 1. Specialize children first let new_callee = self.visit_node(callee); let new_args: Vec = args.into_iter().map(|a| self.visit_node(a)).collect(); @@ -115,14 +116,14 @@ impl Specializer { *addr } else { // Not a direct call to a named function/variable - return (BoundKind::Call { callee: Box::new(new_callee), args: new_args }, original_ty); + return (new_callee, new_args, original_ty); }; // 3. Check if all argument types are statically known let arg_types: Vec = 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); + return (new_callee, new_args, original_ty); } // --- Optimization Candidate --- @@ -139,14 +140,27 @@ impl Specializer { ret: ret_ty.clone(), })), }; - return (BoundKind::Call { callee: Box::new(specialized_callee), args: new_args }, ret_ty.clone()); + return (specialized_callee, 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. + // 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: 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)) { @@ -154,11 +168,11 @@ impl Specializer { 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); + return (new_callee, new_args, original_ty); } } else { // Not a lambda? - return (BoundKind::Call { callee: Box::new(new_callee), args: new_args }, original_ty); + return (new_callee, new_args, original_ty); } // 7. Compile Specialization (User Code) @@ -176,7 +190,7 @@ impl Specializer { ret: ret_ty.clone(), })), }; - return (BoundKind::Call { callee: Box::new(specialized_callee), args: new_args }, ret_ty); + return (specialized_callee, new_args, ret_ty); }, Err(_) => { // Fallback on error @@ -186,7 +200,7 @@ impl Specializer { } // Fallback: Dynamic Call - (BoundKind::Call { callee: Box::new(new_callee), args: new_args }, original_ty) + (new_callee, new_args, original_ty) } } @@ -354,4 +368,83 @@ mod tests { 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)); + + 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 call_node = make_typed_node( + BoundKind::Call { callee: Box::new(callee), args: vec![arg] }, + 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); + + 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); + + // Use TailCall here + let call_node = make_typed_node( + BoundKind::TailCall { callee: Box::new(callee), args: vec![arg] }, + 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); + } + } }