Brummel/RustAst
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor specialization to handle RTL lookup

Browse Source 
The specialization logic has been refactored to include a call to
`rtl_lookup` for host functions. This allows the specializer to
recognize and optimize calls to built-in runtime functions by directly
returning their specialized values.

Additionally, the `_rtl_lookup` field in `Specializer` has been renamed
to `rtl_lookup` for clarity. The `visit_call` method has also been
renamed to `specialize_call_logic` to better reflect its purpose, and it
now correctly handles `TailCall` nodes by preserving them.
This commit is contained in:
Michael Schimmel
2026-02-19 16:39:39 +01:00
parent e7628e5cdf
commit b8390a3431
1 changed files with 116 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ast/compiler/specializer.rs
+116 -23
View File
@@ -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<HashMap<MonoCacheKey, (Value, StaticType)>>,
registry: Option<Rc<dyn FunctionRegistry>>,
compiler: Option<CompileFunc>,
_rtl_lookup: Option<RtlLookupFunc>,
rtl_lookup: Option<RtlLookupFunc>,
}
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<TypedNode>, original_ty: StaticType) -> (BoundKind<StaticType>, StaticType) {
fn specialize_call_logic(&self, callee: TypedNode, args: Vec<TypedNode>, original_ty: StaticType) -> (TypedNode, Vec<TypedNode>, 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();
@@ -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<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);
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);
}
}
}