- 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
+16
View File
@@ -0,0 +1,16 @@
- Macro expansion
- Binding
- var lambdas
- Linking(var lambdas)
-Typechecker::check_with_args(template, ())
- lambdas = Lambda Collection <- fills registry
- Specializing call, say (tak int int int)
- found in lambdas: (tak any any any)->any
- var local_lambdas
- recurse Linking(var local_lambdas):
- TypeChecker::check_with_args(template, args)
- local_lambdas = Lambda Collection
- Specializing
....
- TCO
+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> {
+29 -38
View File
@@ -15,23 +15,23 @@ use crate::ast::compiler::macros::{MacroExpander, MacroRegistry, MacroEvaluator}
use crate::ast::compiler::specializer::{Specializer, MonoCache, FunctionRegistry};
use crate::ast::rtl;
use crate::ast::rtl::intrinsics;
use crate::ast::compiler::bound_nodes::{BoundKind, Address};
use crate::ast::compiler::bound_nodes::{Address, BoundNode};
pub struct Environment {
pub global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
pub global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
pub global_values: Rc<RefCell<Vec<Value>>>,
pub function_registry: Rc<RefCell<HashMap<u32, TypedNode>>>,
pub function_registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
pub monomorph_cache: Rc<RefCell<MonoCache>>,
pub debug_mode: bool,
}
struct EnvFunctionRegistry {
registry: Rc<RefCell<HashMap<u32, TypedNode>>>,
registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
}
impl FunctionRegistry for EnvFunctionRegistry {
fn resolve(&self, addr: Address) -> Option<TypedNode> {
fn resolve(&self, addr: Address) -> Option<BoundNode> {
if let Address::Global(idx) = addr {
self.registry.borrow().get(&idx).cloned()
} else {
@@ -40,15 +40,15 @@ impl FunctionRegistry for EnvFunctionRegistry {
}
}
/// Evaluator used during macro expansion to allow compile-time logic.
struct RuntimeMacroEvaluator {
global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
global_values: Rc<RefCell<Vec<Value>>>,
function_registry: Rc<RefCell<HashMap<u32, TypedNode>>>,
function_registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
}
impl MacroEvaluator for RuntimeMacroEvaluator {
fn evaluate(&self, node: &Node<UntypedKind>, bindings: &HashMap<Rc<str>, Node<UntypedKind>>) -> Result<Value, String> {
// 1. Check if it's a simple parameter substitution
@@ -61,7 +61,7 @@ impl MacroEvaluator for RuntimeMacroEvaluator {
let bound_ast = Binder::bind_root(self.global_names.clone(), node)?;
let checker = TypeChecker::new(self.global_types.clone());
let typed_ast = checker.check(bound_ast)?;
let typed_ast = checker.check(bound_ast, &[])?;
let mut vm = VM::new(self.global_values.clone());
vm.run(&typed_ast)
@@ -152,23 +152,22 @@ impl Environment {
// 4. Bind
let bound_ast = Binder::bind_root(self.global_names.clone(), &expanded_ast)?;
// 5. Type Check
// 5. Collect Lambdas (Populate the registry with untyped templates)
LambdaCollector::collect(&bound_ast, &mut self.function_registry.borrow_mut());
// 6. Type Check
let checker = TypeChecker::new(self.global_types.clone());
let typed_ast = checker.check(bound_ast)?;
let typed_ast = checker.check(bound_ast, &[])?;
Ok(typed_ast)
}
/// Backend: Optimization (TCO, etc.)
pub fn link(&self, node: TypedNode) -> TypedNode {
// 1. Collect Lambdas (Populate the registry for the specializer)
LambdaCollector::collect(&node, &mut self.function_registry.borrow_mut());
// 2. Specialize
// 1. Specialize
let specialized = self.specialize_node(node);
// let specialized = node;
// 3. Optimize
// 2. Optimize
TCO::optimize(specialized)
}
@@ -179,50 +178,42 @@ impl Environment {
let rtl_lookup = Rc::new(|name: &str, args: &[StaticType]| intrinsics::lookup(name, args));
// We need to construct a compiler callback that can recursively specialize and compile.
// To avoid complex self-capturing, we reconstruct the environment context needed.
let func_reg = self.function_registry.clone();
let mono_cache = self.monomorph_cache.clone();
let global_values = self.global_values.clone(); // Needed for VM/Closure creation
let global_values = self.global_values.clone();
let global_types = self.global_types.clone();
let compiler = Rc::new(move |func_node: TypedNode, _arg_types: &[StaticType]| -> Result<(Value, StaticType), String> {
// 1. Specialize the body (Recursive)
// We recreate the specializer context here.
// Note: This creates a new Specializer for each recursion, but they SHARE the 'mono_cache'.
let compiler = Rc::new(move |func_template: BoundNode, arg_types: &[StaticType]| -> Result<(Value, StaticType), String> {
// 1. Re-TypeCheck the template with concrete argument types
let checker = TypeChecker::new(global_types.clone());
let retyped_ast = checker.check(func_template, arg_types)?;
// 2. Specialize (Recursive)
let sub_registry = Rc::new(EnvFunctionRegistry { registry: func_reg.clone() });
let sub_rtl_lookup = Rc::new(|name: &str, args: &[StaticType]| intrinsics::lookup(name, args));
// Note: We are passing 'None' as compiler to the inner specializer for now to prevent infinite recursion on cycles.
// A robust implementation would handle the recursion cycle or use a shared compiler reference.
// For 'tak', the recursion is handled by the cache or dynamic fallback.
let sub_specializer = Specializer::new(
Some(sub_registry),
None, // recursive compilation limit (depth 1) for safety
None,
Some(sub_rtl_lookup),
Some(mono_cache.clone())
);
let specialized_ast = sub_specializer.specialize(func_node);
let specialized_ast = sub_specializer.specialize(retyped_ast);
// 2. Optimize (TCO)
// 3. Optimize (TCO)
let optimized_ast = TCO::optimize(specialized_ast);
// 3. Compile to Closure (VM)
// We run the VM once to evaluate the Lambda definition, producing a closure Value.
// 4. Compile to Value (VM)
let mut vm = VM::new(global_values.clone());
let compiled_val = match vm.run(&optimized_ast) {
Ok(v) => v,
Err(e) => return Err(format!("VM Error during specialization: {}", e)),
};
// We need the return type.
// For a Lambda, the type is stored in the node.
// But we need the return type of the FUNCTION (e.g. Int), not the type of the Lambda node (Method).
// Actually, Specializer expects (Value, ReturnType).
// If the specialized function returns Int, we return Int.
let ret_type = if let BoundKind::Lambda { body, .. } = &optimized_ast.kind {
body.ty.clone()
// 5. Determine correct return type from the newly inferred function signature
let ret_type = if let StaticType::Function(sig) = &optimized_ast.ty {
sig.ret.clone()
} else {
StaticType::Any
};