Formatting
This commit is contained in:
Michael Schimmel
+313
-232
@@ -1,232 +1,313 @@
|
||||
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)
|
||||
},
|
||||
|
||||
// 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::Record { fields } => {
|
||||
let fields = fields.into_iter().map(|(k, v)| (self.visit_node(k), self.visit_node(v))).collect();
|
||||
(BoundKind::Record { fields }, 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() {
|
||||
return (new_callee, new_args, original_ty);
|
||||
}
|
||||
} else {
|
||||
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.
|
||||
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],
|
||||
}
|
||||
}
|
||||
}
|
||||
use crate::ast::compiler::bound_nodes::{Address, BoundKind, BoundNode, TypedNode};
|
||||
use crate::ast::nodes::Node;
|
||||
use crate::ast::types::{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,
|
||||
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,
|
||||
)
|
||||
}
|
||||
|
||||
// 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::Record { fields } => {
|
||||
let fields = fields
|
||||
.into_iter()
|
||||
.map(|(k, v)| (self.visit_node(k), self.visit_node(v)))
|
||||
.collect();
|
||||
(BoundKind::Record { fields }, 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() {
|
||||
return (new_callee, new_args, original_ty);
|
||||
}
|
||||
} else {
|
||||
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.
|
||||
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],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user