From 3c0f2ec8ce8586f8b47152d2fa29486da3c9ece9 Mon Sep 17 00:00:00 2001 From: Michael Schimmel Date: Thu, 19 Feb 2026 23:18:04 +0100 Subject: [PATCH] - 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. --- docs/Piepline.md | 16 +++++++ src/ast/compiler/dumper.rs | 28 +++++++++++- src/ast/compiler/lambda_collector.rs | 13 +++--- src/ast/compiler/specializer.rs | 17 ++++--- src/ast/compiler/type_checker.rs | 66 ++++++++++++++++++++++++--- src/ast/environment.rs | 67 ++++++++++++---------------- 6 files changed, 147 insertions(+), 60 deletions(-) create mode 100644 docs/Piepline.md diff --git a/docs/Piepline.md b/docs/Piepline.md new file mode 100644 index 0000000..38eb0f4 --- /dev/null +++ b/docs/Piepline.md @@ -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 diff --git a/src/ast/compiler/dumper.rs b/src/ast/compiler/dumper.rs index 84efaed..faace0f 100644 --- a/src/ast/compiler/dumper.rs +++ b/src/ast/compiler/dumper.rs @@ -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(&mut self, node: &Node, 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::() { + 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 } => { diff --git a/src/ast/compiler/lambda_collector.rs b/src/ast/compiler/lambda_collector.rs index 257e2f2..6cf4e0a 100644 --- a/src/ast/compiler/lambda_collector.rs +++ b/src/ast/compiler/lambda_collector.rs @@ -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, + registry: &'a mut HashMap, } 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) { + pub fn collect(node: &BoundNode, registry: &'a mut HashMap) { 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 } } } diff --git a/src/ast/compiler/specializer.rs b/src/ast/compiler/specializer.rs index 3c7f4b5..c2ae00c 100644 --- a/src/ast/compiler/specializer.rs +++ b/src/ast/compiler/specializer.rs @@ -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, } -pub type CompileFunc = Rc Result<(Value, StaticType), String>>; +pub type CompileFunc = Rc Result<(Value, StaticType), String>>; pub type RtlLookupFunc = Rc Option<(Value, StaticType)>>; pub trait FunctionRegistry { - fn resolve(&self, addr: Address) -> Option; + fn resolve(&self, addr: Address) -> Option; } pub type MonoCache = HashMap; @@ -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 diff --git a/src/ast/compiler/type_checker.rs b/src/ast/compiler/type_checker.rs index 49bf810..59a42b7 100644 --- a/src/ast/compiler/type_checker.rs +++ b/src/ast/compiler/type_checker.rs @@ -46,12 +46,66 @@ impl TypeChecker { Self { global_types } } - pub fn check(&self, node: BoundNode) -> Result { - // 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 { + 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 { diff --git a/src/ast/environment.rs b/src/ast/environment.rs index 44dbd73..9d77d68 100644 --- a/src/ast/environment.rs +++ b/src/ast/environment.rs @@ -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>>, pub global_types: Rc>>, pub global_values: Rc>>, - pub function_registry: Rc>>, + pub function_registry: Rc>>, pub monomorph_cache: Rc>, pub debug_mode: bool, } struct EnvFunctionRegistry { - registry: Rc>>, + registry: Rc>>, } impl FunctionRegistry for EnvFunctionRegistry { - fn resolve(&self, addr: Address) -> Option { + fn resolve(&self, addr: Address) -> Option { 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>>, global_types: Rc>>, global_values: Rc>>, - function_registry: Rc>>, + function_registry: Rc>>, } + impl MacroEvaluator for RuntimeMacroEvaluator { fn evaluate(&self, node: &Node, bindings: &HashMap, Node>) -> Result { // 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 };