- Adjust Environment to use BoundNode for the function registry and

correctly initialize the `TypeChecker` with argument types during
  macro expansion.
- Refactor `Specializer::compile` to perform type checking with provided
  arguments before specialization and to correctly extract the return
  type.
- Enhance the `Dumper` to introspect and display specialized closure
  bodies.
- Update `LambdaCollector` to use `BoundNode` consistently.
- Modify `TypeChecker` to accept and inject specialized argument types
  for lambdas.
This commit is contained in:
Michael Schimmel
2026-02-19 23:18:04 +01:00
parent 16d9d41e3d
commit 3c0f2ec8ce
6 changed files with 147 additions and 60 deletions
+27 -1
View File
@@ -1,5 +1,7 @@
use crate::ast::nodes::Node;
use crate::ast::compiler::bound_nodes::BoundKind;
use crate::ast::types::Value;
use crate::ast::vm::Closure;
use std::fmt::Debug;
/// Human-readable AST dumper for the bound AST.
@@ -34,7 +36,31 @@ impl Dumper {
fn visit<T: Debug>(&mut self, node: &Node<BoundKind<T>, T>) {
match &node.kind {
BoundKind::Nop => self.log("Nop", node),
BoundKind::Constant(v) => self.log(&format!("Constant: {}", v), node),
BoundKind::Constant(v) => {
self.log(&format!("Constant: {}", v), node);
// Introspect Closure AST if possible
if let Value::Object(obj) = v {
if let Some(closure) = obj.as_any().downcast_ref::<Closure>() {
self.indent += 1;
self.write_indent();
self.output.push_str("--- Specialized Body ---\n");
// We need to cast the inner TypedNode to the generic T required by visit.
// Since Dumper is generic over T, but Closure stores TypedNode (where T = StaticType),
// we can only fully dump if T is StaticType.
// However, we can hack it by creating a new Dumper for the inner AST string.
// We can't call self.visit because types mismatch if T != StaticType.
// So we just recursively dump to string and append.
let inner_dump = Dumper::dump(&closure.function_node);
for line in inner_dump.lines() {
self.write_indent();
self.output.push_str(line);
self.output.push('\n');
}
self.indent -= 1;
}
}
},
BoundKind::Get { addr, name } => self.log(&format!("Get: {} ({:?})", name.name, addr), node),
BoundKind::Set { addr, value } => {
+5 -8
View File
@@ -1,21 +1,20 @@
use std::collections::HashMap;
use crate::ast::compiler::TypedNode;
use crate::ast::compiler::bound_nodes::{BoundKind, Address};
use crate::ast::compiler::bound_nodes::{BoundKind, Address, BoundNode};
/// A pass that collects all global function definitions (lambdas) into a registry.
/// This allows the Specializer to retrieve the original AST of a function for monomorphization.
pub struct LambdaCollector<'a> {
registry: &'a mut HashMap<u32, TypedNode>,
registry: &'a mut HashMap<u32, BoundNode>,
}
impl<'a> LambdaCollector<'a> {
/// Performs a full traversal of the AST and populates the provided registry.
pub fn collect(node: &TypedNode, registry: &'a mut HashMap<u32, TypedNode>) {
pub fn collect(node: &BoundNode, registry: &'a mut HashMap<u32, BoundNode>) {
let mut collector = Self { registry };
collector.visit(node);
}
fn visit(&mut self, node: &TypedNode) {
fn visit(&mut self, node: &BoundNode) {
match &node.kind {
BoundKind::Block { exprs } => {
for expr in exprs {
@@ -50,8 +49,6 @@ impl<'a> LambdaCollector<'a> {
}
BoundKind::Lambda { body, .. } => {
// Nested functions are not yet supported for global specialization
// but we traverse them to find potential global definitions inside (if allowed).
self.visit(body);
}
@@ -83,7 +80,7 @@ impl<'a> LambdaCollector<'a> {
self.visit(bound_expanded);
}
_ => {} // Leaf nodes (Constant, Get, Nop, etc.)
_ => {} // Leaf nodes
}
}
}
+10 -7
View File
@@ -2,7 +2,7 @@ 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};
use crate::ast::compiler::bound_nodes::{BoundKind, Address, TypedNode, BoundNode};
use crate::ast::nodes::Node;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -11,11 +11,11 @@ pub struct MonoCacheKey {
pub arg_types: Vec<StaticType>,
}
pub type CompileFunc = Rc<dyn Fn(TypedNode, &[StaticType]) -> Result<(Value, StaticType), String>>;
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<TypedNode>;
fn resolve(&self, addr: Address) -> Option<BoundNode>;
}
pub type MonoCache = HashMap<MonoCacheKey, (Value, StaticType)>;
@@ -182,18 +182,21 @@ impl Specializer {
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, (compiled_val.clone(), ret_ty.clone()));
self.cache.borrow_mut().insert(key, (res_val.clone(), res_ty.clone()));
let specialized_callee = Node {
identity: new_callee.identity.clone(),
kind: BoundKind::Constant(compiled_val),
kind: BoundKind::Constant(res_val),
ty: StaticType::Function(Box::new(Signature {
params: arg_types,
ret: ret_ty.clone(),
ret: res_ty.clone(),
})),
};
return (specialized_callee, new_args, ret_ty);
return (specialized_callee, new_args, res_ty);
},
Err(_) => {
// Fallback on error
+60 -6
View File
@@ -46,12 +46,66 @@ impl TypeChecker {
Self { global_types }
}
pub fn check(&self, node: BoundNode) -> Result<TypedNode, String> {
// Start with a root context. Root scope has no upvalues.
// We assume 1000 slots for global script level if needed, but Binder handles it.
// Actually, Binder already assigned slot indices. We need to know the max slot.
let mut ctx = TypeContext::new(256, vec![], None);
self.check_node(node, &mut ctx)
pub fn check(&self, node: BoundNode, arg_types: &[StaticType]) -> Result<TypedNode, String> {
match node.kind {
BoundKind::Lambda { param_count, upvalues, body } => {
// 1. Determine types of captured variables (Root lambdas have none)
let mut upvalue_types = Vec::with_capacity(upvalues.len());
for _ in &upvalues {
upvalue_types.push(StaticType::Any);
}
// 2. Create the specialized context
let root_ctx = TypeContext::new(0, vec![], None);
let mut lambda_ctx = TypeContext::new(param_count + 64, upvalue_types, Some(&root_ctx));
// 3. INJECT specialized argument types into slots
for (i, ty) in arg_types.iter().enumerate() {
if (i as u32) < param_count {
lambda_ctx.set_local_type(i as u32, ty.clone());
}
}
// 4. Check body with the new types
let body_typed = self.check_node((*body).clone(), &mut lambda_ctx)?;
let ret_ty = body_typed.ty.clone();
// 5. Construct specialized function type
let final_params = if arg_types.is_empty() {
vec![StaticType::Any; param_count as usize]
} else {
arg_types.to_vec()
};
let fn_ty = StaticType::Function(Box::new(crate::ast::types::Signature {
params: final_params,
ret: ret_ty,
}));
Ok(Node {
identity: node.identity,
kind: BoundKind::Lambda {
param_count,
upvalues,
body: Rc::new(body_typed)
},
ty: fn_ty,
})
}
_ => {
// Fallback: Wrap in implicit lambda
let virtual_lambda = BoundNode {
identity: node.identity.clone(),
kind: BoundKind::Lambda {
param_count: 0,
upvalues: vec![],
body: Rc::new(node)
},
ty: (),
};
self.check(virtual_lambda, &[])
}
}
}
fn check_node(&self, node: BoundNode, ctx: &mut TypeContext) -> Result<TypedNode, String> {