From e65402364dafc9619c73232963196f9b070371db Mon Sep 17 00:00:00 2001 From: Brummel Date: Sat, 21 Mar 2026 14:12:14 +0100 Subject: [PATCH] Refactor: Use new NodeKind and clean up Binder definitions The Binder and related types have been refactored to use the new `NodeKind` enum instead of the previous `BoundKind`. This commit updates all references to use the new structure, ensuring consistency across the compiler's AST representation. Key changes include: - Replacing `BoundKind` with `NodeKind` in the Binder's `bind` and `visit` methods. - Updating pattern matching and field access to reflect the new enum variants (e.g., `NodeKind::Def` instead of `BoundKind::Define`). - Adjusting identifier bindings to use the new `IdentifierBinding` enum. - Reflecting changes in `DefBinding`, `AssignBinding`, and `LambdaBinding` structures. - Ensuring all newly created nodes use `NodeKind` and the appropriate metadata. --- examples/repeat.myc | 2 - src/ast/compiler/analyzer.rs | 773 +++---- src/ast/compiler/binder.rs | 1617 +++++++------- src/ast/compiler/bound_nodes.rs | 34 +- src/ast/compiler/captures.rs | 251 +-- src/ast/compiler/dumper.rs | 479 ++-- src/ast/compiler/lambda_collector.rs | 210 +- src/ast/compiler/lowering.rs | 514 +++-- src/ast/compiler/macros.rs | 1691 +++++++------- src/ast/compiler/optimizer/engine.rs | 1552 +++++++------ src/ast/compiler/optimizer/folder.rs | 204 +- src/ast/compiler/optimizer/inliner.rs | 387 ++-- .../compiler/optimizer/substitution_map.rs | 447 ++-- src/ast/compiler/optimizer/utils.rs | 401 ++-- src/ast/compiler/specializer.rs | 552 ++--- src/ast/compiler/type_checker.rs | 1986 +++++++++-------- src/ast/environment.rs | 37 +- src/ast/nodes.rs | 199 +- src/ast/parser.rs | 1066 ++++----- src/ast/vm.rs | 189 +- 20 files changed, 6523 insertions(+), 6068 deletions(-) diff --git a/examples/repeat.myc b/examples/repeat.myc index c6c4f9e..aa143db 100644 --- a/examples/repeat.myc +++ b/examples/repeat.myc @@ -1,5 +1,3 @@ -;; Benchmark: 19.9us -;; Benchmark-Repeat: 110 (do (repeat n 10 (print n) ) ) diff --git a/src/ast/compiler/analyzer.rs b/src/ast/compiler/analyzer.rs index 80032b3..d5f6ef3 100644 --- a/src/ast/compiler/analyzer.rs +++ b/src/ast/compiler/analyzer.rs @@ -1,386 +1,387 @@ -use crate::ast::compiler::bound_nodes::{ - Address, AnalyzedNode, BoundKind, GlobalIdx, Node, NodeMetrics, TypedNode, TypedPhase, -}; -use crate::ast::types::Purity; -use std::collections::{HashMap, HashSet}; -use std::rc::Rc; - -pub struct Analyzer<'a> { - root_purity: &'a [Purity], - /// Stack of currently visiting lambdas to detect direct recursion. - lambda_stack: Vec, - /// Map of global index to its Lambda identity if known. - globals_to_lambdas: HashMap, - /// Set of identities that were found to be recursive. - recursive_identities: HashSet, -} - -impl<'a> Analyzer<'a> { - pub fn analyze( - node: &TypedNode, - root_purity: &'a [Purity], - ) -> AnalyzedNode { - let mut analyzer = Self { - root_purity, - lambda_stack: Vec::new(), - globals_to_lambdas: HashMap::new(), - recursive_identities: HashSet::new(), - }; - - // First pass: map globals to their lambda identities - analyzer.collect_globals(node); - - // Second pass: full analysis (decorating TypedNode into AnalyzedNode) - analyzer.visit(Rc::new(node.clone())) - } - - fn collect_globals(&mut self, node: &TypedNode) { - match &node.kind { - BoundKind::Define { - addr: Address::Global(global_index), - value, - .. - } => { - if let BoundKind::Lambda { .. } = &value.kind { - self.globals_to_lambdas - .insert(*global_index, value.identity.clone()); - } - self.collect_globals(value); - } - BoundKind::Block { exprs } => { - for e in exprs { - self.collect_globals(e); - } - } - _ => { - node.kind - .for_each_child(|child| self.collect_globals(child)); - } - } - } - - fn visit(&mut self, node_rc: Rc) -> AnalyzedNode { - let node = &*node_rc; - let mut is_recursive = false; - - let (new_kind, purity) = match &node.kind { - BoundKind::Constant(v) => (BoundKind::Constant(v.clone()), Purity::Pure), - BoundKind::Nop => (BoundKind::Nop, Purity::Pure), - - BoundKind::Get { addr, name } => { - let p = match addr { - Address::Global(idx) => { - self.root_purity.get(idx.0 as usize).cloned().unwrap_or(Purity::Pure) - } - _ => Purity::Pure, - }; - ( - BoundKind::Get { - addr: *addr, - name: name.clone(), - }, - p, - ) - } - - BoundKind::FieldAccessor(k) => (BoundKind::FieldAccessor(*k), Purity::Pure), - - BoundKind::GetField { rec, field } => { - let rec_m = self.visit(rec.clone()); - let p = rec_m.ty.purity; - ( - BoundKind::GetField { - rec: Rc::new(rec_m), - field: *field, - }, - p, - ) - } - - BoundKind::Set { addr, value } => { - let val_m = self.visit(value.clone()); - ( - BoundKind::Set { - addr: *addr, - value: Rc::new(val_m), - }, - Purity::Impure, - ) - } - - BoundKind::Define { - name, - addr, - kind, - value, - captured_by, - } => { - let val_m = self.visit(value.clone()); - ( - BoundKind::Define { - name: name.clone(), - addr: *addr, - kind: *kind, - value: Rc::new(val_m), - captured_by: captured_by.clone(), - }, - Purity::Impure, - ) - } - - BoundKind::If { - cond, - then_br, - else_br, - } => { - let cond_m = self.visit(cond.clone()); - let then_m = self.visit(then_br.clone()); - let else_m = else_br.as_ref().map(|e| self.visit(e.clone())); - - let mut p = cond_m.ty.purity.min(then_m.ty.purity); - if let Some(ref em) = else_m { - p = p.min(em.ty.purity); - } - ( - BoundKind::If { - cond: Rc::new(cond_m), - then_br: Rc::new(then_m), - else_br: else_m.map(Rc::new), - }, - p, - ) - } - - BoundKind::Lambda { - params, - upvalues, - body, - positional_count, - } => { - self.lambda_stack.push(node.identity.clone()); - let params_m = self.visit(params.clone()); - let body_m = self.visit(body.clone()); - self.lambda_stack.pop(); - - is_recursive = self.recursive_identities.contains(&node.identity); - ( - BoundKind::Lambda { - params: Rc::new(params_m), - upvalues: upvalues.clone(), - body: Rc::new(body_m), - positional_count: *positional_count, - }, - Purity::Pure, - ) - } - - BoundKind::Destructure { pattern, value } => { - let pat_m = self.visit(pattern.clone()); - let val_m = self.visit(value.clone()); - ( - BoundKind::Destructure { - pattern: Rc::new(pat_m), - value: Rc::new(val_m), - }, - Purity::Impure, - ) - } - - BoundKind::Call { callee, args } => { - let callee_m = self.visit(callee.clone()); - let args_m = self.visit(args.clone()); - - if let BoundKind::Get { - addr: Address::Global(idx), - .. - } = &callee.kind - && let Some(lambda_id) = self.globals_to_lambdas.get(idx) - && self.lambda_stack.contains(lambda_id) - { - self.recursive_identities.insert(lambda_id.clone()); - is_recursive = true; - } - - let p_func = if let BoundKind::Get { - addr: Address::Global(idx), - .. - } = &callee.kind - { - self.root_purity - .get(idx.0 as usize) - .cloned() - .unwrap_or(Purity::Impure) - } else { - Purity::Impure - }; - let p = callee_m.ty.purity.min(args_m.ty.purity).min(p_func); - ( - BoundKind::Call { - callee: Rc::new(callee_m), - args: Rc::new(args_m), - }, - p, - ) - } - - BoundKind::Again { args } => { - let args_m = self.visit(args.clone()); - if let Some(lambda_id) = self.lambda_stack.last() { - self.recursive_identities.insert(lambda_id.clone()); - is_recursive = true; - } - ( - BoundKind::Again { - args: Rc::new(args_m), - }, - Purity::Impure, - ) - } - BoundKind::Pipe { - inputs, - lambda, - out_type, - } => { - let mut analyzed_inputs = Vec::with_capacity(inputs.len()); - for input in inputs { - analyzed_inputs.push(Rc::new(self.visit(input.clone()))); - } - let a_lambda = Rc::new(self.visit(lambda.clone())); - ( - BoundKind::Pipe { - inputs: analyzed_inputs, - lambda: a_lambda, - out_type: out_type.clone(), - }, - Purity::Impure, - ) - } - - BoundKind::Block { exprs } => { - let mut new_exprs = Vec::with_capacity(exprs.len()); - let mut p = Purity::Pure; - for e in exprs { - let em = self.visit(e.clone()); - p = p.min(em.ty.purity); - new_exprs.push(Rc::new(em)); - } - (BoundKind::Block { exprs: new_exprs }, p) - } - - BoundKind::Tuple { elements } => { - let mut new_elements = Vec::with_capacity(elements.len()); - let mut p = Purity::Pure; - for e in elements { - let em = self.visit(e.clone()); - p = p.min(em.ty.purity); - new_elements.push(Rc::new(em)); - } - ( - BoundKind::Tuple { - elements: new_elements, - }, - p, - ) - } - - BoundKind::Record { layout, values } => { - let mut new_values = Vec::with_capacity(values.len()); - let mut p = Purity::Pure; - for v in values { - let vm = self.visit(v.clone()); - p = p.min(vm.ty.purity); - new_values.push(Rc::new(vm)); - } - ( - BoundKind::Record { - layout: layout.clone(), - values: new_values, - }, - p, - ) - } - - BoundKind::Expansion { - original_call, - bound_expanded, - } => { - let expanded_m = self.visit(bound_expanded.clone()); - ( - BoundKind::Expansion { - original_call: original_call.clone(), - bound_expanded: Rc::new(expanded_m.clone()), - }, - expanded_m.ty.purity, - ) - } - - BoundKind::Extension(_) => (BoundKind::Nop, Purity::Impure), - BoundKind::Error => (BoundKind::Error, Purity::Impure), - }; - - Node { - identity: node.identity.clone(), - kind: new_kind, - ty: NodeMetrics { - original: node_rc, - purity, - is_recursive, - }, - } - } -} - -trait NodeExt { - fn for_each_child(&self, f: F); -} - -impl NodeExt for BoundKind { - fn for_each_child(&self, mut f: F) { - match self { - BoundKind::If { - cond, - then_br, - else_br, - } => { - f(cond); - f(then_br); - if let Some(e) = else_br { - f(e); - } - } - BoundKind::Define { value, .. } | BoundKind::Set { value, .. } => { - f(value); - } - BoundKind::GetField { rec, .. } => { - f(rec); - } - BoundKind::Lambda { params, body, .. } => { - f(params); - f(body); - } - BoundKind::Call { callee, args } => { - f(callee); - f(args); - } - BoundKind::Block { exprs } => { - for e in exprs { - f(e); - } - } - BoundKind::Tuple { elements } => { - for e in elements { - f(e); - } - } - BoundKind::Record { values, .. } => { - for v in values { - f(v); - } - } - BoundKind::Expansion { bound_expanded, .. } => { - f(bound_expanded); - } - _ => {} - } - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AnalyzedNode, GlobalIdx, IdentifierBinding, Node, + NodeKind, NodeMetrics, TypedNode, TypedPhase, +}; +use crate::ast::types::Purity; +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; + +pub struct Analyzer<'a> { + root_purity: &'a [Purity], + /// Stack of currently visiting lambdas to detect direct recursion. + lambda_stack: Vec, + /// Map of global index to its Lambda identity if known. + globals_to_lambdas: HashMap, + /// Set of identities that were found to be recursive. + recursive_identities: HashSet, +} + +impl<'a> Analyzer<'a> { + pub fn analyze( + node: &TypedNode, + root_purity: &'a [Purity], + ) -> AnalyzedNode { + let mut analyzer = Self { + root_purity, + lambda_stack: Vec::new(), + globals_to_lambdas: HashMap::new(), + recursive_identities: HashSet::new(), + }; + + // First pass: map globals to their lambda identities + analyzer.collect_globals(node); + + // Second pass: full analysis (decorating TypedNode into AnalyzedNode) + analyzer.visit(Rc::new(node.clone())) + } + + fn collect_globals(&mut self, node: &TypedNode) { + match &node.kind { + NodeKind::Def { + pattern, + value, + .. + } => { + if let NodeKind::Identifier { + binding: IdentifierBinding::Declaration { addr: Address::Global(global_index), .. }, + .. + } = &pattern.kind + && let NodeKind::Lambda { .. } = &value.kind + { + self.globals_to_lambdas + .insert(*global_index, value.identity.clone()); + } + self.collect_globals(value); + } + NodeKind::Block { exprs } => { + for e in exprs { + self.collect_globals(e); + } + } + _ => { + node.kind + .for_each_child(|child| self.collect_globals(child)); + } + } + } + + fn visit(&mut self, node_rc: Rc) -> AnalyzedNode { + let node = &*node_rc; + let mut is_recursive = false; + + let (new_kind, purity) = match &node.kind { + NodeKind::Constant(v) => (NodeKind::Constant(v.clone()), Purity::Pure), + NodeKind::Nop => (NodeKind::Nop, Purity::Pure), + + NodeKind::Identifier { symbol, binding } => { + let p = if let IdentifierBinding::Reference(Address::Global(idx)) = binding { + self.root_purity.get(idx.0 as usize).cloned().unwrap_or(Purity::Pure) + } else { + Purity::Pure + }; + ( + NodeKind::Identifier { + symbol: symbol.clone(), + binding: binding.clone(), + }, + p, + ) + } + + NodeKind::FieldAccessor(k) => (NodeKind::FieldAccessor(*k), Purity::Pure), + + NodeKind::GetField { rec, field } => { + let rec_m = self.visit(rec.clone()); + let p = rec_m.ty.purity; + ( + NodeKind::GetField { + rec: Rc::new(rec_m), + field: *field, + }, + p, + ) + } + + NodeKind::Assign { target, value, info } => { + let target_m = self.visit(target.clone()); + let val_m = self.visit(value.clone()); + ( + NodeKind::Assign { + target: Rc::new(target_m), + value: Rc::new(val_m), + info: info.clone(), + }, + Purity::Impure, + ) + } + + NodeKind::Def { + pattern, + value, + info, + } => { + let pat_m = self.visit(pattern.clone()); + let val_m = self.visit(value.clone()); + ( + NodeKind::Def { + pattern: Rc::new(pat_m), + value: Rc::new(val_m), + info: info.clone(), + }, + Purity::Impure, + ) + } + + NodeKind::If { + cond, + then_br, + else_br, + } => { + let cond_m = self.visit(cond.clone()); + let then_m = self.visit(then_br.clone()); + let else_m = else_br.as_ref().map(|e| self.visit(e.clone())); + + let mut p = cond_m.ty.purity.min(then_m.ty.purity); + if let Some(ref em) = else_m { + p = p.min(em.ty.purity); + } + ( + NodeKind::If { + cond: Rc::new(cond_m), + then_br: Rc::new(then_m), + else_br: else_m.map(Rc::new), + }, + p, + ) + } + + NodeKind::Lambda { + params, + body, + info, + } => { + self.lambda_stack.push(node.identity.clone()); + let params_m = self.visit(params.clone()); + let body_m = self.visit(body.clone()); + self.lambda_stack.pop(); + + is_recursive = self.recursive_identities.contains(&node.identity); + ( + NodeKind::Lambda { + params: Rc::new(params_m), + body: Rc::new(body_m), + info: info.clone(), + }, + Purity::Pure, + ) + } + + NodeKind::Call { callee, args } => { + let callee_m = self.visit(callee.clone()); + let args_m = self.visit(args.clone()); + + if let NodeKind::Identifier { + binding: IdentifierBinding::Reference(Address::Global(idx)), + .. + } = &callee.kind + && let Some(lambda_id) = self.globals_to_lambdas.get(idx) + && self.lambda_stack.contains(lambda_id) + { + self.recursive_identities.insert(lambda_id.clone()); + is_recursive = true; + } + + let p_func = if let NodeKind::Identifier { + binding: IdentifierBinding::Reference(Address::Global(idx)), + .. + } = &callee.kind + { + self.root_purity + .get(idx.0 as usize) + .cloned() + .unwrap_or(Purity::Impure) + } else { + Purity::Impure + }; + let p = callee_m.ty.purity.min(args_m.ty.purity).min(p_func); + ( + NodeKind::Call { + callee: Rc::new(callee_m), + args: Rc::new(args_m), + }, + p, + ) + } + + NodeKind::Again { args } => { + let args_m = self.visit(args.clone()); + if let Some(lambda_id) = self.lambda_stack.last() { + self.recursive_identities.insert(lambda_id.clone()); + is_recursive = true; + } + ( + NodeKind::Again { + args: Rc::new(args_m), + }, + Purity::Impure, + ) + } + NodeKind::Pipe { + inputs, + lambda, + } => { + let mut analyzed_inputs = Vec::with_capacity(inputs.len()); + for input in inputs { + analyzed_inputs.push(Rc::new(self.visit(input.clone()))); + } + let a_lambda = Rc::new(self.visit(lambda.clone())); + ( + NodeKind::Pipe { + inputs: analyzed_inputs, + lambda: a_lambda, + }, + Purity::Impure, + ) + } + + NodeKind::Block { exprs } => { + let mut new_exprs = Vec::with_capacity(exprs.len()); + let mut p = Purity::Pure; + for e in exprs { + let em = self.visit(e.clone()); + p = p.min(em.ty.purity); + new_exprs.push(Rc::new(em)); + } + (NodeKind::Block { exprs: new_exprs }, p) + } + + NodeKind::Tuple { elements } => { + let mut new_elements = Vec::with_capacity(elements.len()); + let mut p = Purity::Pure; + for e in elements { + let em = self.visit(e.clone()); + p = p.min(em.ty.purity); + new_elements.push(Rc::new(em)); + } + ( + NodeKind::Tuple { + elements: new_elements, + }, + p, + ) + } + + NodeKind::Record { fields, layout } => { + let mut new_fields = Vec::with_capacity(fields.len()); + let mut p = Purity::Pure; + for (key_node, val_node) in fields { + let km = self.visit(key_node.clone()); + let vm = self.visit(val_node.clone()); + p = p.min(vm.ty.purity); + new_fields.push((Rc::new(km), Rc::new(vm))); + } + ( + NodeKind::Record { + fields: new_fields, + layout: layout.clone(), + }, + p, + ) + } + + NodeKind::Expansion { + original_call, + expanded, + } => { + let expanded_m = self.visit(expanded.clone()); + ( + NodeKind::Expansion { + original_call: original_call.clone(), + expanded: Rc::new(expanded_m.clone()), + }, + expanded_m.ty.purity, + ) + } + + NodeKind::Extension(_) => (NodeKind::Nop, Purity::Impure), + NodeKind::Error => (NodeKind::Error, Purity::Impure), + + // Syntax-only variants should not appear in typed phases + NodeKind::MacroDecl { .. } + | NodeKind::Template(_) + | NodeKind::Placeholder(_) + | NodeKind::Splice(_) => (NodeKind::Error, Purity::Impure), + }; + + Node { + identity: node.identity.clone(), + kind: new_kind, + ty: NodeMetrics { + original: node_rc, + purity, + is_recursive, + }, + } + } +} + +trait NodeExt { + fn for_each_child(&self, f: F); +} + +impl NodeExt for NodeKind { + fn for_each_child(&self, mut f: F) { + match self { + NodeKind::If { + cond, + then_br, + else_br, + } => { + f(cond); + f(then_br); + if let Some(e) = else_br { + f(e); + } + } + NodeKind::Def { pattern, value, .. } => { + f(pattern); + f(value); + } + NodeKind::Assign { target, value, .. } => { + f(target); + f(value); + } + NodeKind::GetField { rec, .. } => { + f(rec); + } + NodeKind::Lambda { params, body, .. } => { + f(params); + f(body); + } + NodeKind::Call { callee, args } => { + f(callee); + f(args); + } + NodeKind::Block { exprs } => { + for e in exprs { + f(e); + } + } + NodeKind::Tuple { elements } => { + for e in elements { + f(e); + } + } + NodeKind::Record { fields, .. } => { + for (key, val) in fields { + f(key); + f(val); + } + } + NodeKind::Expansion { expanded, .. } => { + f(expanded); + } + _ => {} + } + } +} diff --git a/src/ast/compiler/binder.rs b/src/ast/compiler/binder.rs index c384312..04c4853 100644 --- a/src/ast/compiler/binder.rs +++ b/src/ast/compiler/binder.rs @@ -1,789 +1,828 @@ -use crate::ast::compiler::bound_nodes::{ - Address, BoundKind, BoundPhase, DeclarationKind, GlobalIdx, Node, UpvalueIdx, VirtualId, -}; -use crate::ast::diagnostics::Diagnostics; -use crate::ast::nodes::{Symbol, SyntaxKind, SyntaxNode}; -use crate::ast::types::{Identity, StaticType, Purity}; -use std::collections::HashMap; -use std::rc::Rc; - -#[derive(Debug, Clone)] -pub struct LocalInfo { - pub addr: Address, - pub identity: Identity, - pub _ty: StaticType, - pub purity: Purity, -} - -#[derive(Debug, Clone)] -pub struct CompilerScope { - pub locals: HashMap, -} - -impl Default for CompilerScope { - fn default() -> Self { - Self::new() - } -} - -impl CompilerScope { - pub fn new() -> Self { - Self { - locals: HashMap::new(), - } - } -} - -struct FunctionCompiler { - identity: Identity, - scopes: Vec, - slot_count: u32, - upvalues: Vec>, -} - -impl FunctionCompiler { - fn new(identity: Identity, initial_scopes: Vec, initial_slot_count: u32) -> Self { - Self { - identity, - scopes: if initial_scopes.is_empty() { - vec![CompilerScope::new()] - } else { - initial_scopes - }, - slot_count: initial_slot_count, - upvalues: Vec::new(), - } - } - - fn push_scope(&mut self) { - self.scopes.push(CompilerScope::new()); - } - - fn pop_scope(&mut self) { - if self.scopes.len() > 1 { - self.scopes.pop(); - } - } - - fn define_variable(&mut self, name: &Symbol, identity: Identity) -> Result, String> { - let current_scope = self.scopes.last_mut().unwrap(); - if current_scope.locals.contains_key(name) { - return Err(format!( - "Variable '{}' is already defined in this scope level.", - name.name - )); - } - - let slot = VirtualId(self.slot_count); - current_scope.locals.insert( - name.clone(), - LocalInfo { - addr: Address::Local(slot), - identity, - _ty: StaticType::Any, - purity: Purity::Impure, - }, - ); - self.slot_count += 1; - Ok(Address::Local(slot)) - } - - fn resolve_local(&self, sym: &Symbol) -> Option<(LocalInfo, usize)> { - for (idx, scope) in self.scopes.iter().enumerate().rev() { - if let Some(info) = scope.locals.get(sym) { - return Some((info.clone(), idx)); - } - } - None - } - - fn add_upvalue(&mut self, addr: Address) -> UpvalueIdx { - if let Some(idx) = self.upvalues.iter().position(|&a| a == addr) { - return UpvalueIdx(idx as u32); - } - let idx = UpvalueIdx(self.upvalues.len() as u32); - self.upvalues.push(addr); - idx - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ExprContext { - Expression, - Statement, -} - -pub type BindingResult = ( - Node, - HashMap>, - Vec, - u32, -); - -pub struct Binder { - functions: Vec, - capture_map: HashMap>, - fixed_scope_idx: i32, -} - -impl Binder { - pub fn new( - initial_scopes: Vec, - initial_slot_count: u32, - fixed_scope_idx: i32, - ) -> Self { - let mut binder = Self { - functions: Vec::new(), - capture_map: HashMap::new(), - fixed_scope_idx, - }; - binder.functions.push(FunctionCompiler::new( - crate::ast::types::NodeIdentity::new(crate::ast::types::SourceLocation { - line: 0, - col: 0, - }), - initial_scopes, - initial_slot_count, - )); - binder - } - - pub fn bind_root( - initial_scopes: Vec, - initial_slot_count: u32, - fixed_scope_idx: i32, - node: &SyntaxNode, - diagnostics: &mut Diagnostics, - ) -> Result { - let mut binder = Self::new(initial_scopes, initial_slot_count, fixed_scope_idx); - let bound = binder.bind(node, ExprContext::Expression, diagnostics); - - let final_captures = binder - .capture_map - .into_iter() - .map(|(k, v)| (k, v.into_iter().collect())) - .collect(); - - let root_compiler = binder.functions.pop().unwrap(); - Ok(( - bound, - final_captures, - root_compiler.scopes, - root_compiler.slot_count, - )) - } - - fn declare_variable( - &mut self, - name: &Symbol, - identity: Identity, - _kind: crate::ast::compiler::bound_nodes::DeclarationKind, - diag: &mut Diagnostics, - ) -> Option> { - let current_fn_idx = self.functions.len() - 1; - - if current_fn_idx == 0 { - let current_scope_idx = self.functions[0].scopes.len() as i32 - 1; - if current_scope_idx <= self.fixed_scope_idx { - diag.push_error( - format!( - "Cannot define '{}': Scope {} is frozen/immutable.", - name.name, current_scope_idx - ), - None, - ); - return None; - } - } - - let current_fn = self.functions.last_mut().unwrap(); - match current_fn.define_variable(name, identity) { - Ok(addr) => Some(addr), - Err(e) => { - diag.push_error(e, None); - None - } - } - } - - pub fn bind( - &mut self, - node: &SyntaxNode, - ctx: ExprContext, - diag: &mut Diagnostics, - ) -> Node { - match &node.kind { - SyntaxKind::Nop => self.make_node(node.identity.clone(), BoundKind::Nop), - SyntaxKind::Constant(v) => { - self.make_node(node.identity.clone(), BoundKind::Constant(v.clone())) - } - - SyntaxKind::Identifier(sym) => { - if let Some(addr) = self.resolve_variable(sym, diag, &node.identity) { - self.make_node( - node.identity.clone(), - BoundKind::Get { - addr, - name: sym.clone(), - }, - ) - } else { - self.make_node(node.identity.clone(), BoundKind::Error) - } - } - - SyntaxKind::FieldAccessor(k) => { - self.make_node(node.identity.clone(), BoundKind::FieldAccessor(*k)) - } - - SyntaxKind::If { - cond, - then_br, - else_br, - } => { - let cond = self.bind(cond.as_ref(), ExprContext::Expression, diag); - - self.functions.last_mut().unwrap().push_scope(); - let then_br = self.bind(then_br.as_ref(), ctx, diag); - self.functions.last_mut().unwrap().pop_scope(); - - let mut else_br_bound = None; - if let Some(e) = else_br { - self.functions.last_mut().unwrap().push_scope(); - else_br_bound = Some(Rc::new(self.bind(e, ctx, diag))); - self.functions.last_mut().unwrap().pop_scope(); - } - - self.make_node( - node.identity.clone(), - BoundKind::If { - cond: Rc::new(cond), - then_br: Rc::new(then_br), - else_br: else_br_bound, - }, - ) - } - - SyntaxKind::Def { target, value } => { - if ctx == ExprContext::Expression { - diag.push_error( - "Statement 'def' cannot be used as an expression.", - Some(node.identity.clone()), - ); - } - if let SyntaxKind::Identifier(ref name) = target.kind { - let addr_opt = self.declare_variable( - name, - node.identity.clone(), - crate::ast::compiler::bound_nodes::DeclarationKind::Variable, - diag, - ); - let val_node = self.bind(value, ExprContext::Expression, diag); - - if let Some(addr) = addr_opt { - self.make_node( - node.identity.clone(), - BoundKind::Define { - name: name.clone(), - addr, - kind: crate::ast::compiler::bound_nodes::DeclarationKind::Variable, - value: Rc::new(val_node), - captured_by: Vec::new(), - }, - ) - } else { - self.make_node(node.identity.clone(), BoundKind::Error) - } - } else { - let val_node = self.bind(value, ExprContext::Expression, diag); - let target_node = self.bind_pattern(target.as_ref(), DeclarationKind::Variable, diag); - - self.make_node( - node.identity.clone(), - BoundKind::Destructure { - pattern: Rc::new(target_node), - value: Rc::new(val_node), - }, - ) - } - } - - SyntaxKind::Assign { target, value } => { - let val_node = self.bind(value, ExprContext::Expression, diag); - - if let SyntaxKind::Identifier(sym) = &target.kind { - if let Some(addr) = self.resolve_variable(sym, diag, &target.identity) { - self.make_node( - node.identity.clone(), - BoundKind::Set { - addr, - value: Rc::new(val_node), - }, - ) - } else { - self.make_node(node.identity.clone(), BoundKind::Error) - } - } else { - let target_node = self.bind_assign_pattern(target.as_ref(), diag); - self.make_node( - node.identity.clone(), - BoundKind::Destructure { - pattern: Rc::new(target_node), - value: Rc::new(val_node), - }, - ) - } - } - - SyntaxKind::Pipe { inputs, lambda } => { - let mut bound_inputs = Vec::with_capacity(inputs.len()); - for input in inputs { - bound_inputs.push(Rc::new(self.bind(input, ExprContext::Expression, diag))); - } - let bound_lambda = - Rc::new(self.bind(lambda.as_ref(), ExprContext::Expression, diag)); - self.make_node( - node.identity.clone(), - BoundKind::Pipe { - inputs: bound_inputs, - lambda: bound_lambda, - out_type: crate::ast::types::StaticType::Any, - }, - ) - } - - SyntaxKind::Lambda { params, body } => { - let identity = node.identity.clone(); - self.functions - .push(FunctionCompiler::new(identity.clone(), vec![], 0)); - - let params_bound = self.bind_pattern(params.as_ref(), DeclarationKind::Parameter, diag); - let body_bound = self.bind(body.as_ref(), ExprContext::Expression, diag); - let compiled_fn = self.functions.pop().unwrap(); - - fn count_params(node: &Node) -> Option { - match &node.kind { - BoundKind::Define { - kind: DeclarationKind::Parameter, - .. - } => Some(1), - BoundKind::Tuple { elements } => { - let mut total = 0; - for e in elements { - total += count_params(e)?; - } - Some(total) - } - BoundKind::Nop => Some(0), - _ => None, - } - } - - let positional_count = count_params(¶ms_bound); - - self.make_node( - identity, - BoundKind::Lambda { - params: Rc::new(params_bound), - upvalues: compiled_fn.upvalues, - body: Rc::new(body_bound), - positional_count, - }, - ) - } - - SyntaxKind::Call { callee, args } => { - let callee = self.bind(callee, ExprContext::Expression, diag); - let args = self.bind(args.as_ref(), ExprContext::Expression, diag); - - self.make_node( - node.identity.clone(), - BoundKind::Call { - callee: Rc::new(callee), - args: Rc::new(args), - }, - ) - } - - SyntaxKind::Again { args } => { - if self.functions.len() <= 1 { - diag.push_error( - "'again' is only allowed inside a function or lambda.", - Some(node.identity.clone()), - ); - return self.make_node(node.identity.clone(), BoundKind::Error); - } - let args = self.bind(args.as_ref(), ExprContext::Expression, diag); - self.make_node( - node.identity.clone(), - BoundKind::Again { - args: Rc::new(args), - }, - ) - } - - SyntaxKind::Block { exprs } => { - self.functions.last_mut().unwrap().push_scope(); - let mut bound_exprs = Vec::new(); - for (i, expr) in exprs.iter().enumerate() { - let expr_ctx = if i == exprs.len() - 1 { - ctx - } else { - ExprContext::Statement - }; - bound_exprs.push(Rc::new(self.bind(expr, expr_ctx, diag))); - } - self.functions.last_mut().unwrap().pop_scope(); - self.make_node( - node.identity.clone(), - BoundKind::Block { exprs: bound_exprs }, - ) - } - - SyntaxKind::Tuple { elements } => { - let mut bound_elems = Vec::new(); - for e in elements { - bound_elems.push(Rc::new(self.bind(e, ExprContext::Expression, diag))); - } - self.make_node( - node.identity.clone(), - BoundKind::Tuple { - elements: bound_elems, - }, - ) - } - - SyntaxKind::Record { fields } => { - let mut bound_values = Vec::new(); - let mut layout_fields = Vec::new(); - - for (k, v) in fields { - let key_node = self.bind(k, ExprContext::Expression, diag); - let val_node = self.bind(v, ExprContext::Expression, diag); - - if let BoundKind::Constant(crate::ast::types::Value::Keyword(kw)) = - key_node.kind - { - layout_fields.push((kw, crate::ast::types::StaticType::Any)); - } else { - diag.push_error( - format!( - "Record keys must be keywords, found at {:?}", - key_node.identity.location - ), - Some(key_node.identity.clone()), - ); - } - bound_values.push(Rc::new(val_node)); - } - - let layout = crate::ast::types::RecordLayout::get_or_create(layout_fields); - - self.make_node( - node.identity.clone(), - BoundKind::Record { - layout, - values: bound_values, - }, - ) - } - - SyntaxKind::Expansion { call, expanded } => { - let bound_expanded = self.bind(expanded.as_ref(), ctx, diag); - self.make_node( - node.identity.clone(), - BoundKind::Expansion { - original_call: Rc::from(call.as_ref().clone()), - bound_expanded: Rc::new(bound_expanded), - }, - ) - } - - SyntaxKind::Template(_) - | SyntaxKind::Placeholder(_) - | SyntaxKind::Splice(_) - | SyntaxKind::MacroDecl { .. } => { - diag.push_error( - format!( - "Macro construct {:?} found in Binder. Macros must be expanded before binding.", - node.kind - ), - Some(node.identity.clone()), - ); - self.make_node(node.identity.clone(), BoundKind::Error) - } - - SyntaxKind::Extension(_) => { - diag.push_error( - "Custom extensions not supported in Binder yet", - Some(node.identity.clone()), - ); - self.make_node(node.identity.clone(), BoundKind::Error) - } - SyntaxKind::Error => Node { - identity: node.identity.clone(), - kind: BoundKind::Error, - ty: (), - }, - } - } - - fn resolve_variable( - &mut self, - sym: &Symbol, - diag: &mut Diagnostics, - identity: &Identity, - ) -> Option> { - let current_fn_idx = self.functions.len() - 1; - - // 1. Try local in current function - if let Some((info, _)) = self.functions[current_fn_idx].resolve_local(sym) { - return Some(info.addr); - } - - // 2. Try enclosing scopes (capture chain) - for i in (0..current_fn_idx).rev() { - if let Some((info, scope_idx)) = self.functions[i].resolve_local(sym) { - // If it's in the root script and within a frozen scope, translate to Global - let is_frozen_root = i == 0 && (scope_idx as i32) <= self.fixed_scope_idx; - - if let Address::Global(_) = info.addr { - return Some(info.addr); - } - - let mut addr = info.addr; - - if is_frozen_root - && let Address::Local(slot) = addr { - addr = Address::Global(GlobalIdx(slot.0)); - } - - if let Address::Global(_) = addr { - return Some(addr); - } - - // Record the capture for each lambda level in between - for k in (i + 1)..=current_fn_idx { - let lambda_id = self.functions[k].identity.clone(); - self.capture_map - .entry(info.identity.clone()) - .or_default() - .insert(lambda_id); - addr = Address::Upvalue(self.functions[k].add_upvalue(addr)); - } - return Some(addr); - } - } - - // 3. Global Fallback (search in root level with context removed) - if sym.context.is_some() { - let fallback_sym = Symbol { - name: sym.name.clone(), - context: None, - }; - // Search again in all functions, primarily we care about Root Scopes - for i in (0..=current_fn_idx).rev() { - if let Some((info, _)) = self.functions[i].resolve_local(&fallback_sym) - && let Address::Global(_) = info.addr - { - return Some(info.addr); - } - } - } - - diag.push_error( - format!("Undefined variable '{}'", sym.name), - Some(identity.clone()), - ); - None - } - - fn bind_pattern( - &mut self, - node: &SyntaxNode, - kind: DeclarationKind, - diag: &mut Diagnostics, - ) -> Node { - match &node.kind { - SyntaxKind::Identifier(sym) => { - if let Some(addr) = self.declare_variable(sym, node.identity.clone(), kind, diag) { - self.make_node( - node.identity.clone(), - BoundKind::Define { - name: sym.clone(), - addr, - kind, - value: Rc::new(self.make_node(node.identity.clone(), BoundKind::Nop)), - captured_by: Vec::new(), - }, - ) - } else { - self.make_node(node.identity.clone(), BoundKind::Error) - } - } - SyntaxKind::Tuple { elements } => { - let mut bound_elems = Vec::new(); - for e in elements { - bound_elems.push(Rc::new(self.bind_pattern(e, kind, diag))); - } - self.make_node( - node.identity.clone(), - BoundKind::Tuple { - elements: bound_elems, - }, - ) - } - _ => { - diag.push_error( - format!("Invalid node in pattern: {:?}", node.kind), - Some(node.identity.clone()), - ); - self.make_node(node.identity.clone(), BoundKind::Error) - } - } - } - - fn bind_assign_pattern( - &mut self, - node: &SyntaxNode, - diag: &mut Diagnostics, - ) -> Node { - match &node.kind { - SyntaxKind::Identifier(sym) => { - if let Some(addr) = self.resolve_variable(sym, diag, &node.identity) { - self.make_node( - node.identity.clone(), - BoundKind::Set { - addr, - value: Rc::new(self.make_node(node.identity.clone(), BoundKind::Nop)), - }, - ) - } else { - self.make_node(node.identity.clone(), BoundKind::Error) - } - } - SyntaxKind::Tuple { elements } => { - let mut bound_elems = Vec::new(); - for e in elements { - bound_elems.push(Rc::new(self.bind_assign_pattern(e, diag))); - } - self.make_node( - node.identity.clone(), - BoundKind::Tuple { - elements: bound_elems, - }, - ) - } - _ => { - diag.push_error( - format!("Invalid node in assignment pattern: {:?}", node.kind), - Some(node.identity.clone()), - ); - self.make_node(node.identity.clone(), BoundKind::Error) - } - } - } - - fn make_node(&self, identity: Identity, kind: BoundKind) -> Node { - Node { - identity, - kind, - ty: (), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::ast::diagnostics::Diagnostics; - use crate::ast::parser::Parser; - - #[test] - fn test_upvalue_capture_sets_is_boxed() { - let source = "(fn [] (do (def x 10) (def f (fn [] x)) x))"; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut diagnostics = Diagnostics::new(); - let (bound, captures, _scopes, _slots) = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics).unwrap(); - - if let BoundKind::Lambda { body, .. } = &bound.kind { - if let BoundKind::Block { exprs } = &body.kind { - let x_decl = &exprs[0]; - if let BoundKind::Define { addr, .. } = &x_decl.kind { - assert!(matches!(addr, Address::Local(_))); - assert!( - captures.contains_key(&x_decl.identity), - "Variable 'x' should have capturers" - ); - } else { - panic!("First expression should be Define"); - } - } else { - panic!("Lambda body should be a Block"); - } - } else { - panic!("Root should be a Lambda"); - } - } - - #[test] - fn test_no_capture_not_boxed() { - let source = "(fn [] (do (def x 10) x))"; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut diagnostics = Diagnostics::new(); - let (bound, captures, _scopes, _slots) = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics).unwrap(); - - if let BoundKind::Lambda { body, .. } = &bound.kind { - if let BoundKind::Block { exprs } = &body.kind { - let x_decl = &exprs[0]; - if let BoundKind::Define { addr, .. } = &x_decl.kind { - assert!(matches!(addr, Address::Local(_))); - assert!( - !captures.contains_key(&x_decl.identity), - "Variable 'x' should NOT have any capturers" - ); - } else { - panic!("First expression should be Define"); - } - } else { - panic!("Lambda body should be a Block"); - } - } else { - panic!("Root should be a Lambda"); - } - } - - #[test] - fn test_redefinition_error() { - let source = "(do (def x 1) (def x 2) x)"; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut diagnostics = Diagnostics::new(); - let _ = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics); - - assert!(diagnostics.has_errors()); - assert!(diagnostics.items.iter().any(|i| i.message.contains("already defined"))); - } - - #[test] - fn test_repro_global_redefinition() { - let source1 = "(def x 1) 1"; - let syntax1 = Parser::new(source1).parse_expression(); - let mut diagnostics = Diagnostics::new(); - let (_, _, scopes1, slots1) = Binder::bind_root(vec![], 0, -1, &syntax1, &mut diagnostics).unwrap(); - - let source2 = "(def x 2) 2"; - let syntax2 = Parser::new(source2).parse_expression(); - let mut diagnostics2 = Diagnostics::new(); - // Here we simulate frozen scope by passing fixed_scope_idx = 0 - let _ = Binder::bind_root(scopes1, slots1, 0, &syntax2, &mut diagnostics2); - - assert!(diagnostics2.has_errors()); - assert!(diagnostics2.items.iter().any(|i| i.message.contains("frozen/immutable"))); - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AssignBinding, BoundPhase, DeclarationKind, DefBinding, GlobalIdx, + IdentifierBinding, LambdaBinding, Node, NodeKind, UpvalueIdx, VirtualId, +}; +use crate::ast::diagnostics::Diagnostics; +use crate::ast::nodes::{SyntaxKind, SyntaxNode, Symbol}; +use crate::ast::types::{Identity, Purity, StaticType}; +use std::collections::HashMap; +use std::rc::Rc; + +#[derive(Debug, Clone)] +pub struct LocalInfo { + pub addr: Address, + pub identity: Identity, + pub _ty: StaticType, + pub purity: Purity, +} + +#[derive(Debug, Clone)] +pub struct CompilerScope { + pub locals: HashMap, +} + +impl Default for CompilerScope { + fn default() -> Self { + Self::new() + } +} + +impl CompilerScope { + pub fn new() -> Self { + Self { + locals: HashMap::new(), + } + } +} + +struct FunctionCompiler { + identity: Identity, + scopes: Vec, + slot_count: u32, + upvalues: Vec>, +} + +impl FunctionCompiler { + fn new(identity: Identity, initial_scopes: Vec, initial_slot_count: u32) -> Self { + Self { + identity, + scopes: if initial_scopes.is_empty() { + vec![CompilerScope::new()] + } else { + initial_scopes + }, + slot_count: initial_slot_count, + upvalues: Vec::new(), + } + } + + fn push_scope(&mut self) { + self.scopes.push(CompilerScope::new()); + } + + fn pop_scope(&mut self) { + if self.scopes.len() > 1 { + self.scopes.pop(); + } + } + + fn define_variable(&mut self, name: &Symbol, identity: Identity) -> Result, String> { + let current_scope = self.scopes.last_mut().unwrap(); + if current_scope.locals.contains_key(name) { + return Err(format!( + "Variable '{}' is already defined in this scope level.", + name.name + )); + } + + let slot = VirtualId(self.slot_count); + current_scope.locals.insert( + name.clone(), + LocalInfo { + addr: Address::Local(slot), + identity, + _ty: StaticType::Any, + purity: Purity::Impure, + }, + ); + self.slot_count += 1; + Ok(Address::Local(slot)) + } + + fn resolve_local(&self, sym: &Symbol) -> Option<(LocalInfo, usize)> { + for (idx, scope) in self.scopes.iter().enumerate().rev() { + if let Some(info) = scope.locals.get(sym) { + return Some((info.clone(), idx)); + } + } + None + } + + fn add_upvalue(&mut self, addr: Address) -> UpvalueIdx { + if let Some(idx) = self.upvalues.iter().position(|&a| a == addr) { + return UpvalueIdx(idx as u32); + } + let idx = UpvalueIdx(self.upvalues.len() as u32); + self.upvalues.push(addr); + idx + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ExprContext { + Expression, + Statement, +} + +pub type BindingResult = ( + Node, + HashMap>, + Vec, + u32, +); + +pub struct Binder { + functions: Vec, + capture_map: HashMap>, + fixed_scope_idx: i32, +} + +impl Binder { + pub fn new( + initial_scopes: Vec, + initial_slot_count: u32, + fixed_scope_idx: i32, + ) -> Self { + let mut binder = Self { + functions: Vec::new(), + capture_map: HashMap::new(), + fixed_scope_idx, + }; + binder.functions.push(FunctionCompiler::new( + crate::ast::types::NodeIdentity::new(crate::ast::types::SourceLocation { + line: 0, + col: 0, + }), + initial_scopes, + initial_slot_count, + )); + binder + } + + pub fn bind_root( + initial_scopes: Vec, + initial_slot_count: u32, + fixed_scope_idx: i32, + node: &SyntaxNode, + diagnostics: &mut Diagnostics, + ) -> Result { + let mut binder = Self::new(initial_scopes, initial_slot_count, fixed_scope_idx); + let bound = binder.bind(node, ExprContext::Expression, diagnostics); + + let final_captures = binder + .capture_map + .into_iter() + .map(|(k, v)| (k, v.into_iter().collect())) + .collect(); + + let root_compiler = binder.functions.pop().unwrap(); + Ok(( + bound, + final_captures, + root_compiler.scopes, + root_compiler.slot_count, + )) + } + + fn declare_variable( + &mut self, + name: &Symbol, + identity: Identity, + _kind: DeclarationKind, + diag: &mut Diagnostics, + ) -> Option> { + let current_fn_idx = self.functions.len() - 1; + + if current_fn_idx == 0 { + let current_scope_idx = self.functions[0].scopes.len() as i32 - 1; + if current_scope_idx <= self.fixed_scope_idx { + diag.push_error( + format!( + "Cannot define '{}': Scope {} is frozen/immutable.", + name.name, current_scope_idx + ), + None, + ); + return None; + } + } + + let current_fn = self.functions.last_mut().unwrap(); + match current_fn.define_variable(name, identity) { + Ok(addr) => Some(addr), + Err(e) => { + diag.push_error(e, None); + None + } + } + } + + pub fn bind( + &mut self, + node: &SyntaxNode, + ctx: ExprContext, + diag: &mut Diagnostics, + ) -> Node { + match &node.kind { + SyntaxKind::Nop => self.make_node(node.identity.clone(), NodeKind::Nop), + SyntaxKind::Constant(v) => { + self.make_node(node.identity.clone(), NodeKind::Constant(v.clone())) + } + + SyntaxKind::Identifier { symbol: sym, .. } => { + if let Some(addr) = self.resolve_variable(sym, diag, &node.identity) { + self.make_node( + node.identity.clone(), + NodeKind::Identifier { + symbol: sym.clone(), + binding: IdentifierBinding::Reference(addr), + }, + ) + } else { + self.make_node(node.identity.clone(), NodeKind::Error) + } + } + + SyntaxKind::FieldAccessor(k) => { + self.make_node(node.identity.clone(), NodeKind::FieldAccessor(*k)) + } + + SyntaxKind::If { + cond, + then_br, + else_br, + } => { + let cond = self.bind(cond.as_ref(), ExprContext::Expression, diag); + + self.functions.last_mut().unwrap().push_scope(); + let then_br = self.bind(then_br.as_ref(), ctx, diag); + self.functions.last_mut().unwrap().pop_scope(); + + let mut else_br_bound = None; + if let Some(e) = else_br { + self.functions.last_mut().unwrap().push_scope(); + else_br_bound = Some(Rc::new(self.bind(e, ctx, diag))); + self.functions.last_mut().unwrap().pop_scope(); + } + + self.make_node( + node.identity.clone(), + NodeKind::If { + cond: Rc::new(cond), + then_br: Rc::new(then_br), + else_br: else_br_bound, + }, + ) + } + + SyntaxKind::Def { pattern: target, value, .. } => { + if ctx == ExprContext::Expression { + diag.push_error( + "Statement 'def' cannot be used as an expression.", + Some(node.identity.clone()), + ); + } + if let SyntaxKind::Identifier { symbol: ref name, .. } = target.kind { + let addr_opt = self.declare_variable( + name, + node.identity.clone(), + DeclarationKind::Variable, + diag, + ); + let val_node = self.bind(value, ExprContext::Expression, diag); + + if let Some(addr) = addr_opt { + let pattern_node = self.make_node( + target.identity.clone(), + NodeKind::Identifier { + symbol: name.clone(), + binding: IdentifierBinding::Declaration { + addr, + kind: DeclarationKind::Variable, + }, + }, + ); + self.make_node( + node.identity.clone(), + NodeKind::Def { + pattern: Rc::new(pattern_node), + value: Rc::new(val_node), + info: DefBinding { + captured_by: Vec::new(), + }, + }, + ) + } else { + self.make_node(node.identity.clone(), NodeKind::Error) + } + } else { + let val_node = self.bind(value, ExprContext::Expression, diag); + let target_node = self.bind_pattern(target.as_ref(), DeclarationKind::Variable, diag); + + self.make_node( + node.identity.clone(), + NodeKind::Def { + pattern: Rc::new(target_node), + value: Rc::new(val_node), + info: DefBinding { + captured_by: Vec::new(), + }, + }, + ) + } + } + + SyntaxKind::Assign { target, value, .. } => { + let val_node = self.bind(value, ExprContext::Expression, diag); + + if let SyntaxKind::Identifier { symbol: sym, .. } = &target.kind { + if let Some(addr) = self.resolve_variable(sym, diag, &target.identity) { + let target_node = self.make_node( + target.identity.clone(), + NodeKind::Identifier { + symbol: sym.clone(), + binding: IdentifierBinding::Reference(addr), + }, + ); + self.make_node( + node.identity.clone(), + NodeKind::Assign { + target: Rc::new(target_node), + value: Rc::new(val_node), + info: AssignBinding { addr: Some(addr) }, + }, + ) + } else { + self.make_node(node.identity.clone(), NodeKind::Error) + } + } else { + let target_node = self.bind_assign_pattern(target.as_ref(), diag); + self.make_node( + node.identity.clone(), + NodeKind::Assign { + target: Rc::new(target_node), + value: Rc::new(val_node), + info: AssignBinding { addr: None }, + }, + ) + } + } + + SyntaxKind::Pipe { inputs, lambda } => { + let mut bound_inputs = Vec::with_capacity(inputs.len()); + for input in inputs { + bound_inputs.push(Rc::new(self.bind(input, ExprContext::Expression, diag))); + } + let bound_lambda = + Rc::new(self.bind(lambda.as_ref(), ExprContext::Expression, diag)); + self.make_node( + node.identity.clone(), + NodeKind::Pipe { + inputs: bound_inputs, + lambda: bound_lambda, + }, + ) + } + + SyntaxKind::Lambda { params, body, .. } => { + let identity = node.identity.clone(); + self.functions + .push(FunctionCompiler::new(identity.clone(), vec![], 0)); + + let params_bound = self.bind_pattern(params.as_ref(), DeclarationKind::Parameter, diag); + let body_bound = self.bind(body.as_ref(), ExprContext::Expression, diag); + let compiled_fn = self.functions.pop().unwrap(); + + fn count_params(node: &Node) -> Option { + match &node.kind { + NodeKind::Identifier { + binding: IdentifierBinding::Declaration { + kind: DeclarationKind::Parameter, + .. + }, + .. + } => Some(1), + NodeKind::Tuple { elements } => { + let mut total = 0; + for e in elements { + total += count_params(e)?; + } + Some(total) + } + NodeKind::Nop => Some(0), + _ => None, + } + } + + let positional_count = count_params(¶ms_bound); + + self.make_node( + identity, + NodeKind::Lambda { + params: Rc::new(params_bound), + body: Rc::new(body_bound), + info: LambdaBinding { + upvalues: compiled_fn.upvalues, + positional_count, + }, + }, + ) + } + + SyntaxKind::Call { callee, args } => { + let callee = self.bind(callee, ExprContext::Expression, diag); + let args = self.bind(args.as_ref(), ExprContext::Expression, diag); + + self.make_node( + node.identity.clone(), + NodeKind::Call { + callee: Rc::new(callee), + args: Rc::new(args), + }, + ) + } + + SyntaxKind::Again { args } => { + if self.functions.len() <= 1 { + diag.push_error( + "'again' is only allowed inside a function or lambda.", + Some(node.identity.clone()), + ); + return self.make_node(node.identity.clone(), NodeKind::Error); + } + let args = self.bind(args.as_ref(), ExprContext::Expression, diag); + self.make_node( + node.identity.clone(), + NodeKind::Again { + args: Rc::new(args), + }, + ) + } + + SyntaxKind::Block { exprs } => { + self.functions.last_mut().unwrap().push_scope(); + let mut bound_exprs = Vec::new(); + for (i, expr) in exprs.iter().enumerate() { + let expr_ctx = if i == exprs.len() - 1 { + ctx + } else { + ExprContext::Statement + }; + bound_exprs.push(Rc::new(self.bind(expr, expr_ctx, diag))); + } + self.functions.last_mut().unwrap().pop_scope(); + self.make_node( + node.identity.clone(), + NodeKind::Block { exprs: bound_exprs }, + ) + } + + SyntaxKind::Tuple { elements } => { + let mut bound_elems = Vec::new(); + for e in elements { + bound_elems.push(Rc::new(self.bind(e, ExprContext::Expression, diag))); + } + self.make_node( + node.identity.clone(), + NodeKind::Tuple { + elements: bound_elems, + }, + ) + } + + SyntaxKind::Record { fields, .. } => { + let mut bound_fields = Vec::new(); + let mut layout_fields = Vec::new(); + + for (k, v) in fields { + let key_node = self.bind(k, ExprContext::Expression, diag); + let val_node = self.bind(v, ExprContext::Expression, diag); + + if let NodeKind::Constant(crate::ast::types::Value::Keyword(kw)) = + &key_node.kind + { + layout_fields.push((*kw, crate::ast::types::StaticType::Any)); + } else { + diag.push_error( + format!( + "Record keys must be keywords, found at {:?}", + key_node.identity.location + ), + Some(key_node.identity.clone()), + ); + } + bound_fields.push((Rc::new(key_node), Rc::new(val_node))); + } + + let layout = crate::ast::types::RecordLayout::get_or_create(layout_fields); + + self.make_node( + node.identity.clone(), + NodeKind::Record { + fields: bound_fields, + layout, + }, + ) + } + + SyntaxKind::Expansion { original_call, expanded } => { + let bound_expanded = self.bind(expanded.as_ref(), ctx, diag); + self.make_node( + node.identity.clone(), + NodeKind::Expansion { + original_call: original_call.clone(), + expanded: Rc::new(bound_expanded), + }, + ) + } + + SyntaxKind::Template(_) + | SyntaxKind::Placeholder(_) + | SyntaxKind::Splice(_) + | SyntaxKind::MacroDecl { .. } => { + diag.push_error( + format!( + "Macro construct {:?} found in Binder. Macros must be expanded before binding.", + node.kind + ), + Some(node.identity.clone()), + ); + self.make_node(node.identity.clone(), NodeKind::Error) + } + + SyntaxKind::Extension(_) => { + diag.push_error( + "Custom extensions not supported in Binder yet", + Some(node.identity.clone()), + ); + self.make_node(node.identity.clone(), NodeKind::Error) + } + SyntaxKind::GetField { .. } => { + diag.push_error( + "GetField not expected in SyntaxPhase", + Some(node.identity.clone()), + ); + self.make_node(node.identity.clone(), NodeKind::Error) + } + SyntaxKind::Error => Node { + identity: node.identity.clone(), + kind: NodeKind::Error, + ty: (), + }, + } + } + + fn resolve_variable( + &mut self, + sym: &Symbol, + diag: &mut Diagnostics, + identity: &Identity, + ) -> Option> { + let current_fn_idx = self.functions.len() - 1; + + // 1. Try local in current function + if let Some((info, _)) = self.functions[current_fn_idx].resolve_local(sym) { + return Some(info.addr); + } + + // 2. Try enclosing scopes (capture chain) + for i in (0..current_fn_idx).rev() { + if let Some((info, scope_idx)) = self.functions[i].resolve_local(sym) { + // If it's in the root script and within a frozen scope, translate to Global + let is_frozen_root = i == 0 && (scope_idx as i32) <= self.fixed_scope_idx; + + if let Address::Global(_) = info.addr { + return Some(info.addr); + } + + let mut addr = info.addr; + + if is_frozen_root + && let Address::Local(slot) = addr { + addr = Address::Global(GlobalIdx(slot.0)); + } + + if let Address::Global(_) = addr { + return Some(addr); + } + + // Record the capture for each lambda level in between + for k in (i + 1)..=current_fn_idx { + let lambda_id = self.functions[k].identity.clone(); + self.capture_map + .entry(info.identity.clone()) + .or_default() + .insert(lambda_id); + addr = Address::Upvalue(self.functions[k].add_upvalue(addr)); + } + return Some(addr); + } + } + + // 3. Global Fallback (search in root level with context removed) + if sym.context.is_some() { + let fallback_sym = Symbol { + name: sym.name.clone(), + context: None, + }; + // Search again in all functions, primarily we care about Root Scopes + for i in (0..=current_fn_idx).rev() { + if let Some((info, _)) = self.functions[i].resolve_local(&fallback_sym) + && let Address::Global(_) = info.addr + { + return Some(info.addr); + } + } + } + + diag.push_error( + format!("Undefined variable '{}'", sym.name), + Some(identity.clone()), + ); + None + } + + fn bind_pattern( + &mut self, + node: &SyntaxNode, + kind: DeclarationKind, + diag: &mut Diagnostics, + ) -> Node { + match &node.kind { + SyntaxKind::Identifier { symbol: sym, .. } => { + if let Some(addr) = self.declare_variable(sym, node.identity.clone(), kind, diag) { + self.make_node( + node.identity.clone(), + NodeKind::Identifier { + symbol: sym.clone(), + binding: IdentifierBinding::Declaration { addr, kind }, + }, + ) + } else { + self.make_node(node.identity.clone(), NodeKind::Error) + } + } + SyntaxKind::Tuple { elements } => { + let mut bound_elems = Vec::new(); + for e in elements { + bound_elems.push(Rc::new(self.bind_pattern(e, kind, diag))); + } + self.make_node( + node.identity.clone(), + NodeKind::Tuple { + elements: bound_elems, + }, + ) + } + _ => { + diag.push_error( + format!("Invalid node in pattern: {:?}", node.kind), + Some(node.identity.clone()), + ); + self.make_node(node.identity.clone(), NodeKind::Error) + } + } + } + + fn bind_assign_pattern( + &mut self, + node: &SyntaxNode, + diag: &mut Diagnostics, + ) -> Node { + match &node.kind { + SyntaxKind::Identifier { symbol: sym, .. } => { + if let Some(addr) = self.resolve_variable(sym, diag, &node.identity) { + self.make_node( + node.identity.clone(), + NodeKind::Identifier { + symbol: sym.clone(), + binding: IdentifierBinding::Reference(addr), + }, + ) + } else { + self.make_node(node.identity.clone(), NodeKind::Error) + } + } + SyntaxKind::Tuple { elements } => { + let mut bound_elems = Vec::new(); + for e in elements { + bound_elems.push(Rc::new(self.bind_assign_pattern(e, diag))); + } + self.make_node( + node.identity.clone(), + NodeKind::Tuple { + elements: bound_elems, + }, + ) + } + _ => { + diag.push_error( + format!("Invalid node in assignment pattern: {:?}", node.kind), + Some(node.identity.clone()), + ); + self.make_node(node.identity.clone(), NodeKind::Error) + } + } + } + + fn make_node(&self, identity: Identity, kind: NodeKind) -> Node { + Node { + identity, + kind, + ty: (), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::diagnostics::Diagnostics; + use crate::ast::parser::Parser; + + #[test] + fn test_upvalue_capture_sets_is_boxed() { + let source = "(fn [] (do (def x 10) (def f (fn [] x)) x))"; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut diagnostics = Diagnostics::new(); + let (bound, captures, _scopes, _slots) = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics).unwrap(); + + if let NodeKind::Lambda { body, .. } = &bound.kind { + if let NodeKind::Block { exprs } = &body.kind { + let x_decl = &exprs[0]; + if let NodeKind::Def { pattern, .. } = &x_decl.kind { + if let NodeKind::Identifier { binding: IdentifierBinding::Declaration { addr, .. }, .. } = &pattern.kind { + assert!(matches!(addr, Address::Local(_))); + assert!( + captures.contains_key(&x_decl.identity), + "Variable 'x' should have capturers" + ); + } else { + panic!("Def pattern should be an Identifier with Declaration binding"); + } + } else { + panic!("First expression should be Def"); + } + } else { + panic!("Lambda body should be a Block"); + } + } else { + panic!("Root should be a Lambda"); + } + } + + #[test] + fn test_no_capture_not_boxed() { + let source = "(fn [] (do (def x 10) x))"; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut diagnostics = Diagnostics::new(); + let (bound, captures, _scopes, _slots) = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics).unwrap(); + + if let NodeKind::Lambda { body, .. } = &bound.kind { + if let NodeKind::Block { exprs } = &body.kind { + let x_decl = &exprs[0]; + if let NodeKind::Def { pattern, .. } = &x_decl.kind { + if let NodeKind::Identifier { binding: IdentifierBinding::Declaration { addr, .. }, .. } = &pattern.kind { + assert!(matches!(addr, Address::Local(_))); + assert!( + !captures.contains_key(&x_decl.identity), + "Variable 'x' should NOT have any capturers" + ); + } else { + panic!("Def pattern should be an Identifier with Declaration binding"); + } + } else { + panic!("First expression should be Def"); + } + } else { + panic!("Lambda body should be a Block"); + } + } else { + panic!("Root should be a Lambda"); + } + } + + #[test] + fn test_redefinition_error() { + let source = "(do (def x 1) (def x 2) x)"; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut diagnostics = Diagnostics::new(); + let _ = Binder::bind_root(vec![], 0, 0, &syntax, &mut diagnostics); + + assert!(diagnostics.has_errors()); + assert!(diagnostics.items.iter().any(|i| i.message.contains("already defined"))); + } + + #[test] + fn test_repro_global_redefinition() { + let source1 = "(def x 1) 1"; + let syntax1 = Parser::new(source1).parse_expression(); + let mut diagnostics = Diagnostics::new(); + let (_, _, scopes1, slots1) = Binder::bind_root(vec![], 0, -1, &syntax1, &mut diagnostics).unwrap(); + + let source2 = "(def x 2) 2"; + let syntax2 = Parser::new(source2).parse_expression(); + let mut diagnostics2 = Diagnostics::new(); + // Here we simulate frozen scope by passing fixed_scope_idx = 0 + let _ = Binder::bind_root(scopes1, slots1, 0, &syntax2, &mut diagnostics2); + + assert!(diagnostics2.has_errors()); + assert!(diagnostics2.items.iter().any(|i| i.message.contains("frozen/immutable"))); + } +} diff --git a/src/ast/compiler/bound_nodes.rs b/src/ast/compiler/bound_nodes.rs index 7f8c06c..b985361 100644 --- a/src/ast/compiler/bound_nodes.rs +++ b/src/ast/compiler/bound_nodes.rs @@ -181,11 +181,39 @@ impl CompilerPhase for RuntimePhase { type RecordLayout = Arc; } -/// A bound AST node, decorated with phase-specific information P. +/// Convenience alias for all post-binding phases that share the same +/// concrete binding, def, assign, lambda, and record-layout types. +/// Covers `BoundPhase`, `TypedPhase`, `AnalyzedPhase` (but not `SyntaxPhase` or `RuntimePhase`). +pub trait BoundLike: + CompilerPhase< + LocalAddress = VirtualId, + Binding = IdentifierBinding, + DefInfo = DefBinding, + AssignInfo = AssignBinding, + LambdaInfo = LambdaBinding, + RecordLayout = Arc, + > +{ +} + +impl

BoundLike for P where + P: CompilerPhase< + LocalAddress = VirtualId, + Binding = IdentifierBinding, + DefInfo = DefBinding, + AssignInfo = AssignBinding, + LambdaInfo = LambdaBinding, + RecordLayout = Arc, + > +{ +} + +/// A unified AST node, decorated with phase-specific information P. +/// Replaces both `SyntaxNode` (parser output) and the old bound AST node. #[derive(Debug, PartialEq)] pub struct Node { pub identity: Identity, - pub kind: BoundKind

, + pub kind: NodeKind

, pub ty: P::Metadata, } @@ -418,7 +446,7 @@ impl BoundKind

{ BoundKind::Destructure { .. } => "DESTRUCTURE".to_string(), BoundKind::Lambda { params, upvalues, .. } => { let p_str = match ¶ms.kind { - BoundKind::Tuple { elements } => format!("p:{}", elements.len()), + NodeKind::Tuple { elements } => format!("p:{}", elements.len()), _ => "p:1".to_string(), }; format!("LAMBDA({}, Captures:{})", p_str, upvalues.len()) diff --git a/src/ast/compiler/captures.rs b/src/ast/compiler/captures.rs index ed8eedb..3d6dada 100644 --- a/src/ast/compiler/captures.rs +++ b/src/ast/compiler/captures.rs @@ -1,129 +1,122 @@ -use crate::ast::compiler::bound_nodes::{BoundKind, CompilerPhase, Node}; -use crate::ast::types::Identity; -use std::collections::HashMap; -use std::rc::Rc; - -pub struct CapturePass; - -impl CapturePass { - pub fn apply(node: Node

, capture_map: &HashMap>) -> Node

{ - Self::transform(node, capture_map) - } - - fn transform(mut node: Node

, capture_map: &HashMap>) -> Node

{ - node.kind = match node.kind { - BoundKind::Define { - name, - addr, - kind, - value, - mut captured_by, - } => { - if let Some(capturers) = capture_map.get(&node.identity) { - captured_by.extend(capturers.iter().cloned()); - } - BoundKind::Define { - name, - addr, - kind, - value: Rc::new(Self::transform(value.as_ref().clone(), capture_map)), - captured_by, - } - } - - BoundKind::Lambda { - params, - upvalues, - body, - positional_count, - } => BoundKind::Lambda { - params: Rc::new(Self::transform(params.as_ref().clone(), capture_map)), - upvalues, - body: Rc::new(Self::transform(body.as_ref().clone(), capture_map)), - positional_count, - }, - - BoundKind::Call { callee, args } => BoundKind::Call { - callee: Rc::new(Self::transform(callee.as_ref().clone(), capture_map)), - args: Rc::new(Self::transform(args.as_ref().clone(), capture_map)), - }, - - BoundKind::Again { args } => BoundKind::Again { - args: Rc::new(Self::transform(args.as_ref().clone(), capture_map)), - }, - - BoundKind::If { - cond, - then_br, - else_br, - } => BoundKind::If { - cond: Rc::new(Self::transform(cond.as_ref().clone(), capture_map)), - then_br: Rc::new(Self::transform(then_br.as_ref().clone(), capture_map)), - else_br: else_br - .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map))), - }, - - BoundKind::Destructure { pattern, value } => BoundKind::Destructure { - pattern: Rc::new(Self::transform(pattern.as_ref().clone(), capture_map)), - value: Rc::new(Self::transform(value.as_ref().clone(), capture_map)), - }, - - BoundKind::Pipe { - inputs, - lambda, - out_type, - } => { - let mut t_inputs = Vec::with_capacity(inputs.len()); - for input in inputs { - t_inputs.push(Rc::new(Self::transform(input.as_ref().clone(), capture_map))); - } - BoundKind::Pipe { - inputs: t_inputs, - lambda: Rc::new(Self::transform(lambda.as_ref().clone(), capture_map)), - out_type, - } - } - - BoundKind::Block { exprs } => BoundKind::Block { - exprs: exprs - .into_iter() - .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map))) - .collect(), - }, - - BoundKind::Tuple { elements } => BoundKind::Tuple { - elements: elements - .into_iter() - .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map))) - .collect(), - }, - - BoundKind::Record { layout, values } => BoundKind::Record { - layout, - values: values - .into_iter() - .map(|v| Rc::new(Self::transform(v.as_ref().clone(), capture_map))) - .collect(), - }, - - BoundKind::Expansion { - original_call, - bound_expanded, - } => BoundKind::Expansion { - original_call, - bound_expanded: Rc::new(Self::transform( - bound_expanded.as_ref().clone(), - capture_map, - )), - }, - - BoundKind::GetField { rec, field } => BoundKind::GetField { - rec: Rc::new(Self::transform(rec.as_ref().clone(), capture_map)), - field, - }, - - other => other, - }; - node - } -} +use crate::ast::compiler::bound_nodes::{CompilerPhase, DefBinding, Node, NodeKind}; +use crate::ast::types::Identity; +use std::collections::HashMap; +use std::rc::Rc; + +pub struct CapturePass; + +impl CapturePass { + pub fn apply>(node: Node

, capture_map: &HashMap>) -> Node

{ + Self::transform(node, capture_map) + } + + fn transform>(mut node: Node

, capture_map: &HashMap>) -> Node

{ + node.kind = match node.kind { + NodeKind::Def { + pattern, + value, + mut info, + } => { + if let Some(capturers) = capture_map.get(&node.identity) { + info.captured_by.extend(capturers.iter().cloned()); + } + NodeKind::Def { + pattern: Rc::new(Self::transform(pattern.as_ref().clone(), capture_map)), + value: Rc::new(Self::transform(value.as_ref().clone(), capture_map)), + info, + } + } + + NodeKind::Lambda { + params, + body, + info, + } => NodeKind::Lambda { + params: Rc::new(Self::transform(params.as_ref().clone(), capture_map)), + body: Rc::new(Self::transform(body.as_ref().clone(), capture_map)), + info, + }, + + NodeKind::Call { callee, args } => NodeKind::Call { + callee: Rc::new(Self::transform(callee.as_ref().clone(), capture_map)), + args: Rc::new(Self::transform(args.as_ref().clone(), capture_map)), + }, + + NodeKind::Again { args } => NodeKind::Again { + args: Rc::new(Self::transform(args.as_ref().clone(), capture_map)), + }, + + NodeKind::If { + cond, + then_br, + else_br, + } => NodeKind::If { + cond: Rc::new(Self::transform(cond.as_ref().clone(), capture_map)), + then_br: Rc::new(Self::transform(then_br.as_ref().clone(), capture_map)), + else_br: else_br + .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map))), + }, + + NodeKind::Assign { target, value, info } => NodeKind::Assign { + target: Rc::new(Self::transform(target.as_ref().clone(), capture_map)), + value: Rc::new(Self::transform(value.as_ref().clone(), capture_map)), + info, + }, + + NodeKind::Pipe { + inputs, + lambda, + } => { + let mut t_inputs = Vec::with_capacity(inputs.len()); + for input in inputs { + t_inputs.push(Rc::new(Self::transform(input.as_ref().clone(), capture_map))); + } + NodeKind::Pipe { + inputs: t_inputs, + lambda: Rc::new(Self::transform(lambda.as_ref().clone(), capture_map)), + } + } + + NodeKind::Block { exprs } => NodeKind::Block { + exprs: exprs + .into_iter() + .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map))) + .collect(), + }, + + NodeKind::Tuple { elements } => NodeKind::Tuple { + elements: elements + .into_iter() + .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map))) + .collect(), + }, + + NodeKind::Record { fields, layout } => NodeKind::Record { + fields: fields + .into_iter() + .map(|(k, v)| (k, Rc::new(Self::transform(v.as_ref().clone(), capture_map)))) + .collect(), + layout, + }, + + NodeKind::Expansion { + original_call, + expanded, + } => NodeKind::Expansion { + original_call, + expanded: Rc::new(Self::transform( + expanded.as_ref().clone(), + capture_map, + )), + }, + + NodeKind::GetField { rec, field } => NodeKind::GetField { + rec: Rc::new(Self::transform(rec.as_ref().clone(), capture_map)), + field, + }, + + other => other, + }; + node + } +} diff --git a/src/ast/compiler/dumper.rs b/src/ast/compiler/dumper.rs index d862097..e6b6ccc 100644 --- a/src/ast/compiler/dumper.rs +++ b/src/ast/compiler/dumper.rs @@ -1,239 +1,240 @@ -use crate::ast::compiler::bound_nodes::{BoundKind, CompilerPhase, Node}; - -/// Human-readable AST dumper for the bound AST. -pub struct Dumper { - output: String, - indent: usize, -} - -impl Dumper { - /// Produces a formatted string representation of the given bound AST node and its children. - pub fn dump(node: &Node

) -> String { - let mut dumper = Self { - output: String::new(), - indent: 0, - }; - dumper.visit(node); - dumper.output - } - - fn write_indent(&mut self) { - for _ in 0..self.indent { - self.output.push_str(" "); - } - } - - fn log(&mut self, label: &str, node: &Node

) { - self.write_indent(); - self.output.push_str(label); - self.output - .push_str(&format!(" \n", node.ty)); - } - - fn visit(&mut self, node: &Node

) { - match &node.kind { - BoundKind::Nop => self.log("Nop", node), - BoundKind::Constant(v) => { - self.log(&format!("Constant: {}", v), node); - } - BoundKind::Get { addr, name } => { - self.log(&format!("Get: {} ({:?})", name.name, addr), node) - } - - BoundKind::FieldAccessor(k) => self.log(&format!("FieldAccessor: .{}", k.name()), node), - - BoundKind::GetField { rec, field } => { - self.log(&format!("GetField: .{}", field.name()), node); - self.indent += 1; - self.visit(rec); - self.indent -= 1; - } - - BoundKind::Set { addr, value } => { - self.log(&format!("Set: {:?}", addr), node); - self.indent += 1; - self.visit(value); - self.indent -= 1; - } - - BoundKind::Define { - name, - addr, - kind, - value, - captured_by, - } => { - let k_str = match kind { - crate::ast::compiler::bound_nodes::DeclarationKind::Variable => "Variable", - crate::ast::compiler::bound_nodes::DeclarationKind::Parameter => "Parameter", - }; - let capture_info = if captured_by.is_empty() { - String::from("not captured") - } else { - format!("captured by {} lambdas", captured_by.len()) - }; - self.log( - &format!( - "Define {} (Name: '{}', Address: {:?}, {})", - k_str, name.name, addr, capture_info - ), - node, - ); - - self.indent += 1; - if !captured_by.is_empty() { - for capturer in captured_by { - self.write_indent(); - let loc = capturer - .location - .unwrap_or(crate::ast::types::SourceLocation { line: 0, col: 0 }); - self.output.push_str(&format!( - "- Capturer: Lambda at line {}, col {}\n", - loc.line, loc.col - )); - } - } - self.visit(value); - self.indent -= 1; - } - - BoundKind::Destructure { pattern, value } => { - self.log("Destructure", node); - self.indent += 1; - self.write_indent(); - self.output.push_str("Pattern:\n"); - self.visit(pattern); - self.write_indent(); - self.output.push_str("Value:\n"); - self.visit(value); - self.indent -= 1; - } - - BoundKind::If { - cond, - then_br, - else_br, - } => { - self.log("If", node); - self.indent += 1; - - self.write_indent(); - self.output.push_str("Condition:\n"); - self.visit(cond); - - self.write_indent(); - self.output.push_str("Then:\n"); - self.visit(then_br); - - if let Some(e) = else_br { - self.write_indent(); - self.output.push_str("Else:\n"); - self.visit(e); - } - self.indent -= 1; - } - - BoundKind::Lambda { - params, - upvalues, - body, - .. - } => { - self.log(&format!("Lambda (Upvalues: {})", upvalues.len()), node); - self.indent += 1; - - self.write_indent(); - self.output.push_str("Parameters:\n"); - self.visit(params); - - if !upvalues.is_empty() { - self.write_indent(); - self.output.push_str(&format!("Upvalues: {:?}\n", upvalues)); - } - self.visit(body); - self.indent -= 1; - } - - BoundKind::Call { callee, args } => { - self.log("Call", node); - self.indent += 1; - - self.write_indent(); - self.output.push_str("Callee:\n"); - self.visit(callee); - - self.write_indent(); - self.output.push_str("Arguments:\n"); - self.visit(args); - - self.indent -= 1; - } - - BoundKind::Again { args } => { - self.log("Again", node); - self.indent += 1; - self.visit(args); - self.indent -= 1; - } - BoundKind::Pipe { inputs, lambda, .. } => { - self.log("Pipe", node); - self.indent += 1; - for input in inputs { - self.visit(input); - } - self.visit(lambda); - self.indent -= 1; - } - - BoundKind::Block { exprs } => { - self.log("Block", node); - self.indent += 1; - for expr in exprs { - self.visit(expr); - } - self.indent -= 1; - } - - BoundKind::Tuple { elements } => { - self.log("Tuple", node); - self.indent += 1; - for el in elements { - self.visit(el); - } - self.indent -= 1; - } - - BoundKind::Record { layout, values } => { - self.log( - &format!("Record (Layout: {} fields)", layout.fields.len()), - node, - ); - self.indent += 1; - for v in values { - self.visit(v); - } - self.indent -= 1; - } - - BoundKind::Expansion { - original_call, - bound_expanded, - } => { - self.log( - &format!("Expansion (Original: {:?})", original_call.kind), - node, - ); - self.indent += 1; - self.visit(bound_expanded); - self.indent -= 1; - } - - BoundKind::Extension(ext) => { - self.log(&ext.display_name(), node); - } - BoundKind::Error => { - self.log("ERROR_NODE", node); - } - } - } -} +use crate::ast::compiler::bound_nodes::{CompilerPhase, Node, NodeKind}; + +/// Human-readable AST dumper for the bound AST. +pub struct Dumper { + output: String, + indent: usize, +} + +impl Dumper { + /// Produces a formatted string representation of the given bound AST node and its children. + pub fn dump(node: &Node

) -> String { + let mut dumper = Self { + output: String::new(), + indent: 0, + }; + dumper.visit(node); + dumper.output + } + + fn write_indent(&mut self) { + for _ in 0..self.indent { + self.output.push_str(" "); + } + } + + fn log(&mut self, label: &str, node: &Node

) { + self.write_indent(); + self.output.push_str(label); + self.output + .push_str(&format!(" \n", node.ty)); + } + + fn visit(&mut self, node: &Node

) { + match &node.kind { + NodeKind::Nop => self.log("Nop", node), + NodeKind::Constant(v) => { + self.log(&format!("Constant: {}", v), node); + } + NodeKind::Identifier { symbol, binding } => { + self.log(&format!("Identifier: {} ({:?})", symbol.name, binding), node) + } + + NodeKind::FieldAccessor(k) => self.log(&format!("FieldAccessor: .{}", k.name()), node), + + NodeKind::GetField { rec, field } => { + self.log(&format!("GetField: .{}", field.name()), node); + self.indent += 1; + self.visit(rec); + self.indent -= 1; + } + + NodeKind::Assign { target, value, info } => { + self.log(&format!("Assign: {:?}", info), node); + self.indent += 1; + self.write_indent(); + self.output.push_str("Target:\n"); + self.visit(target); + self.write_indent(); + self.output.push_str("Value:\n"); + self.visit(value); + self.indent -= 1; + } + + NodeKind::Def { + pattern, + value, + info, + } => { + self.log(&format!("Def ({:?})", info), node); + self.indent += 1; + self.write_indent(); + self.output.push_str("Pattern:\n"); + self.visit(pattern); + self.write_indent(); + self.output.push_str("Value:\n"); + self.visit(value); + self.indent -= 1; + } + + NodeKind::If { + cond, + then_br, + else_br, + } => { + self.log("If", node); + self.indent += 1; + + self.write_indent(); + self.output.push_str("Condition:\n"); + self.visit(cond); + + self.write_indent(); + self.output.push_str("Then:\n"); + self.visit(then_br); + + if let Some(e) = else_br { + self.write_indent(); + self.output.push_str("Else:\n"); + self.visit(e); + } + self.indent -= 1; + } + + NodeKind::Lambda { + params, + body, + info, + } => { + self.log(&format!("Lambda ({:?})", info), node); + self.indent += 1; + + self.write_indent(); + self.output.push_str("Parameters:\n"); + self.visit(params); + + self.visit(body); + self.indent -= 1; + } + + NodeKind::Call { callee, args } => { + self.log("Call", node); + self.indent += 1; + + self.write_indent(); + self.output.push_str("Callee:\n"); + self.visit(callee); + + self.write_indent(); + self.output.push_str("Arguments:\n"); + self.visit(args); + + self.indent -= 1; + } + + NodeKind::Again { args } => { + self.log("Again", node); + self.indent += 1; + self.visit(args); + self.indent -= 1; + } + NodeKind::Pipe { inputs, lambda } => { + self.log("Pipe", node); + self.indent += 1; + for input in inputs { + self.visit(input); + } + self.visit(lambda); + self.indent -= 1; + } + + NodeKind::Block { exprs } => { + self.log("Block", node); + self.indent += 1; + for expr in exprs { + self.visit(expr); + } + self.indent -= 1; + } + + NodeKind::Tuple { elements } => { + self.log("Tuple", node); + self.indent += 1; + for el in elements { + self.visit(el); + } + self.indent -= 1; + } + + NodeKind::Record { fields, layout } => { + self.log( + &format!("Record (Layout: {:?})", layout), + node, + ); + self.indent += 1; + for (key, val) in fields { + self.write_indent(); + self.output.push_str("Key:\n"); + self.indent += 1; + self.visit(key); + self.indent -= 1; + self.write_indent(); + self.output.push_str("Value:\n"); + self.indent += 1; + self.visit(val); + self.indent -= 1; + } + self.indent -= 1; + } + + NodeKind::Expansion { + original_call, + expanded, + } => { + self.log( + &format!("Expansion (Original: {:?})", original_call.kind), + node, + ); + self.indent += 1; + self.visit(expanded); + self.indent -= 1; + } + + NodeKind::MacroDecl { name, params, body } => { + self.log(&format!("MacroDecl: {}", name.name), node); + self.indent += 1; + self.visit(params); + self.visit(body); + self.indent -= 1; + } + + NodeKind::Template(inner) => { + self.log("Template", node); + self.indent += 1; + self.visit(inner); + self.indent -= 1; + } + + NodeKind::Placeholder(inner) => { + self.log("Placeholder", node); + self.indent += 1; + self.visit(inner); + self.indent -= 1; + } + + NodeKind::Splice(inner) => { + self.log("Splice", node); + self.indent += 1; + self.visit(inner); + self.indent -= 1; + } + + NodeKind::Extension(ext) => { + self.log(&ext.display_name(), node); + } + NodeKind::Error => { + self.log("ERROR_NODE", node); + } + } + } +} diff --git a/src/ast/compiler/lambda_collector.rs b/src/ast/compiler/lambda_collector.rs index 85dd565..8af592e 100644 --- a/src/ast/compiler/lambda_collector.rs +++ b/src/ast/compiler/lambda_collector.rs @@ -1,99 +1,111 @@ -use crate::ast::compiler::bound_nodes::{Address, BoundKind, CompilerPhase, GlobalIdx, Node}; -use std::collections::HashMap; -use std::rc::Rc; - -/// 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, P: CompilerPhase> { - registry: &'a mut HashMap>>, -} - -impl<'a, P: CompilerPhase> LambdaCollector<'a, P> { - /// Performs a full traversal of the AST and populates the provided registry. - pub fn collect(node: &Node

, registry: &'a mut HashMap>>) { - let mut collector = Self { registry }; - collector.visit(node); - } - - fn visit(&mut self, node: &Node

) { - match &node.kind { - BoundKind::Block { exprs } => { - for expr in exprs { - self.visit(expr); - } - } - - BoundKind::Define { addr, value, .. } => { - // Register global function definitions (lambdas) - if let Address::Global(global_index) = addr { - let mut current = value; - while let BoundKind::Expansion { bound_expanded, .. } = ¤t.kind { - current = bound_expanded; - } - - if let BoundKind::Lambda { .. } = ¤t.kind { - self.registry - .insert(*global_index, (*current).clone()); - } - } - self.visit(value); - } - - BoundKind::Set { addr, value } => { - // Also track assignments to globals if they hold lambdas. - if let Address::Global(global_index) = addr { - let mut current = value; - while let BoundKind::Expansion { bound_expanded, .. } = ¤t.kind { - current = bound_expanded; - } - - if let BoundKind::Lambda { .. } = ¤t.kind { - self.registry - .insert(*global_index, (*current).clone()); - } - } - self.visit(value); - } - - BoundKind::If { - cond, - then_br, - else_br, - } => { - self.visit(cond); - self.visit(then_br); - if let Some(e) = else_br { - self.visit(e); - } - } - - BoundKind::Lambda { params, body, .. } => { - self.visit(params); - self.visit(body); - } - - BoundKind::Call { callee, args } => { - self.visit(callee); - self.visit(args); - } - - BoundKind::Tuple { elements } => { - for el in elements { - self.visit(el); - } - } - - BoundKind::Record { values, .. } => { - for v in values { - self.visit(v); - } - } - - BoundKind::Expansion { bound_expanded, .. } => { - self.visit(bound_expanded); - } - - _ => {} // Leaf nodes - } - } -} +use crate::ast::compiler::bound_nodes::{ + Address, BoundLike, GlobalIdx, IdentifierBinding, Node, NodeKind, +}; +use std::collections::HashMap; +use std::rc::Rc; + +/// 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, P: BoundLike> { + registry: &'a mut HashMap>>, +} + +impl<'a, P> LambdaCollector<'a, P> +where + P: BoundLike, +{ + /// Performs a full traversal of the AST and populates the provided registry. + pub fn collect(node: &Node

, registry: &'a mut HashMap>>) { + let mut collector = Self { registry }; + collector.visit(node); + } + + fn visit(&mut self, node: &Node

) { + match &node.kind { + NodeKind::Block { exprs } => { + for expr in exprs { + self.visit(expr); + } + } + + NodeKind::Def { pattern, value, .. } => { + // Register global function definitions (lambdas) + if let NodeKind::Identifier { + binding: IdentifierBinding::Declaration { + addr: Address::Global(global_index), + .. + }, + .. + } = &pattern.kind + { + let mut current = value; + while let NodeKind::Expansion { expanded, .. } = ¤t.kind { + current = expanded; + } + + if let NodeKind::Lambda { .. } = ¤t.kind { + self.registry + .insert(*global_index, (*current).clone()); + } + } + self.visit(value); + } + + NodeKind::Assign { value, info, .. } => { + // Also track assignments to globals if they hold lambdas. + if let Some(Address::Global(global_index)) = &info.addr { + let mut current = value; + while let NodeKind::Expansion { expanded, .. } = ¤t.kind { + current = expanded; + } + + if let NodeKind::Lambda { .. } = ¤t.kind { + self.registry + .insert(*global_index, (*current).clone()); + } + } + self.visit(value); + } + + NodeKind::If { + cond, + then_br, + else_br, + } => { + self.visit(cond); + self.visit(then_br); + if let Some(e) = else_br { + self.visit(e); + } + } + + NodeKind::Lambda { params, body, .. } => { + self.visit(params); + self.visit(body); + } + + NodeKind::Call { callee, args } => { + self.visit(callee); + self.visit(args); + } + + NodeKind::Tuple { elements } => { + for el in elements { + self.visit(el); + } + } + + NodeKind::Record { fields, .. } => { + for (_, v) in fields { + self.visit(v); + } + } + + NodeKind::Expansion { expanded, .. } => { + self.visit(expanded); + } + + _ => {} // Leaf nodes + } + } +} diff --git a/src/ast/compiler/lowering.rs b/src/ast/compiler/lowering.rs index 1b0155b..552a4e9 100644 --- a/src/ast/compiler/lowering.rs +++ b/src/ast/compiler/lowering.rs @@ -1,231 +1,283 @@ -use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, ExecNode, Node, RuntimeMetadata, StackOffset, VirtualId}; -use std::collections::HashMap; -use std::rc::Rc; - -struct StackAllocator { - mapping: HashMap, - next_slot: u32, -} - -impl StackAllocator { - fn new() -> Self { - Self { - mapping: HashMap::new(), - next_slot: 0, - } - } - - fn map_slot(&mut self, slot: VirtualId) -> StackOffset { - let entry = self.mapping.entry(slot.0).or_insert_with(|| { - let s = self.next_slot; - self.next_slot += 1; - s - }); - StackOffset(*entry) - } - - fn map_address(&mut self, addr: Address) -> Address { - match addr { - Address::Local(slot) => Address::Local(self.map_slot(slot)), - Address::Upvalue(idx) => Address::Upvalue(idx), - Address::Global(idx) => Address::Global(idx), - } - } -} - -pub struct Lowering; - -impl Lowering { - /// Lowers an AnalyzedNode to an ExecNode, marking tail positions and calculating stack sizes. - pub fn lower(node: AnalyzedNode) -> ExecNode { - let mut allocator = StackAllocator::new(); - let mut exec_node = Self::transform(Rc::new(node), true, &mut allocator); - - // If the top-level node is a Lambda, it already has its internal stack_size - // calculated during transform(). For non-lambdas (like raw expressions), - // we use the allocator's next_slot to determine the required root stack size. - if !matches!(exec_node.kind, BoundKind::Lambda { .. }) { - exec_node.ty.stack_size = allocator.next_slot; - } - exec_node - } - - fn transform( - node_rc: Rc, - is_tail_position: bool, - allocator: &mut StackAllocator, - ) -> ExecNode { - let node = &*node_rc; - let mut lambda_stack_size = 0; - - let new_kind = match &node.kind { - BoundKind::Call { callee, args } => BoundKind::Call { - callee: Rc::new(Self::transform(callee.clone(), false, allocator)), - args: Rc::new(Self::transform(args.clone(), false, allocator)), - }, - - BoundKind::Again { args } => { - if !is_tail_position { - panic!("'again' is only allowed in tail position to avoid dead code."); - } - BoundKind::Again { - args: Rc::new(Self::transform(args.clone(), false, allocator)), - } - } - BoundKind::Pipe { - inputs, - lambda, - out_type, - } => { - let mut t_inputs = Vec::with_capacity(inputs.len()); - for input in inputs { - t_inputs.push(Rc::new(Self::transform(input.clone(), false, allocator))); - } - BoundKind::Pipe { - inputs: t_inputs, - lambda: Rc::new(Self::transform(lambda.clone(), false, allocator)), - out_type: out_type.clone(), - } - } - - BoundKind::If { - cond, - then_br, - else_br, - } => BoundKind::If { - cond: Rc::new(Self::transform(cond.clone(), false, allocator)), - then_br: Rc::new(Self::transform(then_br.clone(), is_tail_position, allocator)), - else_br: else_br - .as_ref() - .map(|e| Rc::new(Self::transform(e.clone(), is_tail_position, allocator))), - }, - - BoundKind::Block { exprs } => { - if exprs.is_empty() { - BoundKind::Block { exprs: vec![] } - } else { - let last_idx = exprs.len() - 1; - let mut new_exprs = Vec::with_capacity(exprs.len()); - - for (i, expr) in exprs.iter().enumerate() { - let is_last = i == last_idx; - new_exprs.push(Rc::new(Self::transform( - expr.clone(), - is_tail_position && is_last, - allocator, - ))); - } - BoundKind::Block { exprs: new_exprs } - } - } - - BoundKind::Lambda { - params, - upvalues, - body, - positional_count, - } => { - // Upvalues refer to the PARENT scope's addresses. - let mapped_upvalues = upvalues - .iter() - .map(|a| allocator.map_address(*a)) - .collect(); - - // New allocator for the lambda's own stack frame - let mut lambda_allocator = StackAllocator::new(); - let t_params = Rc::new(Self::transform(params.clone(), false, &mut lambda_allocator)); - let t_body = Rc::new(Self::transform(body.clone(), true, &mut lambda_allocator)); - lambda_stack_size = lambda_allocator.next_slot; - - BoundKind::Lambda { - params: t_params, - upvalues: mapped_upvalues, - body: t_body, - positional_count: *positional_count, - } - } - - BoundKind::Set { addr, value } => BoundKind::Set { - addr: allocator.map_address(*addr), - value: Rc::new(Self::transform(value.clone(), false, allocator)), - }, - BoundKind::Define { - name, - addr, - kind, - value, - captured_by, - } => BoundKind::Define { - name: name.clone(), - addr: allocator.map_address(*addr), - kind: *kind, - value: Rc::new(Self::transform(value.clone(), false, allocator)), - captured_by: captured_by.clone(), - }, - BoundKind::Destructure { pattern, value } => BoundKind::Destructure { - pattern: Rc::new(Self::transform(pattern.clone(), false, allocator)), - value: Rc::new(Self::transform(value.clone(), false, allocator)), - }, - BoundKind::Record { layout, values } => { - let new_values = values - .iter() - .map(|v| Rc::new(Self::transform(v.clone(), false, allocator))) - .collect(); - BoundKind::Record { - layout: layout.clone(), - values: new_values, - } - } - BoundKind::Tuple { elements } => { - let new_elements = elements - .iter() - .map(|e| Rc::new(Self::transform(e.clone(), false, allocator))) - .collect(); - BoundKind::Tuple { - elements: new_elements, - } - } - BoundKind::Constant(v) => BoundKind::Constant(v.clone()), - BoundKind::Get { addr, name } => BoundKind::Get { - addr: allocator.map_address(*addr), - name: name.clone(), - }, - BoundKind::FieldAccessor(k) => BoundKind::FieldAccessor(*k), - BoundKind::GetField { rec, field } => BoundKind::GetField { - rec: Rc::new(Self::transform(rec.clone(), false, allocator)), - field: *field, - }, - BoundKind::Nop => BoundKind::Nop, - BoundKind::Expansion { - original_call, - bound_expanded, - } => BoundKind::Expansion { - original_call: original_call.clone(), - bound_expanded: Rc::new(Self::transform( - bound_expanded.clone(), - is_tail_position, - allocator, - )), - }, - BoundKind::Extension(_) => BoundKind::Nop, - BoundKind::Error => BoundKind::Error, - }; - - let stack_size = if let BoundKind::Lambda { .. } = &new_kind { - lambda_stack_size - } else { - 0 - }; - - Node { - identity: node.identity.clone(), - kind: new_kind, - ty: RuntimeMetadata { - ty: node.ty.original.ty.clone(), - is_tail: is_tail_position, - original: node_rc, - stack_size, - }, - } - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AnalyzedNode, AssignBinding, ExecNode, + IdentifierBinding, LambdaBinding, Node, NodeKind, RuntimeMetadata, + StackOffset, VirtualId, +}; +use std::collections::HashMap; +use std::rc::Rc; + +struct StackAllocator { + mapping: HashMap, + next_slot: u32, +} + +impl StackAllocator { + fn new() -> Self { + Self { + mapping: HashMap::new(), + next_slot: 0, + } + } + + fn map_slot(&mut self, slot: VirtualId) -> StackOffset { + let entry = self.mapping.entry(slot.0).or_insert_with(|| { + let s = self.next_slot; + self.next_slot += 1; + s + }); + StackOffset(*entry) + } + + fn map_address(&mut self, addr: Address) -> Address { + match addr { + Address::Local(slot) => Address::Local(self.map_slot(slot)), + Address::Upvalue(idx) => Address::Upvalue(idx), + Address::Global(idx) => Address::Global(idx), + } + } +} + +pub struct Lowering; + +impl Lowering { + /// Lowers an AnalyzedNode to an ExecNode, marking tail positions and calculating stack sizes. + pub fn lower(node: AnalyzedNode) -> ExecNode { + let mut allocator = StackAllocator::new(); + let mut exec_node = Self::transform(Rc::new(node), true, &mut allocator); + + // If the top-level node is a Lambda, it already has its internal stack_size + // calculated during transform(). For non-lambdas (like raw expressions), + // we use the allocator's next_slot to determine the required root stack size. + if !matches!(exec_node.kind, NodeKind::Lambda { .. }) { + exec_node.ty.stack_size = allocator.next_slot; + } + exec_node + } + + fn transform( + node_rc: Rc, + is_tail_position: bool, + allocator: &mut StackAllocator, + ) -> ExecNode { + let node = &*node_rc; + let mut lambda_stack_size = 0; + + let new_kind = match &node.kind { + NodeKind::Call { callee, args } => { + // Optimized Field Access: Call { callee: FieldAccessor, args: Tuple[1] } → GetField + if let NodeKind::FieldAccessor(k) = &callee.kind + && let NodeKind::Tuple { elements } = &args.kind + && elements.len() == 1 + { + let rec = Rc::new(Self::transform(elements[0].clone(), false, allocator)); + return Node { + identity: node.identity.clone(), + kind: NodeKind::GetField { + rec, + field: *k, + }, + ty: RuntimeMetadata { + ty: node.ty.original.ty.clone(), + is_tail: is_tail_position, + original: node_rc, + stack_size: 0, + }, + }; + } + + NodeKind::Call { + callee: Rc::new(Self::transform(callee.clone(), false, allocator)), + args: Rc::new(Self::transform(args.clone(), false, allocator)), + } + } + + NodeKind::Again { args } => { + if !is_tail_position { + panic!("'again' is only allowed in tail position to avoid dead code."); + } + NodeKind::Again { + args: Rc::new(Self::transform(args.clone(), false, allocator)), + } + } + NodeKind::Pipe { + inputs, + lambda, + } => { + let mut t_inputs = Vec::with_capacity(inputs.len()); + for input in inputs { + t_inputs.push(Rc::new(Self::transform(input.clone(), false, allocator))); + } + NodeKind::Pipe { + inputs: t_inputs, + lambda: Rc::new(Self::transform(lambda.clone(), false, allocator)), + } + } + + NodeKind::If { + cond, + then_br, + else_br, + } => NodeKind::If { + cond: Rc::new(Self::transform(cond.clone(), false, allocator)), + then_br: Rc::new(Self::transform(then_br.clone(), is_tail_position, allocator)), + else_br: else_br + .as_ref() + .map(|e| Rc::new(Self::transform(e.clone(), is_tail_position, allocator))), + }, + + NodeKind::Block { exprs } => { + if exprs.is_empty() { + NodeKind::Block { exprs: vec![] } + } else { + let last_idx = exprs.len() - 1; + let mut new_exprs = Vec::with_capacity(exprs.len()); + + for (i, expr) in exprs.iter().enumerate() { + let is_last = i == last_idx; + new_exprs.push(Rc::new(Self::transform( + expr.clone(), + is_tail_position && is_last, + allocator, + ))); + } + NodeKind::Block { exprs: new_exprs } + } + } + + NodeKind::Lambda { + params, + body, + info: lambda_info, + } => { + // Upvalues refer to the PARENT scope's addresses. + let mapped_upvalues = lambda_info + .upvalues + .iter() + .map(|a| allocator.map_address(*a)) + .collect(); + + // New allocator for the lambda's own stack frame + let mut lambda_allocator = StackAllocator::new(); + let t_params = Rc::new(Self::transform(params.clone(), false, &mut lambda_allocator)); + let t_body = Rc::new(Self::transform(body.clone(), true, &mut lambda_allocator)); + lambda_stack_size = lambda_allocator.next_slot; + + NodeKind::Lambda { + params: t_params, + body: t_body, + info: LambdaBinding { + upvalues: mapped_upvalues, + positional_count: lambda_info.positional_count, + }, + } + } + + NodeKind::Assign { target, value, info } => { + let new_info = if let Some(addr) = info.addr { + AssignBinding { + addr: Some(allocator.map_address(addr)), + } + } else { + AssignBinding { addr: None } + }; + NodeKind::Assign { + target: Rc::new(Self::transform(target.clone(), false, allocator)), + value: Rc::new(Self::transform(value.clone(), false, allocator)), + info: new_info, + } + } + NodeKind::Def { + pattern, + value, + info, + } => NodeKind::Def { + pattern: Rc::new(Self::transform(pattern.clone(), false, allocator)), + value: Rc::new(Self::transform(value.clone(), false, allocator)), + info: info.clone(), + }, + NodeKind::Identifier { symbol, binding } => { + let new_binding = match binding { + IdentifierBinding::Reference(addr) => { + IdentifierBinding::Reference(allocator.map_address(*addr)) + } + IdentifierBinding::Declaration { addr, kind } => { + IdentifierBinding::Declaration { + addr: allocator.map_address(*addr), + kind: *kind, + } + } + }; + NodeKind::Identifier { + symbol: symbol.clone(), + binding: new_binding, + } + } + NodeKind::Record { fields, layout } => { + let new_fields = fields + .iter() + .map(|(k, v)| { + ( + Rc::new(Self::transform(k.clone(), false, allocator)), + Rc::new(Self::transform(v.clone(), false, allocator)), + ) + }) + .collect(); + NodeKind::Record { + fields: new_fields, + layout: layout.clone(), + } + } + NodeKind::Tuple { elements } => { + let new_elements = elements + .iter() + .map(|e| Rc::new(Self::transform(e.clone(), false, allocator))) + .collect(); + NodeKind::Tuple { + elements: new_elements, + } + } + NodeKind::Constant(v) => NodeKind::Constant(v.clone()), + NodeKind::FieldAccessor(k) => NodeKind::FieldAccessor(*k), + NodeKind::GetField { rec, field } => NodeKind::GetField { + rec: Rc::new(Self::transform(rec.clone(), false, allocator)), + field: *field, + }, + NodeKind::Nop => NodeKind::Nop, + NodeKind::Expansion { + original_call, + expanded, + } => NodeKind::Expansion { + original_call: original_call.clone(), + expanded: Rc::new(Self::transform( + expanded.clone(), + is_tail_position, + allocator, + )), + }, + NodeKind::Extension(_) => NodeKind::Nop, + NodeKind::Error => NodeKind::Error, + // Syntax-only variants should not appear in AnalyzedPhase + NodeKind::MacroDecl { .. } + | NodeKind::Template(_) + | NodeKind::Placeholder(_) + | NodeKind::Splice(_) => NodeKind::Nop, + }; + + let stack_size = if let NodeKind::Lambda { .. } = &new_kind { + lambda_stack_size + } else { + 0 + }; + + Node { + identity: node.identity.clone(), + kind: new_kind, + ty: RuntimeMetadata { + ty: node.ty.original.ty.clone(), + is_tail: is_tail_position, + original: node_rc, + stack_size, + }, + } + } +} diff --git a/src/ast/compiler/macros.rs b/src/ast/compiler/macros.rs index eeec12b..454e646 100644 --- a/src/ast/compiler/macros.rs +++ b/src/ast/compiler/macros.rs @@ -1,836 +1,855 @@ -use crate::ast::nodes::{Symbol, SyntaxKind, SyntaxNode}; -use crate::ast::types::{Identity, Value}; -use std::collections::HashMap; -use std::rc::Rc; - -/// Trait for evaluating expressions during macro expansion. -pub trait MacroEvaluator { - fn evaluate( - &self, - node: &SyntaxNode, - bindings: &HashMap, SyntaxNode>, - ) -> Result; -} - -/// Internal state for template expansion (Hygiene context + Parameter binding) -struct ExpansionState<'a> { - bindings: &'a HashMap, SyntaxNode>, - /// The identity of the current expansion instance, used to "color" internal symbols. - expansion_id: Identity, -} - -type MacroMap = HashMap, (Vec>, SyntaxNode)>; - -/// A registry for macro declarations. -#[derive(Clone)] -pub struct MacroRegistry { - scopes: Vec>, -} - -impl Default for MacroRegistry { - fn default() -> Self { - Self::new() - } -} - -impl MacroRegistry { - pub fn new() -> Self { - Self { - scopes: vec![Rc::new(HashMap::new())], - } - } - - pub fn push(&mut self) { - self.scopes.push(Rc::new(HashMap::new())); - } - - pub fn pop(&mut self) { - if self.scopes.len() > 1 { - self.scopes.pop(); - } - } - - pub fn define(&mut self, name: Rc, params: Vec>, body: SyntaxNode) { - if let Some(current_rc) = self.scopes.last_mut() { - // Copy-on-Write: If the Rc is shared, clone the map before modifying. - let current = Rc::make_mut(current_rc); - current.insert(name, (params, body)); - } - } - - pub fn lookup(&self, name: &str) -> Option<(Vec>, SyntaxNode)> { - for scope in self.scopes.iter().rev() { - if let Some(m) = scope.get(name) { - return Some(m.clone()); - } - } - None - } -} - -pub struct MacroExpander { - registry: MacroRegistry, - evaluator: E, -} - -impl MacroExpander { - pub fn new(registry: MacroRegistry, evaluator: E) -> Self { - Self { - registry, - evaluator, - } - } - - pub fn into_registry(self) -> MacroRegistry { - self.registry - } - - pub fn expand(&mut self, node: SyntaxNode) -> Result { - self.expand_recursive(node) - } - - fn expand_recursive(&mut self, node: SyntaxNode) -> Result { - match node.kind { - SyntaxKind::MacroDecl { name, params, body } => { - let p_names = self.extract_param_names(¶ms)?; - self.registry.define(name.name, p_names, *body); - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Nop, - - }) - } - - SyntaxKind::Call { callee, args } => { - if let SyntaxKind::Identifier(ref sym) = callee.kind - && let Some((params, body)) = self.registry.lookup(&sym.name) - { - let expanded_args = self.expand_recursive(*args)?; - - // Extract argument nodes from the expanded tuple - let arg_elements = if let SyntaxKind::Tuple { elements } = expanded_args.kind { - elements - } else { - vec![expanded_args] - }; - - let expanded = self.expand_call( - node.identity.clone(), - &sym.name, - params, - arg_elements, - body, - )?; - return self.expand_recursive(expanded); - } - - let expanded_callee = self.expand_recursive(*callee)?; - let expanded_args = self.expand_recursive(*args)?; - - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Call { - callee: Box::new(expanded_callee), - args: Box::new(expanded_args), - }, - - }) - } - - SyntaxKind::Block { exprs } => { - self.registry.push(); - let mut expanded_exprs = Vec::new(); - for expr in exprs { - expanded_exprs.push(self.expand_recursive(expr)?); - } - self.registry.pop(); - - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Block { - exprs: expanded_exprs, - }, - - }) - } - - SyntaxKind::Pipe { inputs, lambda } => { - let mut expanded_inputs = Vec::with_capacity(inputs.len()); - for input in inputs { - expanded_inputs.push(self.expand_recursive(input)?); - } - let expanded_lambda = Box::new(self.expand_recursive(*lambda)?); - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Pipe { - inputs: expanded_inputs, - lambda: expanded_lambda, - }, - - }) - } - - SyntaxKind::Lambda { params, body } => { - self.registry.push(); - let expanded_params = self.expand_recursive(*params)?; - let expanded_body = self.expand_recursive((*body).clone())?; - self.registry.pop(); - - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Lambda { - params: Box::new(expanded_params), - body: Rc::new(expanded_body), - }, - - }) - } - - SyntaxKind::If { - cond, - then_br, - else_br, - } => { - let cond = self.expand_recursive(*cond)?; - let then_br = self.expand_recursive(*then_br)?; - let else_br = if let Some(e) = else_br { - Some(Box::new(self.expand_recursive(*e)?)) - } else { - None - }; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::If { - cond: Box::new(cond), - then_br: Box::new(then_br), - else_br, - }, - - }) - } - - SyntaxKind::Def { target, value } => { - let target = self.expand_recursive(*target)?; - let value = self.expand_recursive(*value)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Def { - target: Box::new(target), - value: Box::new(value), - }, - - }) - } - - SyntaxKind::Assign { target, value } => { - let target = self.expand_recursive(*target)?; - let value = self.expand_recursive(*value)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Assign { - target: Box::new(target), - value: Box::new(value), - }, - - }) - } - - SyntaxKind::Tuple { elements } => { - let mut expanded_elements = Vec::new(); - for e in elements { - expanded_elements.push(self.expand_recursive(e)?); - } - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Tuple { - elements: expanded_elements, - }, - - }) - } - - SyntaxKind::Record { fields } => { - let mut expanded_fields = Vec::new(); - for (k, v) in fields { - expanded_fields.push((self.expand_recursive(k)?, self.expand_recursive(v)?)); - } - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Record { - fields: expanded_fields, - }, - - }) - } - - SyntaxKind::Expansion { call, expanded } => { - let expanded = self.expand_recursive(*expanded)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Expansion { - call, - expanded: Box::new(expanded), - }, - - }) - } - - SyntaxKind::Again { args } => { - let expanded_args = self.expand_recursive(*args)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Again { - args: Box::new(expanded_args), - }, - - }) - } - - _ => Ok(node), - } - } - - fn expand_call( - &self, - identity: Identity, - name: &str, - params: Vec>, - args: Vec, - body: SyntaxNode, - ) -> Result { - if params.len() != args.len() { - return Err(format!( - "Macro {} expects {} arguments, but got {}", - name, - params.len(), - args.len() - )); - } - - let mut bindings = HashMap::new(); - for (p, a) in params.into_iter().zip(args.into_iter()) { - bindings.insert(p, a); - } - - // AST-Authority: Substitution and Hygiene ONLY happen if there is a Template node. - let expanded_body = if let SyntaxKind::Template(inner) = body.kind { - let mut state = ExpansionState { - bindings: &bindings, - expansion_id: identity.clone(), - }; - self.expand_template(*inner, &mut state)? - } else { - // Static AST fragment: No substitution, no hygiene. - body - }; - - let original_call = SyntaxNode { - identity: identity.clone(), - kind: SyntaxKind::Call { - callee: Box::new(SyntaxNode { - identity: identity.clone(), - kind: SyntaxKind::Identifier(Symbol::from(name)), - - }), - args: Box::new(SyntaxNode { - identity: identity.clone(), - kind: SyntaxKind::Tuple { - elements: bindings.into_values().collect(), - }, - - }), - }, - - }; - - Ok(SyntaxNode { - identity, - kind: SyntaxKind::Expansion { - call: Box::new(original_call), - expanded: Box::new(expanded_body), - }, - - }) - } - - fn extract_param_names(&self, node: &SyntaxNode) -> Result>, String> { - match &node.kind { - SyntaxKind::Identifier(sym) => Ok(vec![sym.name.clone()]), - SyntaxKind::Tuple { elements } => { - let mut names = Vec::new(); - for el in elements { - names.extend(self.extract_param_names(el)?); - } - Ok(names) - } - _ => Err("Invalid parameter pattern in macro declaration".to_string()), - } - } - - fn expand_template( - &self, - node: SyntaxNode, - state: &mut ExpansionState, - ) -> Result { - match node.kind { - SyntaxKind::Identifier(mut sym) => { - // Inside a template, all internal identifiers are colored. - sym.context = Some(state.expansion_id.clone()); - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Identifier(sym), - - }) - } - - SyntaxKind::Placeholder(inner) => { - // Break out of template for substitution/evaluation - if let SyntaxKind::Identifier(ref sym) = inner.kind - && let Some(arg) = state.bindings.get(&sym.name) - { - return Ok(arg.clone()); - } - let val = self.evaluator.evaluate(&inner, state.bindings)?; - Ok(self.value_to_node(val, node.identity)) - } - - SyntaxKind::Splice(inner) => { - // Splicing is also a form of substitution - if let SyntaxKind::Identifier(ref sym) = inner.kind - && let Some(arg) = state.bindings.get(&sym.name) - { - return Ok(arg.clone()); - } - let val = self.evaluator.evaluate(&inner, state.bindings)?; - Ok(self.value_to_node(val, node.identity)) - } - - SyntaxKind::Def { target, value } => { - let target = self.expand_template(*target, state)?; - let val = self.expand_template(*value, state)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Def { - target: Box::new(target), - value: Box::new(val), - }, - - }) - } - - SyntaxKind::Assign { target, value } => { - let target = self.expand_template(*target, state)?; - let value = self.expand_template(*value, state)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Assign { - target: Box::new(target), - value: Box::new(value), - }, - - }) - } - - SyntaxKind::Tuple { elements } => { - let mut new_elements = Vec::new(); - for e in elements { - if let SyntaxKind::Splice(ref inner) = e.kind { - let val = self.evaluator.evaluate(inner, state.bindings)?; - self.handle_splice_value(val, node.identity.clone(), &mut new_elements)?; - } else { - new_elements.push(self.expand_template(e, state)?); - } - } - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Tuple { - elements: new_elements, - }, - - }) - } - - SyntaxKind::Block { exprs } => { - let mut new_exprs = Vec::new(); - for e in exprs { - if let SyntaxKind::Splice(ref inner) = e.kind { - let val = self.evaluator.evaluate(inner, state.bindings)?; - self.handle_splice_value(val, node.identity.clone(), &mut new_exprs)?; - } else { - new_exprs.push(self.expand_template(e, state)?); - } - } - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Block { exprs: new_exprs }, - - }) - } - - SyntaxKind::Record { fields } => { - let mut new_fields = Vec::new(); - for (k, v) in fields { - new_fields.push(( - self.expand_template(k, state)?, - self.expand_template(v, state)?, - )); - } - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Record { fields: new_fields }, - - }) - } - - SyntaxKind::Call { callee, args } => { - let expanded_callee = self.expand_template(*callee, state)?; - let expanded_args = self.expand_template(*args, state)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Call { - callee: Box::new(expanded_callee), - args: Box::new(expanded_args), - }, - - }) - } - - SyntaxKind::If { - cond, - then_br, - else_br, - } => { - let cond = self.expand_template(*cond, state)?; - let then_br = self.expand_template(*then_br, state)?; - let else_br = if let Some(e) = else_br { - Some(Box::new(self.expand_template(*e, state)?)) - } else { - None - }; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::If { - cond: Box::new(cond), - then_br: Box::new(then_br), - else_br, - }, - - }) - } - - SyntaxKind::Pipe { inputs, lambda } => { - let mut expanded_inputs = Vec::with_capacity(inputs.len()); - for input in inputs { - expanded_inputs.push(self.expand_template(input, state)?); - } - let expanded_lambda = Box::new(self.expand_template(*lambda, state)?); - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Pipe { - inputs: expanded_inputs, - lambda: expanded_lambda, - }, - - }) - } - - SyntaxKind::Lambda { params, body } => { - let expanded_params = self.expand_template(*params, state)?; - let body = self.expand_template((*body).clone(), state)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Lambda { - params: Box::new(expanded_params), - body: Rc::new(body), - }, - - }) - } - - SyntaxKind::Again { args } => { - let expanded_args = self.expand_template(*args, state)?; - Ok(SyntaxNode { - identity: node.identity, - kind: SyntaxKind::Again { - args: Box::new(expanded_args), - }, - - }) - } - - _ => Ok(node), - } - } - - fn handle_splice_value( - &self, - val: Value, - _identity: Identity, - target: &mut Vec, - ) -> Result<(), String> { - match val { - Value::Object(obj) => { - if let Some(node) = obj.as_any().downcast_ref::() { - match &node.kind { - SyntaxKind::Tuple { elements } => { - for e in elements { - target.push(e.clone()); - } - return Ok(()); - } - SyntaxKind::Block { exprs } => { - for e in exprs { - target.push(e.clone()); - } - return Ok(()); - } - _ => {} - } - } - Err(format!( - "Strict Splicing: Expected Tuple or Block node, but found {}", - obj.type_name() - )) - } - _ => Err(format!( - "Strict Splicing: Expected sequence (Tuple/Block Node), but found Value: {}", - val - )), - } - } - - fn value_to_node(&self, val: Value, identity: Identity) -> SyntaxNode { - match val { - Value::Object(obj) => { - if let Some(node) = obj.as_any().downcast_ref::() { - return node.clone(); - } - SyntaxNode { - identity, - kind: SyntaxKind::Constant(Value::Object(obj)), - - } - } - _ => SyntaxNode { - identity, - kind: SyntaxKind::Constant(val), - - }, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::ast::compiler::bound_nodes::Address; - use crate::ast::compiler::Binder; - use crate::ast::parser::Parser; - use crate::ast::types::{Object, Value}; - - struct SimpleEvaluator; - impl MacroEvaluator for SimpleEvaluator { - fn evaluate( - &self, - node: &SyntaxNode, - bindings: &HashMap, SyntaxNode>, - ) -> Result { - if let SyntaxKind::Identifier(ref sym) = node.kind - && let Some(arg) = bindings.get(&sym.name) - { - return Ok(Value::Object(Rc::new(arg.clone()) as Rc)); - } - Err(format!( - "SimpleEvaluator cannot evaluate complex expression: {:?}", - node.kind - )) - } - } - - #[test] - fn test_macro_expansion_hygiene() { - let source = " - (do - (macro m [] `(def y 10)) - (m)) - "; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); - let expanded = expander.expand(syntax).unwrap(); - - if let SyntaxKind::Block { exprs } = &expanded.kind - && let SyntaxKind::Expansion { - expanded: result, - call, - } = &exprs[1].kind - { - if let SyntaxKind::Def { target, .. } = &result.kind { - if let SyntaxKind::Identifier(sym) = &target.kind { - assert_eq!(sym.context, Some(call.identity.clone())); - assert_eq!(sym.name.as_ref(), "y"); - } else { - panic!("Expected Identifier target, got {:?}", target.kind); - } - } else { - panic!("Expected Def result, got {:?}", result.kind); - } - } - } - - #[test] - fn test_macro_expansion_unless() { - let source = " - (do - (macro unless [c b] `(if ~c ... ~b)) - (unless false 42)) - "; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); - let expanded = expander.expand(syntax).unwrap(); - - if let SyntaxKind::Block { exprs } = &expanded.kind - && let SyntaxKind::Expansion { - call: _, - expanded: result, - } = &exprs[1].kind - && let SyntaxKind::If { - cond, - then_br, - else_br, - } = &result.kind - { - if let SyntaxKind::Identifier(sym) = &cond.kind { - assert_eq!(sym.name.as_ref(), "false"); - } else { - panic!("Expected identifier 'false', got {:?}", cond.kind); - } - assert!(matches!(then_br.kind, SyntaxKind::Nop)); - if let Some(eb) = else_br { - if let SyntaxKind::Constant(Value::Int(n)) = &eb.kind { - assert_eq!(*n, 42); - } else { - panic!("Expected 42, got {:?}", eb.kind); - } - } else { - panic!("Else branch missing"); - } - } - } - - #[test] - fn test_macro_expansion_splice() { - let source = " - (do - (macro wrap [items] `[0 ~@items 4]) - (def t (wrap [1 2 3]))) - "; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); - let expanded = expander.expand(syntax).unwrap(); - - if let SyntaxKind::Block { exprs } = &expanded.kind - && let SyntaxKind::Def { value, .. } = &exprs[1].kind - && let SyntaxKind::Expansion { - expanded: result, .. - } = &value.kind - && let SyntaxKind::Tuple { elements } = &result.kind - { - assert_eq!(elements.len(), 5); - if let SyntaxKind::Constant(Value::Int(n)) = &elements[0].kind { - assert_eq!(*n, 0); - } - if let SyntaxKind::Constant(Value::Int(n)) = &elements[4].kind { - assert_eq!(*n, 4); - } - } - } - - #[test] - fn test_repro_macro_global_lookup() { - let source = " - (do - (macro square [x] `(* ~x ~x)) - (square 3)) - "; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); - let expanded = expander.expand(syntax).unwrap(); - - // Convert test globals into a CompilerScope - let mut locals = HashMap::new(); - locals.insert( - Symbol::from("*"), - crate::ast::compiler::binder::LocalInfo { - addr: Address::Local(crate::ast::compiler::bound_nodes::VirtualId(0)), - identity: crate::ast::types::NodeIdentity::new(crate::ast::types::SourceLocation { - line: 0, - col: 0, - }), - _ty: crate::ast::types::StaticType::Any, - purity: crate::ast::types::Purity::Pure, - }, - ); - let initial_scopes = vec![crate::ast::compiler::binder::CompilerScope { locals }]; - - let mut diag = crate::ast::diagnostics::Diagnostics::new(); - // fixed_scope_idx = 0 means scope 0 is frozen (Global) - let result = Binder::bind_root(initial_scopes, 1, 0, &expanded, &mut diag); - assert!( - result.is_ok(), - "Should find global '*' Error: {:?}", - result.err() - ); - } - - #[test] - fn test_repro_macro_hygiene_success() { - let source = " - (do - (def y 1) - (macro m [] `(def y 10)) - (m) - y) - "; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); - let expanded = expander.expand(syntax).unwrap(); - - let mut diag = crate::ast::diagnostics::Diagnostics::new(); - let result = Binder::bind_root(vec![], 0, 0, &expanded, &mut diag); - assert!(result.is_ok(), "Hygiene failed: {:?}", result.err()); - } - - #[test] - fn test_repro_macro_hygiene_clash_fixed() { - let source = " - (do - (def y 1) - (macro m1 [x] `(do (def y 10) (assign y 4))) - (m1 8) - y) - "; - let mut parser = Parser::new(source); - let syntax = parser.parse_expression(); - - let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); - let expanded = expander.expand(syntax).unwrap(); - - let mut diag = crate::ast::diagnostics::Diagnostics::new(); - let result = Binder::bind_root(vec![], 0, 0, &expanded, &mut diag); - assert!( - result.is_ok(), - "Explicit Hygiene with Backticks failed: {:?}", - result.err() - ); - } -} +use crate::ast::nodes::{Symbol, SyntaxKind, SyntaxNode}; +use crate::ast::types::{Identity, Value}; +use std::collections::HashMap; +use std::rc::Rc; + +/// Trait for evaluating expressions during macro expansion. +pub trait MacroEvaluator { + fn evaluate( + &self, + node: &SyntaxNode, + bindings: &HashMap, SyntaxNode>, + ) -> Result; +} + +/// Internal state for template expansion (Hygiene context + Parameter binding) +struct ExpansionState<'a> { + bindings: &'a HashMap, SyntaxNode>, + /// The identity of the current expansion instance, used to "color" internal symbols. + expansion_id: Identity, +} + +type MacroMap = HashMap, (Vec>, SyntaxNode)>; + +/// A registry for macro declarations. +#[derive(Clone)] +pub struct MacroRegistry { + scopes: Vec>, +} + +impl Default for MacroRegistry { + fn default() -> Self { + Self::new() + } +} + +impl MacroRegistry { + pub fn new() -> Self { + Self { + scopes: vec![Rc::new(HashMap::new())], + } + } + + pub fn push(&mut self) { + self.scopes.push(Rc::new(HashMap::new())); + } + + pub fn pop(&mut self) { + if self.scopes.len() > 1 { + self.scopes.pop(); + } + } + + pub fn define(&mut self, name: Rc, params: Vec>, body: SyntaxNode) { + if let Some(current_rc) = self.scopes.last_mut() { + // Copy-on-Write: If the Rc is shared, clone the map before modifying. + let current = Rc::make_mut(current_rc); + current.insert(name, (params, body)); + } + } + + pub fn lookup(&self, name: &str) -> Option<(Vec>, SyntaxNode)> { + for scope in self.scopes.iter().rev() { + if let Some(m) = scope.get(name) { + return Some(m.clone()); + } + } + None + } +} + +pub struct MacroExpander { + registry: MacroRegistry, + evaluator: E, +} + +impl MacroExpander { + pub fn new(registry: MacroRegistry, evaluator: E) -> Self { + Self { + registry, + evaluator, + } + } + + pub fn into_registry(self) -> MacroRegistry { + self.registry + } + + pub fn expand(&mut self, node: SyntaxNode) -> Result { + self.expand_recursive(node) + } + + fn expand_recursive(&mut self, node: SyntaxNode) -> Result { + match node.kind { + SyntaxKind::MacroDecl { name, params, body } => { + let p_names = self.extract_param_names(¶ms)?; + self.registry.define(name.name, p_names, (*body).clone()); + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Nop, + ty: (), + }) + } + + SyntaxKind::Call { callee, args } => { + if let SyntaxKind::Identifier { symbol: ref sym, .. } = callee.kind + && let Some((params, body)) = self.registry.lookup(&sym.name) + { + let expanded_args = self.expand_recursive((*args).clone())?; + + // Extract argument nodes from the expanded tuple + let arg_elements = if let SyntaxKind::Tuple { elements } = expanded_args.kind { + elements.into_iter().map(|e| (*e).clone()).collect() + } else { + vec![expanded_args] + }; + + let expanded = self.expand_call( + node.identity.clone(), + &sym.name, + params, + arg_elements, + body, + )?; + return self.expand_recursive(expanded); + } + + let expanded_callee = self.expand_recursive((*callee).clone())?; + let expanded_args = self.expand_recursive((*args).clone())?; + + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Call { + callee: Rc::new(expanded_callee), + args: Rc::new(expanded_args), + }, + ty: (), + }) + } + + SyntaxKind::Block { exprs } => { + self.registry.push(); + let mut expanded_exprs = Vec::new(); + for expr in exprs { + expanded_exprs.push(Rc::new(self.expand_recursive((*expr).clone())?)); + } + self.registry.pop(); + + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Block { + exprs: expanded_exprs, + }, + ty: (), + }) + } + + SyntaxKind::Pipe { inputs, lambda } => { + let mut expanded_inputs = Vec::with_capacity(inputs.len()); + for input in inputs { + expanded_inputs.push(Rc::new(self.expand_recursive((*input).clone())?)); + } + let expanded_lambda = Rc::new(self.expand_recursive((*lambda).clone())?); + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Pipe { + inputs: expanded_inputs, + lambda: expanded_lambda, + }, + ty: (), + }) + } + + SyntaxKind::Lambda { params, body, .. } => { + self.registry.push(); + let expanded_params = self.expand_recursive((*params).clone())?; + let expanded_body = self.expand_recursive((*body).clone())?; + self.registry.pop(); + + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Lambda { + params: Rc::new(expanded_params), + body: Rc::new(expanded_body), + info: (), + }, + ty: (), + }) + } + + SyntaxKind::If { + cond, + then_br, + else_br, + } => { + let cond = self.expand_recursive((*cond).clone())?; + let then_br = self.expand_recursive((*then_br).clone())?; + let else_br = if let Some(e) = else_br { + Some(Rc::new(self.expand_recursive((*e).clone())?)) + } else { + None + }; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::If { + cond: Rc::new(cond), + then_br: Rc::new(then_br), + else_br, + }, + ty: (), + }) + } + + SyntaxKind::Def { pattern, value, .. } => { + let pattern = self.expand_recursive((*pattern).clone())?; + let value = self.expand_recursive((*value).clone())?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Def { + pattern: Rc::new(pattern), + value: Rc::new(value), + info: (), + }, + ty: (), + }) + } + + SyntaxKind::Assign { target, value, .. } => { + let target = self.expand_recursive((*target).clone())?; + let value = self.expand_recursive((*value).clone())?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Assign { + target: Rc::new(target), + value: Rc::new(value), + info: (), + }, + ty: (), + }) + } + + SyntaxKind::Tuple { elements } => { + let mut expanded_elements = Vec::new(); + for e in elements { + expanded_elements.push(Rc::new(self.expand_recursive((*e).clone())?)); + } + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Tuple { + elements: expanded_elements, + }, + ty: (), + }) + } + + SyntaxKind::Record { fields, .. } => { + let mut expanded_fields = Vec::new(); + for (k, v) in fields { + expanded_fields.push(( + Rc::new(self.expand_recursive((*k).clone())?), + Rc::new(self.expand_recursive((*v).clone())?), + )); + } + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Record { + fields: expanded_fields, + layout: (), + }, + ty: (), + }) + } + + SyntaxKind::Expansion { original_call, expanded } => { + let expanded = self.expand_recursive((*expanded).clone())?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Expansion { + original_call, + expanded: Rc::new(expanded), + }, + ty: (), + }) + } + + SyntaxKind::Again { args } => { + let expanded_args = self.expand_recursive((*args).clone())?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Again { + args: Rc::new(expanded_args), + }, + ty: (), + }) + } + + _ => Ok(node), + } + } + + fn expand_call( + &self, + identity: Identity, + name: &str, + params: Vec>, + args: Vec, + body: SyntaxNode, + ) -> Result { + if params.len() != args.len() { + return Err(format!( + "Macro {} expects {} arguments, but got {}", + name, + params.len(), + args.len() + )); + } + + let mut bindings = HashMap::new(); + for (p, a) in params.into_iter().zip(args.into_iter()) { + bindings.insert(p, a); + } + + // AST-Authority: Substitution and Hygiene ONLY happen if there is a Template node. + let expanded_body = if let SyntaxKind::Template(inner) = body.kind { + let mut state = ExpansionState { + bindings: &bindings, + expansion_id: identity.clone(), + }; + self.expand_template((*inner).clone(), &mut state)? + } else { + // Static AST fragment: No substitution, no hygiene. + body + }; + + let original_call = SyntaxNode { + identity: identity.clone(), + kind: SyntaxKind::Call { + callee: Rc::new(SyntaxNode { + identity: identity.clone(), + kind: SyntaxKind::Identifier { + symbol: Symbol::from(name), + binding: (), + }, + ty: (), + }), + args: Rc::new(SyntaxNode { + identity: identity.clone(), + kind: SyntaxKind::Tuple { + elements: bindings.into_values().map(Rc::new).collect(), + }, + ty: (), + }), + }, + ty: (), + }; + + Ok(SyntaxNode { + identity, + kind: SyntaxKind::Expansion { + original_call: Rc::new(original_call), + expanded: Rc::new(expanded_body), + }, + ty: (), + }) + } + + fn extract_param_names(&self, node: &SyntaxNode) -> Result>, String> { + match &node.kind { + SyntaxKind::Identifier { symbol: sym, .. } => Ok(vec![sym.name.clone()]), + SyntaxKind::Tuple { elements } => { + let mut names = Vec::new(); + for el in elements { + names.extend(self.extract_param_names(el)?); + } + Ok(names) + } + _ => Err("Invalid parameter pattern in macro declaration".to_string()), + } + } + + fn expand_template( + &self, + node: SyntaxNode, + state: &mut ExpansionState, + ) -> Result { + match node.kind { + SyntaxKind::Identifier { mut symbol, .. } => { + // Inside a template, all internal identifiers are colored. + symbol.context = Some(state.expansion_id.clone()); + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Identifier { + symbol, + binding: (), + }, + ty: (), + }) + } + + SyntaxKind::Placeholder(inner) => { + // Break out of template for substitution/evaluation + if let SyntaxKind::Identifier { symbol: ref sym, .. } = inner.kind + && let Some(arg) = state.bindings.get(&sym.name) + { + return Ok(arg.clone()); + } + let val = self.evaluator.evaluate(&inner, state.bindings)?; + Ok(self.value_to_node(val, node.identity)) + } + + SyntaxKind::Splice(inner) => { + // Splicing is also a form of substitution + if let SyntaxKind::Identifier { symbol: ref sym, .. } = inner.kind + && let Some(arg) = state.bindings.get(&sym.name) + { + return Ok(arg.clone()); + } + let val = self.evaluator.evaluate(&inner, state.bindings)?; + Ok(self.value_to_node(val, node.identity)) + } + + SyntaxKind::Def { pattern, value, .. } => { + let pattern = self.expand_template((*pattern).clone(), state)?; + let val = self.expand_template((*value).clone(), state)?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Def { + pattern: Rc::new(pattern), + value: Rc::new(val), + info: (), + }, + ty: (), + }) + } + + SyntaxKind::Assign { target, value, .. } => { + let target = self.expand_template((*target).clone(), state)?; + let value = self.expand_template((*value).clone(), state)?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Assign { + target: Rc::new(target), + value: Rc::new(value), + info: (), + }, + ty: (), + }) + } + + SyntaxKind::Tuple { elements } => { + let mut new_elements = Vec::new(); + for e in elements { + if let SyntaxKind::Splice(ref inner) = e.kind { + let val = self.evaluator.evaluate(inner, state.bindings)?; + self.handle_splice_value(val, node.identity.clone(), &mut new_elements)?; + } else { + new_elements.push(Rc::new(self.expand_template((*e).clone(), state)?)); + } + } + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Tuple { + elements: new_elements, + }, + ty: (), + }) + } + + SyntaxKind::Block { exprs } => { + let mut new_exprs = Vec::new(); + for e in exprs { + if let SyntaxKind::Splice(ref inner) = e.kind { + let val = self.evaluator.evaluate(inner, state.bindings)?; + self.handle_splice_value(val, node.identity.clone(), &mut new_exprs)?; + } else { + new_exprs.push(Rc::new(self.expand_template((*e).clone(), state)?)); + } + } + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Block { exprs: new_exprs }, + ty: (), + }) + } + + SyntaxKind::Record { fields, .. } => { + let mut new_fields = Vec::new(); + for (k, v) in fields { + new_fields.push(( + Rc::new(self.expand_template((*k).clone(), state)?), + Rc::new(self.expand_template((*v).clone(), state)?), + )); + } + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Record { + fields: new_fields, + layout: (), + }, + ty: (), + }) + } + + SyntaxKind::Call { callee, args } => { + let expanded_callee = self.expand_template((*callee).clone(), state)?; + let expanded_args = self.expand_template((*args).clone(), state)?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Call { + callee: Rc::new(expanded_callee), + args: Rc::new(expanded_args), + }, + ty: (), + }) + } + + SyntaxKind::If { + cond, + then_br, + else_br, + } => { + let cond = self.expand_template((*cond).clone(), state)?; + let then_br = self.expand_template((*then_br).clone(), state)?; + let else_br = if let Some(e) = else_br { + Some(Rc::new(self.expand_template((*e).clone(), state)?)) + } else { + None + }; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::If { + cond: Rc::new(cond), + then_br: Rc::new(then_br), + else_br, + }, + ty: (), + }) + } + + SyntaxKind::Pipe { inputs, lambda } => { + let mut expanded_inputs = Vec::with_capacity(inputs.len()); + for input in inputs { + expanded_inputs.push(Rc::new(self.expand_template((*input).clone(), state)?)); + } + let expanded_lambda = Rc::new(self.expand_template((*lambda).clone(), state)?); + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Pipe { + inputs: expanded_inputs, + lambda: expanded_lambda, + }, + ty: (), + }) + } + + SyntaxKind::Lambda { params, body, .. } => { + let expanded_params = self.expand_template((*params).clone(), state)?; + let body = self.expand_template((*body).clone(), state)?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Lambda { + params: Rc::new(expanded_params), + body: Rc::new(body), + info: (), + }, + ty: (), + }) + } + + SyntaxKind::Again { args } => { + let expanded_args = self.expand_template((*args).clone(), state)?; + Ok(SyntaxNode { + identity: node.identity, + kind: SyntaxKind::Again { + args: Rc::new(expanded_args), + }, + ty: (), + }) + } + + _ => Ok(node), + } + } + + fn handle_splice_value( + &self, + val: Value, + _identity: Identity, + target: &mut Vec>, + ) -> Result<(), String> { + match val { + Value::Object(obj) => { + if let Some(node) = obj.as_any().downcast_ref::() { + match &node.kind { + SyntaxKind::Tuple { elements } => { + for e in elements { + target.push(e.clone()); + } + return Ok(()); + } + SyntaxKind::Block { exprs } => { + for e in exprs { + target.push(e.clone()); + } + return Ok(()); + } + _ => {} + } + } + Err(format!( + "Strict Splicing: Expected Tuple or Block node, but found {}", + obj.type_name() + )) + } + _ => Err(format!( + "Strict Splicing: Expected sequence (Tuple/Block Node), but found Value: {}", + val + )), + } + } + + fn value_to_node(&self, val: Value, identity: Identity) -> SyntaxNode { + match val { + Value::Object(obj) => { + if let Some(node) = obj.as_any().downcast_ref::() { + return node.clone(); + } + SyntaxNode { + identity, + kind: SyntaxKind::Constant(Value::Object(obj)), + ty: (), + } + } + _ => SyntaxNode { + identity, + kind: SyntaxKind::Constant(val), + ty: (), + }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::compiler::bound_nodes::Address; + use crate::ast::compiler::Binder; + use crate::ast::parser::Parser; + use crate::ast::types::{Object, Value}; + + struct SimpleEvaluator; + impl MacroEvaluator for SimpleEvaluator { + fn evaluate( + &self, + node: &SyntaxNode, + bindings: &HashMap, SyntaxNode>, + ) -> Result { + if let SyntaxKind::Identifier { symbol: ref sym, .. } = node.kind + && let Some(arg) = bindings.get(&sym.name) + { + return Ok(Value::Object(Rc::new(arg.clone()) as Rc)); + } + Err(format!( + "SimpleEvaluator cannot evaluate complex expression: {:?}", + node.kind + )) + } + } + + #[test] + fn test_macro_expansion_hygiene() { + let source = " + (do + (macro m [] `(def y 10)) + (m)) + "; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); + let expanded = expander.expand(syntax).unwrap(); + + if let SyntaxKind::Block { exprs } = &expanded.kind + && let SyntaxKind::Expansion { + expanded: result, + original_call, + } = &exprs[1].kind + { + if let SyntaxKind::Def { pattern, .. } = &result.kind { + if let SyntaxKind::Identifier { symbol: sym, .. } = &pattern.kind { + assert_eq!(sym.context, Some(original_call.identity.clone())); + assert_eq!(sym.name.as_ref(), "y"); + } else { + panic!("Expected Identifier target, got {:?}", pattern.kind); + } + } else { + panic!("Expected Def result, got {:?}", result.kind); + } + } + } + + #[test] + fn test_macro_expansion_unless() { + let source = " + (do + (macro unless [c b] `(if ~c ... ~b)) + (unless false 42)) + "; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); + let expanded = expander.expand(syntax).unwrap(); + + if let SyntaxKind::Block { exprs } = &expanded.kind + && let SyntaxKind::Expansion { + original_call: _, + expanded: result, + } = &exprs[1].kind + && let SyntaxKind::If { + cond, + then_br, + else_br, + } = &result.kind + { + if let SyntaxKind::Identifier { symbol: sym, .. } = &cond.kind { + assert_eq!(sym.name.as_ref(), "false"); + } else { + panic!("Expected identifier 'false', got {:?}", cond.kind); + } + assert!(matches!(then_br.kind, SyntaxKind::Nop)); + if let Some(eb) = else_br { + if let SyntaxKind::Constant(Value::Int(n)) = &eb.kind { + assert_eq!(*n, 42); + } else { + panic!("Expected 42, got {:?}", eb.kind); + } + } else { + panic!("Else branch missing"); + } + } + } + + #[test] + fn test_macro_expansion_splice() { + let source = " + (do + (macro wrap [items] `[0 ~@items 4]) + (def t (wrap [1 2 3]))) + "; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); + let expanded = expander.expand(syntax).unwrap(); + + if let SyntaxKind::Block { exprs } = &expanded.kind + && let SyntaxKind::Def { value, .. } = &exprs[1].kind + && let SyntaxKind::Expansion { + expanded: result, .. + } = &value.kind + && let SyntaxKind::Tuple { elements } = &result.kind + { + assert_eq!(elements.len(), 5); + if let SyntaxKind::Constant(Value::Int(n)) = &elements[0].kind { + assert_eq!(*n, 0); + } + if let SyntaxKind::Constant(Value::Int(n)) = &elements[4].kind { + assert_eq!(*n, 4); + } + } + } + + #[test] + fn test_repro_macro_global_lookup() { + let source = " + (do + (macro square [x] `(* ~x ~x)) + (square 3)) + "; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); + let expanded = expander.expand(syntax).unwrap(); + + // Convert test globals into a CompilerScope + let mut locals = HashMap::new(); + locals.insert( + Symbol::from("*"), + crate::ast::compiler::binder::LocalInfo { + addr: Address::Local(crate::ast::compiler::bound_nodes::VirtualId(0)), + identity: crate::ast::types::NodeIdentity::new(crate::ast::types::SourceLocation { + line: 0, + col: 0, + }), + _ty: crate::ast::types::StaticType::Any, + purity: crate::ast::types::Purity::Pure, + }, + ); + let initial_scopes = vec![crate::ast::compiler::binder::CompilerScope { locals }]; + + let mut diag = crate::ast::diagnostics::Diagnostics::new(); + // fixed_scope_idx = 0 means scope 0 is frozen (Global) + let result = Binder::bind_root(initial_scopes, 1, 0, &expanded, &mut diag); + assert!( + result.is_ok(), + "Should find global '*' Error: {:?}", + result.err() + ); + } + + #[test] + fn test_repro_macro_hygiene_success() { + let source = " + (do + (def y 1) + (macro m [] `(def y 10)) + (m) + y) + "; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); + let expanded = expander.expand(syntax).unwrap(); + + let mut diag = crate::ast::diagnostics::Diagnostics::new(); + let result = Binder::bind_root(vec![], 0, 0, &expanded, &mut diag); + assert!(result.is_ok(), "Hygiene failed: {:?}", result.err()); + } + + #[test] + fn test_repro_macro_hygiene_clash_fixed() { + let source = " + (do + (def y 1) + (macro m1 [x] `(do (def y 10) (assign y 4))) + (m1 8) + y) + "; + let mut parser = Parser::new(source); + let syntax = parser.parse_expression(); + + let mut expander = MacroExpander::new(MacroRegistry::new(), SimpleEvaluator); + let expanded = expander.expand(syntax).unwrap(); + + let mut diag = crate::ast::diagnostics::Diagnostics::new(); + let result = Binder::bind_root(vec![], 0, 0, &expanded, &mut diag); + assert!( + result.is_ok(), + "Explicit Hygiene with Backticks failed: {:?}", + result.err() + ); + } +} diff --git a/src/ast/compiler/optimizer/engine.rs b/src/ast/compiler/optimizer/engine.rs index 8a2e522..96c760c 100644 --- a/src/ast/compiler/optimizer/engine.rs +++ b/src/ast/compiler/optimizer/engine.rs @@ -1,740 +1,812 @@ -use crate::ast::compiler::bound_nodes::{ - Address, AnalyzedNode, BoundKind, GlobalAnalyzedRegistry, Node, UpvalueIdx, -}; -use crate::ast::types::{Purity, Value}; -use crate::ast::vm::Closure; -use std::cell::RefCell; -use std::rc::Rc; - -use super::folder::Folder; -use super::inliner::Inliner; -use super::substitution_map::SubstitutionMap; -use super::utils::{PathTracker, UsageInfo}; - -pub struct Optimizer { - pub enabled: bool, - max_passes: usize, - pub globals: Option>>>, - pub root_purity: Option>>>, - pub lambda_registry: Option>>, -} - -impl Optimizer { - pub fn new(enabled: bool) -> Self { - Self { - enabled, - max_passes: 5, - globals: None, - root_purity: None, - lambda_registry: None, - } - } - - pub fn with_globals(mut self, globals: Rc>>) -> Self { - self.globals = Some(globals); - self - } - - pub fn with_purity(mut self, purity: Rc>>) -> Self { - self.root_purity = Some(purity); - self - } - - pub fn with_registry( - mut self, - registry: Rc>, - ) -> Self { - self.lambda_registry = Some(registry); - self - } - - pub fn optimize(&self, node: AnalyzedNode) -> AnalyzedNode { - if !self.enabled { - return node; - } - - let mut current = Rc::new(node); - for _ in 0..self.max_passes { - let mut sub = SubstitutionMap::new(); - let mut path = PathTracker::new(); - let next = self.visit_node(current.clone(), &mut sub, &mut path); - if Rc::ptr_eq(&next, ¤t) { - break; - } - current = next; - } - // Unwrap the final Rc if we are at the end, or return a clone of the inner node. - // Since AnalyzedNode is small now (header + Rcs), cloning is cheap. - (*current).clone() - } - - fn try_inline( - &self, - params: &AnalyzedNode, - arg_nodes: &[Rc], - body: &AnalyzedNode, - sub: &mut SubstitutionMap, - path: &mut PathTracker, - base_sub: Option, - ) -> Option> { - let inliner = Inliner::new(&self.globals, &self.root_purity); - let mut inner_sub = base_sub.unwrap_or_else(|| sub.new_for_inlining()); - - if inliner.prepare_beta_reduction(params, arg_nodes, body, &mut inner_sub).is_some() { - let res = self.visit_node(Rc::new(body.clone()), &mut inner_sub, path); - - // Sync back state to parent substitution map - sub.next_slot = inner_sub.next_slot; - sub.used.extend(inner_sub.used.iter().cloned()); - sub.assigned.extend(inner_sub.assigned.iter().cloned()); - sub.captured_slots.extend(inner_sub.captured_slots.iter().cloned()); - - Some(res) - } else { - None - } - } - - fn visit_node( - &self, - node_rc: Rc, - sub: &mut SubstitutionMap, - path: &mut PathTracker, - ) -> Rc { - let node = &*node_rc; - let folder = Folder::new(&self.globals); - let inliner = Inliner::new(&self.globals, &self.root_purity); - - let (new_kind, metrics) = match &node.kind { - BoundKind::Get { addr, name } => { - let addr = *addr; - if !sub.assigned.contains(&addr) { - // 1. Try inlining from current value substitution map (locals/globals/upvalues) - if let Some(val) = sub.get_value(&addr) - && inliner.is_inlinable_value(val, addr) - { - return Rc::new(folder.make_constant_node(val.clone(), node)); - } - - // 2. Try inlining from AST substitution map (pure expressions) - if let Some(inlined_node) = sub.ast_substitutions.get(&addr) { - return inlined_node.clone(); - } - - // 3. Fallback for Globals: check the actual VM environment - if let Address::Global(idx) = addr - && let Some(globals_rc) = &self.globals - { - let globals = globals_rc.borrow(); - if let Some(val) = globals.get(idx.0 as usize) - && inliner.is_inlinable_value(val, addr) - { - return Rc::new(folder.make_constant_node(val.clone(), node)); - } - } - } - - sub.used.insert(addr); - ( - BoundKind::Get { - addr: sub.map_address(addr), - name: name.clone(), - }, - node.ty.clone(), - ) - } - - BoundKind::FieldAccessor(k) => (BoundKind::FieldAccessor(*k), node.ty.clone()), - - BoundKind::GetField { rec, field } => { - let rec_opt = self.visit_node(rec.clone(), sub, path); - - // Constant folding for Field Access - if let BoundKind::Constant(Value::Record(layout, values)) = &rec_opt.kind - && let Some(idx) = layout.index_of(*field) - { - return Rc::new(folder.make_constant_node(values[idx].clone(), node)); - } - - if Rc::ptr_eq(&rec_opt, rec) { - return node_rc; - } - - ( - BoundKind::GetField { - rec: rec_opt, - field: *field, - }, - node.ty.clone(), - ) - } - - BoundKind::Set { addr, value } => { - let value_opt = self.visit_node(value.clone(), sub, path); - if let BoundKind::Constant(val) = &value_opt.kind { - sub.add_value(*addr, val.clone()); - } else { - sub.remove_value(addr); - } - - if Rc::ptr_eq(&value_opt, value) { - return node_rc; - } - - ( - BoundKind::Set { - addr: sub.map_address(*addr), - value: value_opt, - }, - node.ty.clone(), - ) - } - - BoundKind::Define { - name, - addr, - kind, - value, - captured_by, - } => { - let value_opt = self.visit_node(value.clone(), sub, path); - - if let Address::Local(slot) = addr - && !captured_by.is_empty() - { - sub.captured_slots.insert(*slot); - } - - if let BoundKind::Constant(val) = &value_opt.kind { - sub.add_value(*addr, val.clone()); - } else { - sub.remove_value(addr); - } - - if let Address::Global(global_index) = addr - && value_opt.ty.purity > Purity::Impure - && let Some(purity_rc) = &self.root_purity - { - let mut pr = purity_rc.borrow_mut(); - let idx = global_index.0 as usize; - if idx < pr.len() { - pr[idx] = value_opt.ty.purity; - } - } - - if Rc::ptr_eq(&value_opt, value) { - return node_rc; - } - - ( - BoundKind::Define { - name: name.clone(), - addr: sub.map_address(*addr), - kind: *kind, - value: value_opt, - captured_by: captured_by.clone(), - }, - node.ty.clone(), - ) - } - - BoundKind::Call { callee, args } => { - let callee_opt = self.visit_node(callee.clone(), sub, path); - let args_opt = self.visit_node(args.clone(), sub, path); - - if self.enabled { - // Optimized Field Access Transformation - if let BoundKind::FieldAccessor(k) = &callee_opt.kind - && let BoundKind::Tuple { elements } = &args_opt.kind - && elements.len() == 1 - { - let rec = elements[0].clone(); - let metrics = node.ty.clone(); - return Rc::new(Node { - identity: node.identity.clone(), - kind: BoundKind::GetField { - rec, - field: *k, - }, - ty: metrics, - }); - } - - let mut arg_nodes = Vec::new(); - self.flatten_tuple(args_opt.clone(), &mut arg_nodes); - - if let BoundKind::Lambda { - params, - body, - upvalues, - positional_count, - } = &callee_opt.kind - && upvalues.is_empty() - && positional_count.is_some() - && path.inlining_depth < 5 - && !callee_opt.ty.is_recursive - && path.enter_lambda(&callee_opt.identity) - { - path.inlining_depth += 1; - let collapsed = self.try_inline(params, &arg_nodes, body, sub, path, None); - path.inlining_depth -= 1; - path.exit_lambda(&callee_opt.identity); - - if let Some(res) = collapsed { - return res; - } - } - - if let BoundKind::Get { - addr: Address::Global(idx), - .. - } = &callee_opt.kind - && let Some(registry_rc) = &self.lambda_registry - && path.inlining_depth < 5 - && !path.inlining_stack.contains(idx) - { - let registry = registry_rc.borrow(); - if let Some(lambda_node) = registry.get(idx) - && let BoundKind::Lambda { - params, - body, - upvalues, - positional_count, - } = &lambda_node.kind - && upvalues.is_empty() - && positional_count.is_some() - && !lambda_node.ty.is_recursive - { - path.inlining_stack.insert(*idx); - path.inlining_depth += 1; - let collapsed = self.try_inline(params, &arg_nodes, body, sub, path, None); - path.inlining_depth -= 1; - path.inlining_stack.remove(idx); - - if let Some(res) = collapsed { - return res; - } - } - } - - if let BoundKind::Constant(Value::Object(ref obj)) = callee_opt.kind - && path.inlining_depth < 5 - && let Some(closure) = obj.as_any().downcast_ref::() - && (closure.upvalues.is_empty() - || closure.function_node.ty.purity >= Purity::SideEffectFree) - && !closure.function_node.ty.is_recursive - && path.enter_lambda(&closure.function_node.identity) - { - let mut closure_sub = sub.new_for_inlining(); - for (i, cell) in closure.upvalues.iter().enumerate() { - closure_sub.add_value( - Address::Upvalue(UpvalueIdx(i as u32)), - cell.borrow().clone(), - ); - } - - path.inlining_depth += 1; - if let BoundKind::Lambda { params, .. } = &closure.function_node.kind { - let collapsed = self.try_inline( - params, - &arg_nodes, - &closure.function_node, - sub, - path, - Some(closure_sub), - ); - path.inlining_depth -= 1; - path.exit_lambda(&closure.function_node.identity); - - if let Some(res) = collapsed { - return res; - } - } else { - path.inlining_depth -= 1; - path.exit_lambda(&closure.function_node.identity); - } - } - - if let Some(folded) = folder.try_fold_pure(&callee_opt, &arg_nodes) { - return Rc::new(folded); - } - } - - if Rc::ptr_eq(&callee_opt, callee) && Rc::ptr_eq(&args_opt, args) { - return node_rc; - } - - ( - BoundKind::Call { - callee: callee_opt, - args: args_opt, - }, - node.ty.clone(), - ) - } - - BoundKind::Again { args } => { - let args_opt = self.visit_node(args.clone(), sub, path); - if Rc::ptr_eq(&args_opt, args) { - return node_rc; - } - ( - BoundKind::Again { - args: args_opt, - }, - node.ty.clone(), - ) - } - - BoundKind::If { - cond, - then_br, - else_br, - } => { - let cond_opt = self.visit_node(cond.clone(), sub, path); - if self.enabled - && let BoundKind::Constant(ref val) = cond_opt.kind - { - if val.is_truthy() { - return self.visit_node(then_br.clone(), sub, path); - } else if let Some(else_node) = else_br { - return self.visit_node(else_node.clone(), sub, path); - } else { - return Rc::new(folder.make_nop_node(node)); - } - } - - let then_br_opt = self.visit_node(then_br.clone(), sub, path); - let else_br_opt = else_br - .as_ref() - .map(|e| self.visit_node(e.clone(), sub, path)); - - let unchanged = Rc::ptr_eq(&cond_opt, cond) && Rc::ptr_eq(&then_br_opt, then_br) && match (&else_br_opt, else_br) { - (Some(a), Some(b)) => Rc::ptr_eq(a, b), - (None, None) => true, - _ => false, - }; - - if unchanged { - return node_rc; - } - - ( - BoundKind::If { - cond: cond_opt, - then_br: then_br_opt, - else_br: else_br_opt, - }, - node.ty.clone(), - ) - } - - BoundKind::Pipe { - inputs, - lambda, - out_type, - } => { - let mut o_inputs = Vec::with_capacity(inputs.len()); - let mut inputs_changed = false; - for input in inputs { - let opt = self.visit_node(input.clone(), sub, path); - if !Rc::ptr_eq(&opt, input) { - inputs_changed = true; - } - o_inputs.push(opt); - } - let o_lambda = self.visit_node(lambda.clone(), sub, path); - - if !inputs_changed && Rc::ptr_eq(&o_lambda, lambda) { - return node_rc; - } - - ( - BoundKind::Pipe { - inputs: o_inputs, - lambda: o_lambda, - out_type: out_type.clone(), - }, - node.ty.clone(), - ) - } - BoundKind::Block { exprs } => { - let mut info = UsageInfo::default(); - if !exprs.is_empty() { - for e in exprs { - info.collect(e); - } - } - - sub.assigned.extend(info.assigned.iter().cloned()); - - let mut new_exprs = Vec::with_capacity(exprs.len()); - let last_idx = exprs.len().saturating_sub(1); - let mut block_changed = false; - - for (i, e) in exprs.iter().enumerate() { - let is_last = i == last_idx; - if self.enabled && !is_last { - let removable = match &e.kind { - BoundKind::Define { addr, value, .. } => { - !info.is_used(addr) - && (if let Address::Local(slot) = addr { - !sub.captured_slots.contains(slot) - } else { - true - }) - && (value.ty.purity >= Purity::SideEffectFree - || matches!(value.kind, BoundKind::Lambda { .. })) - } - BoundKind::Set { addr, value, .. } => { - !info.is_used(addr) - && (if let Address::Local(slot) = addr { - !sub.captured_slots.contains(slot) - } else { - true - }) - && value.ty.purity >= Purity::SideEffectFree - } - _ => e.ty.purity >= Purity::SideEffectFree, - }; - - if removable { - block_changed = true; - continue; - } - } - - let unmapped_addr = if let BoundKind::Define { addr, .. } = &e.kind { - Some(*addr) - } else { - None - }; - - let opt = self.visit_node(e.clone(), sub, path); - - if let BoundKind::Define { value, .. } = &opt.kind - && let Some(orig_addr) = unmapped_addr - && !info.assigned.contains(&orig_addr) - { - let mut core_value = value.as_ref(); - while let BoundKind::Expansion { bound_expanded, .. } = &core_value.kind { - core_value = bound_expanded.as_ref(); - } - - if let BoundKind::Constant(val) = &core_value.kind { - sub.add_value(orig_addr, val.clone()); - } else if let BoundKind::Lambda { upvalues, .. } = &core_value.kind - && upvalues.is_empty() - { - sub.add_ast_substitution(orig_addr, core_value.clone()); - } - } - - if self.enabled && matches!(opt.kind, BoundKind::Nop) && !is_last { - block_changed = true; - continue; - } - if !Rc::ptr_eq(&opt, e) { - block_changed = true; - } - new_exprs.push(opt); - } - - if !block_changed { - return node_rc; - } - - if self.enabled { - if new_exprs.is_empty() { - return Rc::new(folder.make_nop_node(node)); - } else if new_exprs.len() == 1 { - return new_exprs.pop().unwrap(); - } - } - - (BoundKind::Block { exprs: new_exprs }, node.ty.clone()) - } - - BoundKind::Lambda { - params, - upvalues, - body, - positional_count, - } => { - let mut info = UsageInfo::default(); - info.collect(node); - - let mut new_upvalues = Vec::new(); - let mut mapping = Vec::new(); - let mut next_inner_subs = sub.new_inner(); - next_inner_subs.assigned = info.assigned; - let mut upvalues_changed = false; - - for (old_idx, capture_addr) in upvalues.iter().enumerate() { - let mut inlined_val = None; - let mut inlined_ast = None; - if !sub.assigned.contains(capture_addr) { - if let Some(val) = sub.get_value(capture_addr) { - inlined_val = Some(val.clone()); - } else if let Some(ast) = sub.ast_substitutions.get(capture_addr) { - inlined_ast = Some(Rc::clone(ast)); - } - } - - if let Address::Local(slot) = capture_addr { - sub.captured_slots.insert(*slot); - } - - if let Some(val) = inlined_val { - next_inner_subs - .add_value(Address::Upvalue(UpvalueIdx(old_idx as u32)), val); - mapping.push(None); - upvalues_changed = true; - } else if let Some(ast) = inlined_ast { - next_inner_subs - .add_ast_substitution(Address::Upvalue(UpvalueIdx(old_idx as u32)), (*ast).clone()); - mapping.push(None); - upvalues_changed = true; - } else { - mapping.push(Some(new_upvalues.len() as u32)); - let mapped = sub.map_address(*capture_addr); - if mapped != *capture_addr { - upvalues_changed = true; - } - new_upvalues.push(mapped); - } - } - - let params_opt = self.visit_node(params.clone(), &mut next_inner_subs, path); - - let body_opt = self.visit_node(body.clone(), &mut next_inner_subs, path); - let reindexed_body = if new_upvalues.len() != upvalues.len() { - sub.reindex_upvalues(body_opt, &mapping) - } else { - body_opt - }; - - if !upvalues_changed && Rc::ptr_eq(¶ms_opt, params) && Rc::ptr_eq(&reindexed_body, body) { - return node_rc; - } - - ( - BoundKind::Lambda { - params: params_opt, - upvalues: new_upvalues, - body: reindexed_body, - positional_count: *positional_count, - }, - node.ty.clone(), - ) - } - - BoundKind::Destructure { pattern, value } => { - let val_opt = self.visit_node(value.clone(), sub, path); - let pat_opt = self.visit_node(pattern.clone(), sub, path); - - let mut info = UsageInfo::default(); - info.collect_pattern(&pat_opt); - for addr in info.assigned { - sub.remove_value(&addr); - } - - if Rc::ptr_eq(&val_opt, value) && Rc::ptr_eq(&pat_opt, pattern) { - return node_rc; - } - - ( - BoundKind::Destructure { - pattern: pat_opt, - value: val_opt, - }, - node.ty.clone(), - ) - } - - BoundKind::Tuple { elements } => { - let mut new_elements = Vec::with_capacity(elements.len()); - let mut changed = false; - for e in elements { - let opt = self.visit_node(e.clone(), sub, path); - if !Rc::ptr_eq(&opt, e) { - changed = true; - } - new_elements.push(opt); - } - if !changed { - return node_rc; - } - (BoundKind::Tuple { elements: new_elements }, node.ty.clone()) - } - BoundKind::Record { layout, values } => { - let mut mapped_values = Vec::with_capacity(values.len()); - let mut changed = false; - for v in values { - let opt = self.visit_node(v.clone(), sub, path); - if !Rc::ptr_eq(&opt, v) { - changed = true; - } - mapped_values.push(opt); - } - - if self.enabled - && let Some(folded) = folder.try_fold_record(layout, &mapped_values, node) - { - return Rc::new(folded); - } - - if !changed { - return node_rc; - } - - ( - BoundKind::Record { - layout: layout.clone(), - values: mapped_values, - }, - node.ty.clone(), - ) - } - BoundKind::Expansion { - original_call, - bound_expanded, - } => { - path.inlining_depth += 1; - let expanded_opt = self.visit_node(bound_expanded.clone(), sub, path); - path.inlining_depth -= 1; - - if Rc::ptr_eq(&expanded_opt, bound_expanded) { - return node_rc; - } - - ( - BoundKind::Expansion { - original_call: original_call.clone(), - bound_expanded: expanded_opt, - }, - node.ty.clone(), - ) - } - k => (k.clone(), node.ty.clone()), - }; - - Rc::new(Node { - identity: node.identity.clone(), - kind: new_kind, - ty: metrics, - }) - } - - fn flatten_tuple(&self, node: Rc, into: &mut Vec>) { - match &node.kind { - BoundKind::Tuple { elements } => { - for el in elements { - self.flatten_tuple(el.clone(), into); - } - } - BoundKind::Nop => {} - _ => into.push(node), - } - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AnalyzedNode, AssignBinding, GlobalAnalyzedRegistry, + IdentifierBinding, LambdaBinding, Node, NodeKind, UpvalueIdx, +}; +use crate::ast::types::{Purity, Value}; +use crate::ast::vm::Closure; +use std::cell::RefCell; +use std::rc::Rc; + +use super::folder::Folder; +use super::inliner::Inliner; +use super::substitution_map::SubstitutionMap; +use super::utils::{PathTracker, UsageInfo}; + +pub struct Optimizer { + pub enabled: bool, + max_passes: usize, + pub globals: Option>>>, + pub root_purity: Option>>>, + pub lambda_registry: Option>>, +} + +impl Optimizer { + pub fn new(enabled: bool) -> Self { + Self { + enabled, + max_passes: 5, + globals: None, + root_purity: None, + lambda_registry: None, + } + } + + pub fn with_globals(mut self, globals: Rc>>) -> Self { + self.globals = Some(globals); + self + } + + pub fn with_purity(mut self, purity: Rc>>) -> Self { + self.root_purity = Some(purity); + self + } + + pub fn with_registry( + mut self, + registry: Rc>, + ) -> Self { + self.lambda_registry = Some(registry); + self + } + + pub fn optimize(&self, node: AnalyzedNode) -> AnalyzedNode { + if !self.enabled { + return node; + } + + let mut current = Rc::new(node); + for _ in 0..self.max_passes { + let mut sub = SubstitutionMap::new(); + let mut path = PathTracker::new(); + let next = self.visit_node(current.clone(), &mut sub, &mut path); + if Rc::ptr_eq(&next, ¤t) { + break; + } + current = next; + } + // Unwrap the final Rc if we are at the end, or return a clone of the inner node. + // Since AnalyzedNode is small now (header + Rcs), cloning is cheap. + (*current).clone() + } + + fn try_inline( + &self, + params: &AnalyzedNode, + arg_nodes: &[Rc], + body: &AnalyzedNode, + sub: &mut SubstitutionMap, + path: &mut PathTracker, + base_sub: Option, + ) -> Option> { + let inliner = Inliner::new(&self.globals, &self.root_purity); + let mut inner_sub = base_sub.unwrap_or_else(|| sub.new_for_inlining()); + + if inliner.prepare_beta_reduction(params, arg_nodes, body, &mut inner_sub).is_some() { + let res = self.visit_node(Rc::new(body.clone()), &mut inner_sub, path); + + // Sync back state to parent substitution map + sub.next_slot = inner_sub.next_slot; + sub.used.extend(inner_sub.used.iter().cloned()); + sub.assigned.extend(inner_sub.assigned.iter().cloned()); + sub.captured_slots.extend(inner_sub.captured_slots.iter().cloned()); + + Some(res) + } else { + None + } + } + + fn visit_node( + &self, + node_rc: Rc, + sub: &mut SubstitutionMap, + path: &mut PathTracker, + ) -> Rc { + let node = &*node_rc; + let folder = Folder::new(&self.globals); + let inliner = Inliner::new(&self.globals, &self.root_purity); + + let (new_kind, metrics) = match &node.kind { + NodeKind::Identifier { symbol, binding } => { + let addr = match binding { + IdentifierBinding::Reference(addr) => *addr, + IdentifierBinding::Declaration { addr, .. } => *addr, + }; + if !sub.assigned.contains(&addr) { + // 1. Try inlining from current value substitution map (locals/globals/upvalues) + if let Some(val) = sub.get_value(&addr) + && inliner.is_inlinable_value(val, addr) + { + return Rc::new(folder.make_constant_node(val.clone(), node)); + } + + // 2. Try inlining from AST substitution map (pure expressions) + if let Some(inlined_node) = sub.ast_substitutions.get(&addr) { + return inlined_node.clone(); + } + + // 3. Fallback for Globals: check the actual VM environment + if let Address::Global(idx) = addr + && let Some(globals_rc) = &self.globals + { + let globals = globals_rc.borrow(); + if let Some(val) = globals.get(idx.0 as usize) + && inliner.is_inlinable_value(val, addr) + { + return Rc::new(folder.make_constant_node(val.clone(), node)); + } + } + } + + sub.used.insert(addr); + let new_binding = match binding { + IdentifierBinding::Reference(_) => { + IdentifierBinding::Reference(sub.map_address(addr)) + } + IdentifierBinding::Declaration { kind, .. } => { + IdentifierBinding::Declaration { + addr: sub.map_address(addr), + kind: *kind, + } + } + }; + ( + NodeKind::Identifier { + symbol: symbol.clone(), + binding: new_binding, + }, + node.ty.clone(), + ) + } + + NodeKind::FieldAccessor(k) => (NodeKind::FieldAccessor(*k), node.ty.clone()), + + NodeKind::GetField { rec, field } => { + let rec_opt = self.visit_node(rec.clone(), sub, path); + + // Constant folding for Field Access + if let NodeKind::Constant(Value::Record(layout, values)) = &rec_opt.kind + && let Some(idx) = layout.index_of(*field) + { + return Rc::new(folder.make_constant_node(values[idx].clone(), node)); + } + + if Rc::ptr_eq(&rec_opt, rec) { + return node_rc; + } + + ( + NodeKind::GetField { + rec: rec_opt, + field: *field, + }, + node.ty.clone(), + ) + } + + NodeKind::Assign { target, value, info } => { + let addr = info.addr; + let value_opt = self.visit_node(value.clone(), sub, path); + if let Some(addr) = addr { + if let NodeKind::Constant(val) = &value_opt.kind { + sub.add_value(addr, val.clone()); + } else { + sub.remove_value(&addr); + } + } + + // Remap target addresses without constant folding (targets are lvalues) + let target_opt = Self::remap_pattern(target, sub); + + if Rc::ptr_eq(&value_opt, value) && Rc::ptr_eq(&target_opt, target) { + return node_rc; + } + + let new_info = if let Some(addr) = addr { + AssignBinding { + addr: Some(sub.map_address(addr)), + } + } else { + info.clone() + }; + + ( + NodeKind::Assign { + target: target_opt, + value: value_opt, + info: new_info, + }, + node.ty.clone(), + ) + } + + NodeKind::Def { + pattern, + value, + info, + } => { + let value_opt = self.visit_node(value.clone(), sub, path); + + // Extract addr from the ORIGINAL pattern (before visiting) + let addr = Self::extract_def_addr(pattern); + + if let Some(addr) = addr { + if let Address::Local(slot) = addr + && !info.captured_by.is_empty() + { + sub.captured_slots.insert(slot); + } + + if let NodeKind::Constant(val) = &value_opt.kind { + sub.add_value(addr, val.clone()); + } else { + sub.remove_value(&addr); + } + + if let Address::Global(global_index) = addr + && value_opt.ty.purity > Purity::Impure + && let Some(purity_rc) = &self.root_purity + { + let mut pr = purity_rc.borrow_mut(); + let idx = global_index.0 as usize; + if idx < pr.len() { + pr[idx] = value_opt.ty.purity; + } + } + } + + // Remap pattern addresses without constant folding (patterns are lvalues) + let pattern_opt = Self::remap_pattern(pattern, sub); + + if Rc::ptr_eq(&value_opt, value) && Rc::ptr_eq(&pattern_opt, pattern) { + return node_rc; + } + + ( + NodeKind::Def { + pattern: pattern_opt, + value: value_opt, + info: info.clone(), + }, + node.ty.clone(), + ) + } + + NodeKind::Call { callee, args } => { + let callee_opt = self.visit_node(callee.clone(), sub, path); + let args_opt = self.visit_node(args.clone(), sub, path); + + if self.enabled { + // Constant folding for Call { callee: FieldAccessor, args: Tuple[1] } + // (field access on a constant record) + if let NodeKind::FieldAccessor(k) = &callee_opt.kind + && let NodeKind::Tuple { elements } = &args_opt.kind + && elements.len() == 1 + && let NodeKind::Constant(Value::Record(layout, values)) = &elements[0].kind + && let Some(idx) = layout.index_of(*k) + { + return Rc::new(folder.make_constant_node(values[idx].clone(), node)); + } + + let mut arg_nodes = Vec::new(); + self.flatten_tuple(args_opt.clone(), &mut arg_nodes); + + if let NodeKind::Lambda { + params, + body, + info: lambda_info, + } = &callee_opt.kind + && lambda_info.upvalues.is_empty() + && lambda_info.positional_count.is_some() + && path.inlining_depth < 5 + && !callee_opt.ty.is_recursive + && path.enter_lambda(&callee_opt.identity) + { + path.inlining_depth += 1; + let collapsed = self.try_inline(params, &arg_nodes, body, sub, path, None); + path.inlining_depth -= 1; + path.exit_lambda(&callee_opt.identity); + + if let Some(res) = collapsed { + return res; + } + } + + if let NodeKind::Identifier { + binding: IdentifierBinding::Reference(Address::Global(idx)), + .. + } = &callee_opt.kind + && let Some(registry_rc) = &self.lambda_registry + && path.inlining_depth < 5 + && !path.inlining_stack.contains(idx) + { + let registry = registry_rc.borrow(); + if let Some(lambda_node) = registry.get(idx) + && let NodeKind::Lambda { + params, + body, + info: lambda_info, + } = &lambda_node.kind + && lambda_info.upvalues.is_empty() + && lambda_info.positional_count.is_some() + && !lambda_node.ty.is_recursive + { + path.inlining_stack.insert(*idx); + path.inlining_depth += 1; + let collapsed = self.try_inline(params, &arg_nodes, body, sub, path, None); + path.inlining_depth -= 1; + path.inlining_stack.remove(idx); + + if let Some(res) = collapsed { + return res; + } + } + } + + if let NodeKind::Constant(Value::Object(ref obj)) = callee_opt.kind + && path.inlining_depth < 5 + && let Some(closure) = obj.as_any().downcast_ref::() + && (closure.upvalues.is_empty() + || closure.function_node.ty.purity >= Purity::SideEffectFree) + && !closure.function_node.ty.is_recursive + && path.enter_lambda(&closure.function_node.identity) + { + let mut closure_sub = sub.new_for_inlining(); + for (i, cell) in closure.upvalues.iter().enumerate() { + closure_sub.add_value( + Address::Upvalue(UpvalueIdx(i as u32)), + cell.borrow().clone(), + ); + } + + path.inlining_depth += 1; + if let NodeKind::Lambda { params, .. } = &closure.function_node.kind { + let collapsed = self.try_inline( + params, + &arg_nodes, + &closure.function_node, + sub, + path, + Some(closure_sub), + ); + path.inlining_depth -= 1; + path.exit_lambda(&closure.function_node.identity); + + if let Some(res) = collapsed { + return res; + } + } else { + path.inlining_depth -= 1; + path.exit_lambda(&closure.function_node.identity); + } + } + + if let Some(folded) = folder.try_fold_pure(&callee_opt, &arg_nodes) { + return Rc::new(folded); + } + } + + if Rc::ptr_eq(&callee_opt, callee) && Rc::ptr_eq(&args_opt, args) { + return node_rc; + } + + ( + NodeKind::Call { + callee: callee_opt, + args: args_opt, + }, + node.ty.clone(), + ) + } + + NodeKind::Again { args } => { + let args_opt = self.visit_node(args.clone(), sub, path); + if Rc::ptr_eq(&args_opt, args) { + return node_rc; + } + ( + NodeKind::Again { + args: args_opt, + }, + node.ty.clone(), + ) + } + + NodeKind::If { + cond, + then_br, + else_br, + } => { + let cond_opt = self.visit_node(cond.clone(), sub, path); + if self.enabled + && let NodeKind::Constant(ref val) = cond_opt.kind + { + if val.is_truthy() { + return self.visit_node(then_br.clone(), sub, path); + } else if let Some(else_node) = else_br { + return self.visit_node(else_node.clone(), sub, path); + } else { + return Rc::new(folder.make_nop_node(node)); + } + } + + let then_br_opt = self.visit_node(then_br.clone(), sub, path); + let else_br_opt = else_br + .as_ref() + .map(|e| self.visit_node(e.clone(), sub, path)); + + let unchanged = Rc::ptr_eq(&cond_opt, cond) && Rc::ptr_eq(&then_br_opt, then_br) && match (&else_br_opt, else_br) { + (Some(a), Some(b)) => Rc::ptr_eq(a, b), + (None, None) => true, + _ => false, + }; + + if unchanged { + return node_rc; + } + + ( + NodeKind::If { + cond: cond_opt, + then_br: then_br_opt, + else_br: else_br_opt, + }, + node.ty.clone(), + ) + } + + NodeKind::Pipe { + inputs, + lambda, + } => { + let mut o_inputs = Vec::with_capacity(inputs.len()); + let mut inputs_changed = false; + for input in inputs { + let opt = self.visit_node(input.clone(), sub, path); + if !Rc::ptr_eq(&opt, input) { + inputs_changed = true; + } + o_inputs.push(opt); + } + let o_lambda = self.visit_node(lambda.clone(), sub, path); + + if !inputs_changed && Rc::ptr_eq(&o_lambda, lambda) { + return node_rc; + } + + ( + NodeKind::Pipe { + inputs: o_inputs, + lambda: o_lambda, + }, + node.ty.clone(), + ) + } + NodeKind::Block { exprs } => { + let mut info = UsageInfo::default(); + if !exprs.is_empty() { + for e in exprs { + info.collect(e); + } + } + + sub.assigned.extend(info.assigned.iter().cloned()); + + let mut new_exprs = Vec::with_capacity(exprs.len()); + let last_idx = exprs.len().saturating_sub(1); + let mut block_changed = false; + + for (i, e) in exprs.iter().enumerate() { + let is_last = i == last_idx; + if self.enabled && !is_last { + let removable = match &e.kind { + NodeKind::Def { pattern, value, .. } => { + let addr = Self::extract_def_addr(pattern); + if let Some(addr) = addr { + !info.is_used(&addr) + && (if let Address::Local(slot) = addr { + !sub.captured_slots.contains(&slot) + } else { + true + }) + && (value.ty.purity >= Purity::SideEffectFree + || matches!(value.kind, NodeKind::Lambda { .. })) + } else { + // Destructuring def without single addr - not removable + false + } + } + NodeKind::Assign { info: assign_info, value, .. } => { + if let Some(addr) = assign_info.addr { + !info.is_used(&addr) + && (if let Address::Local(slot) = addr { + !sub.captured_slots.contains(&slot) + } else { + true + }) + && value.ty.purity >= Purity::SideEffectFree + } else { + false + } + } + _ => e.ty.purity >= Purity::SideEffectFree, + }; + + if removable { + block_changed = true; + continue; + } + } + + let unmapped_addr = if let NodeKind::Def { pattern, .. } = &e.kind { + Self::extract_def_addr(pattern) + } else { + None + }; + + let opt = self.visit_node(e.clone(), sub, path); + + if let NodeKind::Def { value, .. } = &opt.kind + && let Some(orig_addr) = unmapped_addr + && !info.assigned.contains(&orig_addr) + { + let mut core_value = value.as_ref(); + while let NodeKind::Expansion { expanded, .. } = &core_value.kind { + core_value = expanded.as_ref(); + } + + if let NodeKind::Constant(val) = &core_value.kind { + sub.add_value(orig_addr, val.clone()); + } else if let NodeKind::Lambda { info: lambda_info, .. } = &core_value.kind + && lambda_info.upvalues.is_empty() + { + sub.add_ast_substitution(orig_addr, core_value.clone()); + } + } + + if self.enabled && matches!(opt.kind, NodeKind::Nop) && !is_last { + block_changed = true; + continue; + } + if !Rc::ptr_eq(&opt, e) { + block_changed = true; + } + new_exprs.push(opt); + } + + if !block_changed { + return node_rc; + } + + if self.enabled { + if new_exprs.is_empty() { + return Rc::new(folder.make_nop_node(node)); + } else if new_exprs.len() == 1 { + return new_exprs.pop().unwrap(); + } + } + + (NodeKind::Block { exprs: new_exprs }, node.ty.clone()) + } + + NodeKind::Lambda { + params, + body, + info: lambda_info, + } => { + let mut usage_info = UsageInfo::default(); + usage_info.collect(node); + + let mut new_upvalues = Vec::new(); + let mut mapping = Vec::new(); + let mut next_inner_subs = sub.new_inner(); + next_inner_subs.assigned = usage_info.assigned; + let mut upvalues_changed = false; + + for (old_idx, capture_addr) in lambda_info.upvalues.iter().enumerate() { + let mut inlined_val = None; + let mut inlined_ast = None; + if !sub.assigned.contains(capture_addr) { + if let Some(val) = sub.get_value(capture_addr) { + inlined_val = Some(val.clone()); + } else if let Some(ast) = sub.ast_substitutions.get(capture_addr) { + inlined_ast = Some(Rc::clone(ast)); + } + } + + if let Address::Local(slot) = capture_addr { + sub.captured_slots.insert(*slot); + } + + if let Some(val) = inlined_val { + next_inner_subs + .add_value(Address::Upvalue(UpvalueIdx(old_idx as u32)), val); + mapping.push(None); + upvalues_changed = true; + } else if let Some(ast) = inlined_ast { + next_inner_subs + .add_ast_substitution(Address::Upvalue(UpvalueIdx(old_idx as u32)), (*ast).clone()); + mapping.push(None); + upvalues_changed = true; + } else { + mapping.push(Some(new_upvalues.len() as u32)); + let mapped = sub.map_address(*capture_addr); + if mapped != *capture_addr { + upvalues_changed = true; + } + new_upvalues.push(mapped); + } + } + + let params_opt = self.visit_node(params.clone(), &mut next_inner_subs, path); + + let body_opt = self.visit_node(body.clone(), &mut next_inner_subs, path); + let reindexed_body = if new_upvalues.len() != lambda_info.upvalues.len() { + sub.reindex_upvalues(body_opt, &mapping) + } else { + body_opt + }; + + if !upvalues_changed && Rc::ptr_eq(¶ms_opt, params) && Rc::ptr_eq(&reindexed_body, body) { + return node_rc; + } + + ( + NodeKind::Lambda { + params: params_opt, + body: reindexed_body, + info: LambdaBinding { + upvalues: new_upvalues, + positional_count: lambda_info.positional_count, + }, + }, + node.ty.clone(), + ) + } + + NodeKind::Tuple { elements } => { + let mut new_elements = Vec::with_capacity(elements.len()); + let mut changed = false; + for e in elements { + let opt = self.visit_node(e.clone(), sub, path); + if !Rc::ptr_eq(&opt, e) { + changed = true; + } + new_elements.push(opt); + } + if !changed { + return node_rc; + } + (NodeKind::Tuple { elements: new_elements }, node.ty.clone()) + } + NodeKind::Record { fields, layout } => { + let mut mapped_fields = Vec::with_capacity(fields.len()); + let mut changed = false; + for (k, v) in fields { + let v_opt = self.visit_node(v.clone(), sub, path); + if !Rc::ptr_eq(&v_opt, v) { + changed = true; + } + mapped_fields.push((k.clone(), v_opt)); + } + + if self.enabled { + let values: Vec<_> = mapped_fields.iter().map(|(_, v)| v.clone()).collect(); + if let Some(folded) = folder.try_fold_record(layout, &values, node) + { + return Rc::new(folded); + } + } + + if !changed { + return node_rc; + } + + ( + NodeKind::Record { + fields: mapped_fields, + layout: layout.clone(), + }, + node.ty.clone(), + ) + } + NodeKind::Expansion { + original_call, + expanded, + } => { + path.inlining_depth += 1; + let expanded_opt = self.visit_node(expanded.clone(), sub, path); + path.inlining_depth -= 1; + + if Rc::ptr_eq(&expanded_opt, expanded) { + return node_rc; + } + + ( + NodeKind::Expansion { + original_call: original_call.clone(), + expanded: expanded_opt, + }, + node.ty.clone(), + ) + } + k => (k.clone(), node.ty.clone()), + }; + + Rc::new(Node { + identity: node.identity.clone(), + kind: new_kind, + ty: metrics, + }) + } + + /// Extracts the address from a Def pattern node (when it's a simple Identifier with Declaration binding). + fn extract_def_addr(pattern: &AnalyzedNode) -> Option> { + if let NodeKind::Identifier { + binding: IdentifierBinding::Declaration { addr, .. }, + .. + } = &pattern.kind + { + Some(*addr) + } else { + None + } + } + + /// Remaps addresses in a pattern/target node without constant folding. + /// Used for Assign targets where identifiers are references to existing variables + /// that should not be replaced by their values. + fn remap_pattern(node_rc: &Rc, sub: &mut SubstitutionMap) -> Rc { + let node = &**node_rc; + let new_kind = match &node.kind { + NodeKind::Identifier { symbol, binding } => { + let addr = match binding { + IdentifierBinding::Reference(addr) => *addr, + IdentifierBinding::Declaration { addr, .. } => *addr, + }; + let new_binding = match binding { + IdentifierBinding::Reference(_) => { + IdentifierBinding::Reference(sub.map_address(addr)) + } + IdentifierBinding::Declaration { kind, .. } => { + IdentifierBinding::Declaration { + addr: sub.map_address(addr), + kind: *kind, + } + } + }; + NodeKind::Identifier { + symbol: symbol.clone(), + binding: new_binding, + } + } + NodeKind::Tuple { elements } => { + let new_elements = elements + .iter() + .map(|e| Self::remap_pattern(e, sub)) + .collect(); + NodeKind::Tuple { + elements: new_elements, + } + } + _ => return node_rc.clone(), + }; + Rc::new(Node { + identity: node.identity.clone(), + kind: new_kind, + ty: node.ty.clone(), + }) + } + + fn flatten_tuple(&self, node: Rc, into: &mut Vec>) { + match &node.kind { + NodeKind::Tuple { elements } => { + for el in elements { + self.flatten_tuple(el.clone(), into); + } + } + NodeKind::Nop => {} + _ => into.push(node), + } + } +} diff --git a/src/ast/compiler/optimizer/folder.rs b/src/ast/compiler/optimizer/folder.rs index 1aef610..d938051 100644 --- a/src/ast/compiler/optimizer/folder.rs +++ b/src/ast/compiler/optimizer/folder.rs @@ -1,101 +1,103 @@ -use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, Node, NodeMetrics}; -use crate::ast::types::{Purity, RecordLayout, StaticType, Value}; -use std::cell::RefCell; -use std::rc::Rc; - -pub struct Folder<'a> { - pub globals: &'a Option>>>, -} - -impl<'a> Folder<'a> { - pub fn new(globals: &'a Option>>>) -> Self { - Self { globals } - } - - pub fn make_constant_node(&self, val: Value, template: &AnalyzedNode) -> AnalyzedNode { - let ty = val.static_type(); - let typed_original = Rc::new(Node { - identity: template.identity.clone(), - kind: BoundKind::Constant(val.clone()), - ty: ty.clone(), - }); - Node { - identity: template.identity.clone(), - kind: BoundKind::Constant(val), - ty: NodeMetrics { - original: typed_original, - purity: Purity::Pure, - is_recursive: false, - }, - } - } - - pub fn make_nop_node(&self, template: &AnalyzedNode) -> AnalyzedNode { - let typed_original = Rc::new(Node { - identity: template.identity.clone(), - kind: BoundKind::Nop, - ty: StaticType::Void, - }); - Node { - identity: template.identity.clone(), - kind: BoundKind::Nop, - ty: NodeMetrics { - original: typed_original, - purity: Purity::Pure, - is_recursive: false, - }, - } - } - - pub fn try_fold_record( - &self, - layout: &std::sync::Arc, - values: &[Rc], - template: &AnalyzedNode, - ) -> Option { - let mut constant_values = Vec::with_capacity(values.len()); - - for v_node in values { - if let BoundKind::Constant(val) = &v_node.kind { - constant_values.push(val.clone()); - } else { - return None; - } - } - - let record_val = Value::Record(layout.clone(), Rc::new(constant_values)); - Some(self.make_constant_node(record_val, template)) - } - - pub fn try_fold_pure( - &self, - callee: &AnalyzedNode, - arg_nodes: &[Rc], - ) -> Option { - if callee.ty.purity < Purity::Pure { - return None; - } - - let mut arg_values = Vec::with_capacity(arg_nodes.len()); - for node in arg_nodes { - if let BoundKind::Constant(val) = &node.kind { - arg_values.push(val.clone()); - } else { - return None; - } - } - let func_val = match &callee.kind { - BoundKind::Get { - addr: Address::Global(idx), - .. - } => self.globals.as_ref()?.borrow().get(idx.0 as usize)?.clone(), - BoundKind::Constant(val) => val.clone(), - _ => return None, - }; - let result = match func_val { - Value::Function(f) => (f.func)(&arg_values), - _ => return None, - }; - Some(self.make_constant_node(result, callee)) - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AnalyzedNode, IdentifierBinding, Node, NodeKind, NodeMetrics, +}; +use crate::ast::types::{Purity, RecordLayout, StaticType, Value}; +use std::cell::RefCell; +use std::rc::Rc; + +pub struct Folder<'a> { + pub globals: &'a Option>>>, +} + +impl<'a> Folder<'a> { + pub fn new(globals: &'a Option>>>) -> Self { + Self { globals } + } + + pub fn make_constant_node(&self, val: Value, template: &AnalyzedNode) -> AnalyzedNode { + let ty = val.static_type(); + let typed_original = Rc::new(Node { + identity: template.identity.clone(), + kind: NodeKind::Constant(val.clone()), + ty: ty.clone(), + }); + Node { + identity: template.identity.clone(), + kind: NodeKind::Constant(val), + ty: NodeMetrics { + original: typed_original, + purity: Purity::Pure, + is_recursive: false, + }, + } + } + + pub fn make_nop_node(&self, template: &AnalyzedNode) -> AnalyzedNode { + let typed_original = Rc::new(Node { + identity: template.identity.clone(), + kind: NodeKind::Nop, + ty: StaticType::Void, + }); + Node { + identity: template.identity.clone(), + kind: NodeKind::Nop, + ty: NodeMetrics { + original: typed_original, + purity: Purity::Pure, + is_recursive: false, + }, + } + } + + pub fn try_fold_record( + &self, + layout: &std::sync::Arc, + values: &[Rc], + template: &AnalyzedNode, + ) -> Option { + let mut constant_values = Vec::with_capacity(values.len()); + + for v_node in values { + if let NodeKind::Constant(val) = &v_node.kind { + constant_values.push(val.clone()); + } else { + return None; + } + } + + let record_val = Value::Record(layout.clone(), Rc::new(constant_values)); + Some(self.make_constant_node(record_val, template)) + } + + pub fn try_fold_pure( + &self, + callee: &AnalyzedNode, + arg_nodes: &[Rc], + ) -> Option { + if callee.ty.purity < Purity::Pure { + return None; + } + + let mut arg_values = Vec::with_capacity(arg_nodes.len()); + for node in arg_nodes { + if let NodeKind::Constant(val) = &node.kind { + arg_values.push(val.clone()); + } else { + return None; + } + } + let func_val = match &callee.kind { + NodeKind::Identifier { + binding: IdentifierBinding::Reference(Address::Global(idx)), + .. + } => self.globals.as_ref()?.borrow().get(idx.0 as usize)?.clone(), + NodeKind::Constant(val) => val.clone(), + _ => return None, + }; + let result = match func_val { + Value::Function(f) => (f.func)(&arg_values), + _ => return None, + }; + Some(self.make_constant_node(result, callee)) + } +} diff --git a/src/ast/compiler/optimizer/inliner.rs b/src/ast/compiler/optimizer/inliner.rs index 0169c61..c0afbba 100644 --- a/src/ast/compiler/optimizer/inliner.rs +++ b/src/ast/compiler/optimizer/inliner.rs @@ -1,165 +1,222 @@ -use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, VirtualId}; -use crate::ast::types::{Purity, Value}; -use crate::ast::vm::Closure; -use std::cell::RefCell; -use std::collections::HashSet; -use std::rc::Rc; - -use super::substitution_map::SubstitutionMap; -use super::utils::UsageInfo; - -pub struct Inliner<'a> { - pub globals: &'a Option>>>, - pub root_purity: &'a Option>>>, -} - -impl<'a> Inliner<'a> { - pub fn new( - globals: &'a Option>>>, - root_purity: &'a Option>>>, - ) -> Self { - Self { - globals, - root_purity, - } - } - - pub fn is_inlinable_value(&self, val: &Value, addr: Address) -> bool { - let type_ok = match val { - Value::Int(_) - | Value::Float(_) - | Value::Bool(_) - | Value::Text(_) - | Value::Keyword(_) - | Value::Record(_, _) - | Value::DateTime(_) => true, - Value::Object(obj) => { - if let Some(closure) = obj.as_any().downcast_ref::() { - closure.upvalues.is_empty() && !closure.function_node.ty.is_recursive - } else { - false - } - } - _ => false, - }; - - if !type_ok { - return false; - } - - if let Address::Global(idx) = addr { - if let Some(purity_rc) = &self.root_purity { - let purity = purity_rc - .borrow() - .get(idx.0 as usize) - .cloned() - .unwrap_or(Purity::Impure); - return purity >= Purity::Pure; - } - return false; - } - - true - } - - pub fn prepare_beta_reduction( - &self, - params: &AnalyzedNode, - arg_vals: &[Rc], - body: &AnalyzedNode, - sub: &mut SubstitutionMap, - ) -> Option<()> { - let mut body_usage = UsageInfo::default(); - body_usage.collect(body); - - let mut slot_index = 0; - self.map_params_to_args(params, arg_vals, &mut slot_index, sub, &body_usage); - - if slot_index != arg_vals.len() { - return None; - } - - let mut param_slots = HashSet::new(); - self.collect_parameter_slots_set(params, &mut param_slots); - - for slot in param_slots { - let addr = Address::Local(slot); - if (body_usage.is_used(&addr) || body_usage.is_assigned(&addr)) - && sub.get_value(&addr).is_none() - && !sub.ast_substitutions.contains_key(&addr) - { - return None; - } - } - - if sub.values.is_empty() && sub.ast_substitutions.is_empty() && !arg_vals.is_empty() { - return None; - } - - Some(()) - } - - pub fn map_params_to_args( - &self, - pattern: &AnalyzedNode, - args: &[Rc], - offset: &mut usize, - sub: &mut SubstitutionMap, - body_usage: &UsageInfo, - ) { - match &pattern.kind { - BoundKind::Define { addr, .. } => { - if let Some(arg) = args.get(*offset) - && !body_usage.is_assigned(addr) - { - let mut core_arg = arg.as_ref(); - while let BoundKind::Expansion { bound_expanded, .. } = &core_arg.kind { - core_arg = bound_expanded.as_ref(); - } - - if let BoundKind::Constant(val) = &core_arg.kind { - sub.add_value(*addr, val.clone()); - } else if let BoundKind::Lambda { upvalues, .. } = &core_arg.kind - && upvalues.is_empty() - { - sub.add_ast_substitution(*addr, core_arg.clone()); - } else if let BoundKind::Get { - addr: Address::Global(_), - .. - } = &core_arg.kind - { - sub.add_ast_substitution(*addr, core_arg.clone()); - } - } - - if let Address::Local(slot) = addr { - sub.map_slot(*slot); - } - *offset += 1; - } - BoundKind::Tuple { elements } => { - for el in elements { - self.map_params_to_args(el, args, offset, sub, body_usage); - } - } - _ => {} - } - } - - pub fn collect_parameter_slots_set(&self, node: &AnalyzedNode, slots: &mut HashSet) { - match &node.kind { - BoundKind::Define { - addr: Address::Local(slot), - .. - } => { - slots.insert(*slot); - } - BoundKind::Tuple { elements } => { - for el in elements { - self.collect_parameter_slots_set(el, slots); - } - } - _ => {} - } - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AnalyzedNode, IdentifierBinding, NodeKind, VirtualId, +}; +use crate::ast::types::{Purity, Value}; +use crate::ast::vm::Closure; +use std::cell::RefCell; +use std::collections::HashSet; +use std::rc::Rc; + +use super::substitution_map::SubstitutionMap; +use super::utils::UsageInfo; + +pub struct Inliner<'a> { + pub globals: &'a Option>>>, + pub root_purity: &'a Option>>>, +} + +impl<'a> Inliner<'a> { + pub fn new( + globals: &'a Option>>>, + root_purity: &'a Option>>>, + ) -> Self { + Self { + globals, + root_purity, + } + } + + pub fn is_inlinable_value(&self, val: &Value, addr: Address) -> bool { + let type_ok = match val { + Value::Int(_) + | Value::Float(_) + | Value::Bool(_) + | Value::Text(_) + | Value::Keyword(_) + | Value::Record(_, _) + | Value::DateTime(_) => true, + Value::Object(obj) => { + if let Some(closure) = obj.as_any().downcast_ref::() { + closure.upvalues.is_empty() && !closure.function_node.ty.is_recursive + } else { + false + } + } + _ => false, + }; + + if !type_ok { + return false; + } + + if let Address::Global(idx) = addr { + if let Some(purity_rc) = &self.root_purity { + let purity = purity_rc + .borrow() + .get(idx.0 as usize) + .cloned() + .unwrap_or(Purity::Impure); + return purity >= Purity::Pure; + } + return false; + } + + true + } + + pub fn prepare_beta_reduction( + &self, + params: &AnalyzedNode, + arg_vals: &[Rc], + body: &AnalyzedNode, + sub: &mut SubstitutionMap, + ) -> Option<()> { + let mut body_usage = UsageInfo::default(); + body_usage.collect(body); + + let mut slot_index = 0; + self.map_params_to_args(params, arg_vals, &mut slot_index, sub, &body_usage); + + if slot_index != arg_vals.len() { + return None; + } + + let mut param_slots = HashSet::new(); + self.collect_parameter_slots_set(params, &mut param_slots); + + for slot in param_slots { + let addr = Address::Local(slot); + if (body_usage.is_used(&addr) || body_usage.is_assigned(&addr)) + && sub.get_value(&addr).is_none() + && !sub.ast_substitutions.contains_key(&addr) + { + return None; + } + } + + if sub.values.is_empty() && sub.ast_substitutions.is_empty() && !arg_vals.is_empty() { + return None; + } + + Some(()) + } + + pub fn map_params_to_args( + &self, + pattern: &AnalyzedNode, + args: &[Rc], + offset: &mut usize, + sub: &mut SubstitutionMap, + body_usage: &UsageInfo, + ) { + match &pattern.kind { + NodeKind::Def { pattern: inner_pattern, .. } => { + // Extract addr from the pattern's Identifier binding + let addr = if let NodeKind::Identifier { + binding: IdentifierBinding::Declaration { addr, .. }, + .. + } = &inner_pattern.kind + { + Some(*addr) + } else { + None + }; + + if let Some(addr) = addr { + if let Some(arg) = args.get(*offset) + && !body_usage.is_assigned(&addr) + { + let mut core_arg = arg.as_ref(); + while let NodeKind::Expansion { expanded, .. } = &core_arg.kind { + core_arg = expanded.as_ref(); + } + + if let NodeKind::Constant(val) = &core_arg.kind { + sub.add_value(addr, val.clone()); + } else if let NodeKind::Lambda { info: lambda_info, .. } = &core_arg.kind + && lambda_info.upvalues.is_empty() + { + sub.add_ast_substitution(addr, core_arg.clone()); + } else if let NodeKind::Identifier { + binding: IdentifierBinding::Reference(Address::Global(_)), + .. + } = &core_arg.kind + { + sub.add_ast_substitution(addr, core_arg.clone()); + } + } + + if let Address::Local(slot) = addr { + sub.map_slot(slot); + } + } + *offset += 1; + } + NodeKind::Identifier { + binding: IdentifierBinding::Declaration { addr, .. }, + .. + } => { + let addr = *addr; + if let Some(arg) = args.get(*offset) + && !body_usage.is_assigned(&addr) + { + let mut core_arg = arg.as_ref(); + while let NodeKind::Expansion { expanded, .. } = &core_arg.kind { + core_arg = expanded.as_ref(); + } + + if let NodeKind::Constant(val) = &core_arg.kind { + sub.add_value(addr, val.clone()); + } else if let NodeKind::Lambda { info: lambda_info, .. } = &core_arg.kind + && lambda_info.upvalues.is_empty() + { + sub.add_ast_substitution(addr, core_arg.clone()); + } else if let NodeKind::Identifier { + binding: IdentifierBinding::Reference(Address::Global(_)), + .. + } = &core_arg.kind + { + sub.add_ast_substitution(addr, core_arg.clone()); + } + } + + if let Address::Local(slot) = addr { + sub.map_slot(slot); + } + *offset += 1; + } + NodeKind::Tuple { elements } => { + for el in elements { + self.map_params_to_args(el, args, offset, sub, body_usage); + } + } + _ => {} + } + } + + pub fn collect_parameter_slots_set(&self, node: &AnalyzedNode, slots: &mut HashSet) { + match &node.kind { + NodeKind::Def { pattern, .. } => { + if let NodeKind::Identifier { + binding: IdentifierBinding::Declaration { addr: Address::Local(slot), .. }, + .. + } = &pattern.kind + { + slots.insert(*slot); + } + } + NodeKind::Identifier { + binding: IdentifierBinding::Declaration { addr: Address::Local(slot), .. }, + .. + } => { + slots.insert(*slot); + } + NodeKind::Tuple { elements } => { + for el in elements { + self.collect_parameter_slots_set(el, slots); + } + } + _ => {} + } + } +} diff --git a/src/ast/compiler/optimizer/substitution_map.rs b/src/ast/compiler/optimizer/substitution_map.rs index 3505d4c..021a0f1 100644 --- a/src/ast/compiler/optimizer/substitution_map.rs +++ b/src/ast/compiler/optimizer/substitution_map.rs @@ -1,210 +1,237 @@ -use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, Node, UpvalueIdx, VirtualId}; -use crate::ast::types::Value; -use std::collections::{HashMap, HashSet}; -use std::rc::Rc; - -#[derive(Default)] -pub struct SubstitutionMap { - pub values: HashMap, Value>, - pub ast_substitutions: HashMap, Rc>, - pub slot_mapping: HashMap, - pub assigned: HashSet>, - pub next_slot: u32, - pub used: HashSet>, - pub captured_slots: HashSet, -} - -impl SubstitutionMap { - pub fn new() -> Self { - Self::default() - } - - /// Creates a new SubstitutionMap for an inner scope (like an inlined Lambda). - /// Safely inherits only Global substitutions, because Local and Upvalue - /// addresses are relative to the specific function frame and would overlap. - pub fn new_inner(&self) -> Self { - let mut inner = Self::new(); - for (k, v) in &self.values { - if matches!(k, Address::Global(_)) { - inner.values.insert(*k, v.clone()); - } - } - for (k, v) in &self.ast_substitutions { - if matches!(k, Address::Global(_)) { - inner.ast_substitutions.insert(*k, v.clone()); - } - } - inner - } - - pub fn new_for_inlining(&self) -> Self { - let mut inner = self.new_inner(); - inner.next_slot = self.next_slot; - inner - } - - pub fn add_ast_substitution(&mut self, addr: Address, node: AnalyzedNode) { - self.ast_substitutions.insert(addr, Rc::new(node)); - } - - pub fn map_slot(&mut self, old_slot: VirtualId) -> VirtualId { - if let Some(&new_slot) = self.slot_mapping.get(&old_slot) { - return new_slot; - } - let new_slot = VirtualId(self.next_slot); - self.slot_mapping.insert(old_slot, new_slot); - self.next_slot += 1; - new_slot - } - - pub fn map_address(&mut self, addr: Address) -> Address { - match addr { - Address::Local(slot) => Address::Local(self.map_slot(slot)), - other => other, - } - } - - pub fn add_value(&mut self, addr: Address, val: Value) { - self.values.insert(addr, val); - } - - pub fn get_value(&self, addr: &Address) -> Option<&Value> { - self.values.get(addr) - } - - pub fn remove_value(&mut self, addr: &Address) { - self.values.remove(addr); - } - - fn reindex_addr(&self, addr: Address, mapping: &[Option]) -> Address { - if let Address::Upvalue(idx) = addr - && let Some(res) = mapping.get(idx.0 as usize) - && let Some(new_idx) = res - { - Address::Upvalue(UpvalueIdx(*new_idx)) - } else { - addr - } - } - - pub fn reindex_upvalues(&self, node_rc: Rc, mapping: &[Option]) -> Rc { - let node = &*node_rc; - let (new_kind, metrics) = match &node.kind { - BoundKind::Get { addr, name } => ( - BoundKind::Get { - addr: self.reindex_addr(*addr, mapping), - name: name.clone(), - }, - node.ty.clone(), - ), - BoundKind::Lambda { - params, - upvalues, - body, - positional_count, - } => { - let mut next_upvalues = Vec::new(); - for addr in upvalues { - next_upvalues.push(self.reindex_addr(*addr, mapping)); - } - ( - BoundKind::Lambda { - params: params.clone(), - upvalues: next_upvalues, - body: body.clone(), - positional_count: *positional_count, - }, - node.ty.clone(), - ) - } - BoundKind::If { - cond, - then_br, - else_br, - } => { - let cond = self.reindex_upvalues(cond.clone(), mapping); - let then_br = self.reindex_upvalues(then_br.clone(), mapping); - let else_br = else_br.as_ref().map(|e| self.reindex_upvalues(e.clone(), mapping)); - ( - BoundKind::If { - cond, - then_br, - else_br, - }, - node.ty.clone(), - ) - } - BoundKind::Block { exprs } => { - let exprs = exprs - .iter() - .map(|e| self.reindex_upvalues(e.clone(), mapping)) - .collect(); - (BoundKind::Block { exprs }, node.ty.clone()) - } - BoundKind::Call { callee, args } => { - let callee = self.reindex_upvalues(callee.clone(), mapping); - let args = self.reindex_upvalues(args.clone(), mapping); - (BoundKind::Call { callee, args }, node.ty.clone()) - } - BoundKind::Define { - name, - addr, - kind, - value, - captured_by, - } => { - let value = self.reindex_upvalues(value.clone(), mapping); - ( - BoundKind::Define { - name: name.clone(), - addr: *addr, - kind: *kind, - value, - captured_by: captured_by.clone(), - }, - node.ty.clone(), - ) - } - BoundKind::Set { addr, value } => { - let value = self.reindex_upvalues(value.clone(), mapping); - ( - BoundKind::Set { - addr: self.reindex_addr(*addr, mapping), - value, - }, - node.ty.clone(), - ) - } - BoundKind::Tuple { elements } => { - let elements = elements - .iter() - .map(|e| self.reindex_upvalues(e.clone(), mapping)) - .collect(); - (BoundKind::Tuple { elements }, node.ty.clone()) - } - BoundKind::Record { layout, values } => { - let values = values - .iter() - .map(|v| self.reindex_upvalues(v.clone(), mapping)) - .collect(); - (BoundKind::Record { layout: layout.clone(), values }, node.ty.clone()) - } - BoundKind::Expansion { original_call, bound_expanded } => { - let bound_expanded = self.reindex_upvalues(bound_expanded.clone(), mapping); - ( - BoundKind::Expansion { - original_call: original_call.clone(), - bound_expanded, - }, - node.ty.clone(), - ) - } - k => (k.clone(), node.ty.clone()), - }; - Rc::new(Node { - identity: node.identity.clone(), - kind: new_kind, - ty: metrics, - }) - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AnalyzedNode, AssignBinding, IdentifierBinding, + LambdaBinding, Node, NodeKind, UpvalueIdx, VirtualId, +}; +use crate::ast::types::Value; +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; + +#[derive(Default)] +pub struct SubstitutionMap { + pub values: HashMap, Value>, + pub ast_substitutions: HashMap, Rc>, + pub slot_mapping: HashMap, + pub assigned: HashSet>, + pub next_slot: u32, + pub used: HashSet>, + pub captured_slots: HashSet, +} + +impl SubstitutionMap { + pub fn new() -> Self { + Self::default() + } + + /// Creates a new SubstitutionMap for an inner scope (like an inlined Lambda). + /// Safely inherits only Global substitutions, because Local and Upvalue + /// addresses are relative to the specific function frame and would overlap. + pub fn new_inner(&self) -> Self { + let mut inner = Self::new(); + for (k, v) in &self.values { + if matches!(k, Address::Global(_)) { + inner.values.insert(*k, v.clone()); + } + } + for (k, v) in &self.ast_substitutions { + if matches!(k, Address::Global(_)) { + inner.ast_substitutions.insert(*k, v.clone()); + } + } + inner + } + + pub fn new_for_inlining(&self) -> Self { + let mut inner = self.new_inner(); + inner.next_slot = self.next_slot; + inner + } + + pub fn add_ast_substitution(&mut self, addr: Address, node: AnalyzedNode) { + self.ast_substitutions.insert(addr, Rc::new(node)); + } + + pub fn map_slot(&mut self, old_slot: VirtualId) -> VirtualId { + if let Some(&new_slot) = self.slot_mapping.get(&old_slot) { + return new_slot; + } + let new_slot = VirtualId(self.next_slot); + self.slot_mapping.insert(old_slot, new_slot); + self.next_slot += 1; + new_slot + } + + pub fn map_address(&mut self, addr: Address) -> Address { + match addr { + Address::Local(slot) => Address::Local(self.map_slot(slot)), + other => other, + } + } + + pub fn add_value(&mut self, addr: Address, val: Value) { + self.values.insert(addr, val); + } + + pub fn get_value(&self, addr: &Address) -> Option<&Value> { + self.values.get(addr) + } + + pub fn remove_value(&mut self, addr: &Address) { + self.values.remove(addr); + } + + fn reindex_addr(&self, addr: Address, mapping: &[Option]) -> Address { + if let Address::Upvalue(idx) = addr + && let Some(res) = mapping.get(idx.0 as usize) + && let Some(new_idx) = res + { + Address::Upvalue(UpvalueIdx(*new_idx)) + } else { + addr + } + } + + pub fn reindex_upvalues(&self, node_rc: Rc, mapping: &[Option]) -> Rc { + let node = &*node_rc; + let (new_kind, metrics) = match &node.kind { + NodeKind::Identifier { symbol, binding } => { + let new_binding = match binding { + IdentifierBinding::Reference(addr) => { + IdentifierBinding::Reference(self.reindex_addr(*addr, mapping)) + } + IdentifierBinding::Declaration { addr, kind } => { + IdentifierBinding::Declaration { + addr: self.reindex_addr(*addr, mapping), + kind: *kind, + } + } + }; + ( + NodeKind::Identifier { + symbol: symbol.clone(), + binding: new_binding, + }, + node.ty.clone(), + ) + } + NodeKind::Lambda { + params, + body, + info: lambda_info, + } => { + let mut next_upvalues = Vec::new(); + for addr in &lambda_info.upvalues { + next_upvalues.push(self.reindex_addr(*addr, mapping)); + } + ( + NodeKind::Lambda { + params: params.clone(), + body: body.clone(), + info: LambdaBinding { + upvalues: next_upvalues, + positional_count: lambda_info.positional_count, + }, + }, + node.ty.clone(), + ) + } + NodeKind::If { + cond, + then_br, + else_br, + } => { + let cond = self.reindex_upvalues(cond.clone(), mapping); + let then_br = self.reindex_upvalues(then_br.clone(), mapping); + let else_br = else_br.as_ref().map(|e| self.reindex_upvalues(e.clone(), mapping)); + ( + NodeKind::If { + cond, + then_br, + else_br, + }, + node.ty.clone(), + ) + } + NodeKind::Block { exprs } => { + let exprs = exprs + .iter() + .map(|e| self.reindex_upvalues(e.clone(), mapping)) + .collect(); + (NodeKind::Block { exprs }, node.ty.clone()) + } + NodeKind::Call { callee, args } => { + let callee = self.reindex_upvalues(callee.clone(), mapping); + let args = self.reindex_upvalues(args.clone(), mapping); + (NodeKind::Call { callee, args }, node.ty.clone()) + } + NodeKind::Def { + pattern, + value, + info, + } => { + let pattern = self.reindex_upvalues(pattern.clone(), mapping); + let value = self.reindex_upvalues(value.clone(), mapping); + ( + NodeKind::Def { + pattern, + value, + info: info.clone(), + }, + node.ty.clone(), + ) + } + NodeKind::Assign { + target, + value, + info: assign_info, + } => { + let target = self.reindex_upvalues(target.clone(), mapping); + let value = self.reindex_upvalues(value.clone(), mapping); + let new_info = if let Some(addr) = assign_info.addr { + AssignBinding { + addr: Some(self.reindex_addr(addr, mapping)), + } + } else { + assign_info.clone() + }; + ( + NodeKind::Assign { + target, + value, + info: new_info, + }, + node.ty.clone(), + ) + } + NodeKind::Tuple { elements } => { + let elements = elements + .iter() + .map(|e| self.reindex_upvalues(e.clone(), mapping)) + .collect(); + (NodeKind::Tuple { elements }, node.ty.clone()) + } + NodeKind::Record { fields, layout } => { + let fields = fields + .iter() + .map(|(k, v)| (k.clone(), self.reindex_upvalues(v.clone(), mapping))) + .collect(); + (NodeKind::Record { fields, layout: layout.clone() }, node.ty.clone()) + } + NodeKind::Expansion { original_call, expanded } => { + let expanded = self.reindex_upvalues(expanded.clone(), mapping); + ( + NodeKind::Expansion { + original_call: original_call.clone(), + expanded, + }, + node.ty.clone(), + ) + } + k => (k.clone(), node.ty.clone()), + }; + Rc::new(Node { + identity: node.identity.clone(), + kind: new_kind, + ty: metrics, + }) + } +} diff --git a/src/ast/compiler/optimizer/utils.rs b/src/ast/compiler/optimizer/utils.rs index 4bd204b..474ed4c 100644 --- a/src/ast/compiler/optimizer/utils.rs +++ b/src/ast/compiler/optimizer/utils.rs @@ -1,182 +1,219 @@ -use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, GlobalIdx, VirtualId}; -use crate::ast::types::{Identity, Value}; -use crate::ast::vm::Closure; -use std::collections::HashSet; - -// --- PathTracker --- -#[derive(Default)] -pub struct PathTracker { - pub inlining_depth: usize, - pub inlining_stack: HashSet, - pub identity_stack: HashSet, -} - -impl PathTracker { - pub fn new() -> Self { - Self::default() - } - - pub fn enter_lambda(&mut self, identity: &Identity) -> bool { - if self.identity_stack.contains(identity) { - return false; - } - self.identity_stack.insert(identity.clone()); - true - } - - pub fn exit_lambda(&mut self, identity: &Identity) { - self.identity_stack.remove(identity); - } -} - -// --- UsageInfo --- -#[derive(Default)] -pub struct UsageInfo { - pub used: HashSet>, - pub assigned: HashSet>, - pub used_identities: HashSet, -} - -impl UsageInfo { - pub fn is_assigned(&self, addr: &Address) -> bool { - self.assigned.contains(addr) - } - - pub fn is_used(&self, addr: &Address) -> bool { - self.used.contains(addr) - } - - pub fn collect(&mut self, node: &AnalyzedNode) { - match &node.kind { - BoundKind::Destructure { pattern, value } => { - self.collect_pattern(pattern); - self.collect(value); - } - BoundKind::Constant(v) => { - if let Value::Object(obj) = v - && let Some(closure) = obj.as_any().downcast_ref::() - { - self.used_identities - .insert(closure.function_node.identity.clone()); - self.collect(&closure.function_node); - } - } - BoundKind::Get { addr, .. } => { - self.used.insert(*addr); - } - BoundKind::Set { addr, value } => { - self.assigned.insert(*addr); - self.collect(value); - } - BoundKind::GetField { rec, .. } => { - self.collect(rec); - } - BoundKind::Lambda { - params, - body, - upvalues, - .. - } => { - self.used_identities.insert(node.identity.clone()); - - let mut inner_info = UsageInfo::default(); - inner_info.collect(params); - inner_info.collect(body); - - // Propagate globals and identities - for addr in &inner_info.used { - if let Address::Global(_) = addr { - self.used.insert(*addr); - } - } - for addr in &inner_info.assigned { - if let Address::Global(_) = addr { - self.assigned.insert(*addr); - } - } - self.used_identities.extend(inner_info.used_identities); - - // Map used upvalues to parent scope - for addr in &inner_info.used { - if let Address::Upvalue(idx) = addr - && let Some(parent_addr) = upvalues.get(idx.0 as usize) - { - self.used.insert(*parent_addr); - } - } - // Map assigned upvalues to parent scope - for addr in &inner_info.assigned { - if let Address::Upvalue(idx) = addr - && let Some(parent_addr) = upvalues.get(idx.0 as usize) - { - self.assigned.insert(*parent_addr); - } - } - } - BoundKind::Block { exprs } => { - for e in exprs { - self.collect(e); - } - } - BoundKind::If { - cond, - then_br, - else_br, - } => { - self.collect(cond); - self.collect(then_br); - if let Some(e) = else_br { - self.collect(e); - } - } - BoundKind::Call { callee, args } => { - self.collect(callee); - self.collect(args); - } - BoundKind::Tuple { elements } => { - for e in elements { - self.collect(e); - } - } - BoundKind::Record { values, .. } => { - for v in values { - self.collect(v); - } - } - BoundKind::Define { value, .. } => { - self.collect(value); - } - BoundKind::Expansion { bound_expanded, .. } => { - self.collect(bound_expanded); - } - BoundKind::Again { args } => { - self.collect(args); - } - BoundKind::Pipe { inputs, lambda, .. } => { - for input in inputs { - self.collect(input); - } - self.collect(lambda); - } - BoundKind::Nop - | BoundKind::FieldAccessor(_) - | BoundKind::Extension(_) - | BoundKind::Error => {} - } - } - - pub fn collect_pattern(&mut self, node: &AnalyzedNode) { - match &node.kind { - BoundKind::Define { .. } => {} - BoundKind::Set { addr, .. } => { - self.assigned.insert(*addr); - } - BoundKind::Tuple { elements } => { - for el in elements { - self.collect_pattern(el); - } - } - _ => {} - } - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AnalyzedNode, GlobalIdx, IdentifierBinding, + NodeKind, VirtualId, +}; +use crate::ast::types::{Identity, Value}; +use crate::ast::vm::Closure; +use std::collections::HashSet; + +// --- PathTracker --- +#[derive(Default)] +pub struct PathTracker { + pub inlining_depth: usize, + pub inlining_stack: HashSet, + pub identity_stack: HashSet, +} + +impl PathTracker { + pub fn new() -> Self { + Self::default() + } + + pub fn enter_lambda(&mut self, identity: &Identity) -> bool { + if self.identity_stack.contains(identity) { + return false; + } + self.identity_stack.insert(identity.clone()); + true + } + + pub fn exit_lambda(&mut self, identity: &Identity) { + self.identity_stack.remove(identity); + } +} + +// --- UsageInfo --- +#[derive(Default)] +pub struct UsageInfo { + pub used: HashSet>, + pub assigned: HashSet>, + pub used_identities: HashSet, +} + +impl UsageInfo { + pub fn is_assigned(&self, addr: &Address) -> bool { + self.assigned.contains(addr) + } + + pub fn is_used(&self, addr: &Address) -> bool { + self.used.contains(addr) + } + + pub fn collect(&mut self, node: &AnalyzedNode) { + match &node.kind { + NodeKind::Def { pattern, value, .. } => { + self.collect_pattern(pattern); + self.collect(value); + } + NodeKind::Constant(v) => { + if let Value::Object(obj) = v + && let Some(closure) = obj.as_any().downcast_ref::() + { + self.used_identities + .insert(closure.function_node.identity.clone()); + self.collect(&closure.function_node); + } + } + NodeKind::Identifier { binding, .. } => { + let addr = match binding { + IdentifierBinding::Reference(addr) => *addr, + IdentifierBinding::Declaration { addr, .. } => *addr, + }; + self.used.insert(addr); + } + NodeKind::Assign { target, value, info, .. } => { + if let Some(addr) = info.addr { + self.assigned.insert(addr); + } else { + // Destructuring assign: collect assigned addresses from target pattern + self.collect_assigned_from_target(target); + } + self.collect(value); + } + NodeKind::GetField { rec, .. } => { + self.collect(rec); + } + NodeKind::Lambda { + params, + body, + info: lambda_info, + } => { + self.used_identities.insert(node.identity.clone()); + + let mut inner_info = UsageInfo::default(); + inner_info.collect(params); + inner_info.collect(body); + + // Propagate globals and identities + for addr in &inner_info.used { + if let Address::Global(_) = addr { + self.used.insert(*addr); + } + } + for addr in &inner_info.assigned { + if let Address::Global(_) = addr { + self.assigned.insert(*addr); + } + } + self.used_identities.extend(inner_info.used_identities); + + // Map used upvalues to parent scope + for addr in &inner_info.used { + if let Address::Upvalue(idx) = addr + && let Some(parent_addr) = lambda_info.upvalues.get(idx.0 as usize) + { + self.used.insert(*parent_addr); + } + } + // Map assigned upvalues to parent scope + for addr in &inner_info.assigned { + if let Address::Upvalue(idx) = addr + && let Some(parent_addr) = lambda_info.upvalues.get(idx.0 as usize) + { + self.assigned.insert(*parent_addr); + } + } + } + NodeKind::Block { exprs } => { + for e in exprs { + self.collect(e); + } + } + NodeKind::If { + cond, + then_br, + else_br, + } => { + self.collect(cond); + self.collect(then_br); + if let Some(e) = else_br { + self.collect(e); + } + } + NodeKind::Call { callee, args } => { + self.collect(callee); + self.collect(args); + } + NodeKind::Tuple { elements } => { + for e in elements { + self.collect(e); + } + } + NodeKind::Record { fields, .. } => { + for (_, v) in fields { + self.collect(v); + } + } + NodeKind::Expansion { expanded, .. } => { + self.collect(expanded); + } + NodeKind::Again { args } => { + self.collect(args); + } + NodeKind::Pipe { inputs, lambda, .. } => { + for input in inputs { + self.collect(input); + } + self.collect(lambda); + } + NodeKind::Nop + | NodeKind::FieldAccessor(_) + | NodeKind::Extension(_) + | NodeKind::Error => {} + // Syntax-only variants that should not appear in AnalyzedPhase + NodeKind::MacroDecl { .. } + | NodeKind::Template(_) + | NodeKind::Placeholder(_) + | NodeKind::Splice(_) => {} + } + } + + pub fn collect_pattern(&mut self, node: &AnalyzedNode) { + match &node.kind { + NodeKind::Identifier { + binding: IdentifierBinding::Declaration { .. }, + .. + } => {} + NodeKind::Def { .. } => {} + NodeKind::Assign { info, .. } => { + if let Some(addr) = info.addr { + self.assigned.insert(addr); + } + } + NodeKind::Tuple { elements } => { + for el in elements { + self.collect_pattern(el); + } + } + _ => {} + } + } + + fn collect_assigned_from_target(&mut self, node: &AnalyzedNode) { + match &node.kind { + NodeKind::Identifier { binding, .. } => { + let addr = match binding { + IdentifierBinding::Reference(addr) + | IdentifierBinding::Declaration { addr, .. } => *addr, + }; + self.assigned.insert(addr); + } + NodeKind::Tuple { elements } => { + for el in elements { + self.collect_assigned_from_target(el); + } + } + _ => {} + } + } +} diff --git a/src/ast/compiler/specializer.rs b/src/ast/compiler/specializer.rs index e967288..07be732 100644 --- a/src/ast/compiler/specializer.rs +++ b/src/ast/compiler/specializer.rs @@ -1,276 +1,276 @@ -use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, AnalyzedPhase, BoundKind, Node, NodeMetrics, VirtualId}; -use crate::ast::types::{Purity, 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, -} - -pub type CompileFunc = Rc>, &[StaticType]) -> Result<(Value, StaticType), String>>; -pub type RtlLookupFunc = Rc Option<(Value, StaticType)>>; - -pub trait FunctionRegistry { - fn resolve(&self, addr: Address) -> Option>>; - fn resolve_analyzed(&self, _addr: Address) -> Option> { - None - } -} - -pub type MonoCache = HashMap; - -pub struct Specializer { - pub cache: Rc>, - registry: Option>, - compiler: Option, - rtl_lookup: Option, -} - -impl Specializer { - pub fn new( - registry: Option>, - compiler: Option, - rtl_lookup: Option, - cache: Option>>, - ) -> Self { - Self { - cache: cache.unwrap_or_else(|| Rc::new(RefCell::new(HashMap::new()))), - registry, - compiler, - rtl_lookup, - } - } - - pub fn specialize(&self, node: AnalyzedNode) -> AnalyzedNode { - self.visit_node(node) - } - - fn visit_node(&self, node: AnalyzedNode) -> AnalyzedNode { - let (new_kind, metrics) = match node.kind { - BoundKind::Call { callee, args } => { - let (new_callee, new_args, _ret_ty) = - self.specialize_call_logic(callee, args, node.ty.original.ty.clone()); - - let new_metrics = node.ty.clone(); - ( - BoundKind::Call { - callee: Rc::new(new_callee), - args: Rc::new(new_args), - }, - new_metrics, - ) - } - - BoundKind::If { - cond, - then_br, - else_br, - } => { - let cond = Rc::new(self.visit_node(cond.as_ref().clone())); - let then_br = Rc::new(self.visit_node(then_br.as_ref().clone())); - let else_br = else_br.map(|e| Rc::new(self.visit_node(e.as_ref().clone()))); - ( - BoundKind::If { - cond, - then_br, - else_br, - }, - node.ty.clone(), - ) - } - BoundKind::Block { exprs } => { - let exprs = exprs.into_iter().map(|e| Rc::new(self.visit_node(e.as_ref().clone()))).collect(); - (BoundKind::Block { exprs }, node.ty.clone()) - } - 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.as_ref().clone())); - ( - BoundKind::Lambda { - params, - upvalues, - body, - positional_count, - }, - node.ty.clone(), - ) - } - BoundKind::Define { - name, - addr, - kind, - value, - captured_by, - } => { - let value = Rc::new(self.visit_node(value.as_ref().clone())); - ( - BoundKind::Define { - name: name.clone(), - addr, - kind, - value, - captured_by: captured_by.clone(), - }, - node.ty.clone(), - ) - } - BoundKind::Set { addr, value } => { - let value = Rc::new(self.visit_node(value.as_ref().clone())); - (BoundKind::Set { addr, value }, node.ty.clone()) - } - BoundKind::Tuple { elements } => { - let elements = elements.into_iter().map(|e| Rc::new(self.visit_node(e.as_ref().clone()))).collect(); - (BoundKind::Tuple { elements }, node.ty.clone()) - } - BoundKind::Record { layout, values } => { - let values = values.into_iter().map(|v| Rc::new(self.visit_node(v.as_ref().clone()))).collect(); - (BoundKind::Record { layout, values }, node.ty.clone()) - } - BoundKind::Expansion { - original_call, - bound_expanded, - } => { - let bound_expanded = Rc::new(self.visit_node(bound_expanded.as_ref().clone())); - ( - BoundKind::Expansion { - original_call, - bound_expanded, - }, - node.ty.clone(), - ) - } - k => (k, node.ty.clone()), - }; - - Node { - identity: node.identity, - kind: new_kind, - ty: metrics, - } - } - - fn specialize_call_logic( - &self, - callee: Rc, - args: Rc, - original_ty: StaticType, - ) -> (AnalyzedNode, AnalyzedNode, StaticType) { - let new_callee = self.visit_node(callee.as_ref().clone()); - let new_args = self.visit_node(args.as_ref().clone()); - - let address = if let BoundKind::Get { addr, .. } = &new_callee.kind { - *addr - } else { - return (new_callee, new_args, original_ty); - }; - - let arg_types: Vec = - if let StaticType::Tuple(elements) = &new_args.ty.original.ty { - elements.clone() - } else { - vec![new_args.ty.original.ty.clone()] - }; - - if arg_types.iter().any(|t| matches!(t, StaticType::Any)) { - return (new_callee, new_args, original_ty); - } - - let key = MonoCacheKey { - address, - arg_types: arg_types.clone(), - }; - - if let Some((val, ret_ty)) = self.cache.borrow().get(&key) { - let specialized_callee = self.make_constant_node( - val.clone(), - StaticType::Function(Box::new(Signature { - params: StaticType::Tuple(arg_types), - ret: ret_ty.clone(), - })), - &new_callee, - ); - return (specialized_callee, new_args, ret_ty.clone()); - } - - 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) - { - self.cache - .borrow_mut() - .insert(key.clone(), (val.clone(), ret_ty.clone())); - let specialized_callee = self.make_constant_node( - val.clone(), - StaticType::Function(Box::new(Signature { - params: StaticType::Tuple(arg_types), - ret: ret_ty.clone(), - })), - &new_callee, - ); - return (specialized_callee, new_args, ret_ty); - } - - if let Some(registry) = &self.registry - && let Some(func_node) = registry.resolve_analyzed(address) - && func_node.ty.is_recursive - { - return (new_callee, new_args, original_ty); - } - - if let Some(compiler) = &self.compiler - && let Some(func_node) = self.registry.as_ref().and_then(|r| r.resolve(address)) - && let Ok((compiled_val, ret_ty)) = compiler(func_node, &arg_types) - { - self.cache - .borrow_mut() - .insert(key, (compiled_val.clone(), ret_ty.clone())); - - // Only replace the callee if the compiled value is actually a function/object. - // If it's a scalar (like 30 from folding), we DON'T fold here. - // We keep the Call but update the callee to the specialized version if it's an object. - if let Value::Object(_) | Value::Function(_) = &compiled_val { - let specialized_callee = self.make_constant_node( - compiled_val, - StaticType::Function(Box::new(Signature { - params: StaticType::Tuple(arg_types), - ret: ret_ty.clone(), - })), - &new_callee, - ); - return (specialized_callee, new_args, ret_ty); - } - } - - (new_callee, new_args, original_ty) - } - - fn make_constant_node( - &self, - val: Value, - ty: StaticType, - template: &AnalyzedNode, - ) -> AnalyzedNode { - let typed_original = Rc::new(Node { - identity: template.identity.clone(), - kind: BoundKind::Constant(val.clone()), - ty: ty.clone(), - }); - Node { - identity: template.identity.clone(), - kind: BoundKind::Constant(val), - ty: NodeMetrics { - original: typed_original, - purity: Purity::Pure, - is_recursive: false, - }, - } - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AnalyzedNode, AnalyzedPhase, IdentifierBinding, Node, NodeKind, NodeMetrics, VirtualId, +}; +use crate::ast::types::{Purity, 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, +} + +pub type CompileFunc = Rc>, &[StaticType]) -> Result<(Value, StaticType), String>>; +pub type RtlLookupFunc = Rc Option<(Value, StaticType)>>; + +pub trait FunctionRegistry { + fn resolve(&self, addr: Address) -> Option>>; + fn resolve_analyzed(&self, _addr: Address) -> Option> { + None + } +} + +pub type MonoCache = HashMap; + +pub struct Specializer { + pub cache: Rc>, + registry: Option>, + compiler: Option, + rtl_lookup: Option, +} + +impl Specializer { + pub fn new( + registry: Option>, + compiler: Option, + rtl_lookup: Option, + cache: Option>>, + ) -> Self { + Self { + cache: cache.unwrap_or_else(|| Rc::new(RefCell::new(HashMap::new()))), + registry, + compiler, + rtl_lookup, + } + } + + pub fn specialize(&self, node: AnalyzedNode) -> AnalyzedNode { + self.visit_node(node) + } + + fn visit_node(&self, node: AnalyzedNode) -> AnalyzedNode { + let (new_kind, metrics) = match node.kind { + NodeKind::Call { callee, args } => { + let (new_callee, new_args, _ret_ty) = + self.specialize_call_logic(callee, args, node.ty.original.ty.clone()); + + let new_metrics = node.ty.clone(); + ( + NodeKind::Call { + callee: Rc::new(new_callee), + args: Rc::new(new_args), + }, + new_metrics, + ) + } + + NodeKind::If { + cond, + then_br, + else_br, + } => { + let cond = Rc::new(self.visit_node(cond.as_ref().clone())); + let then_br = Rc::new(self.visit_node(then_br.as_ref().clone())); + let else_br = else_br.map(|e| Rc::new(self.visit_node(e.as_ref().clone()))); + ( + NodeKind::If { + cond, + then_br, + else_br, + }, + node.ty.clone(), + ) + } + NodeKind::Block { exprs } => { + let exprs = exprs.into_iter().map(|e| Rc::new(self.visit_node(e.as_ref().clone()))).collect(); + (NodeKind::Block { exprs }, node.ty.clone()) + } + NodeKind::Lambda { + params, + body, + info, + } => { + let params = Rc::new(self.visit_node(params.as_ref().clone())); + let body = Rc::new(self.visit_node(body.as_ref().clone())); + ( + NodeKind::Lambda { + params, + body, + info, + }, + node.ty.clone(), + ) + } + NodeKind::Def { + pattern, + value, + info, + } => { + let value = Rc::new(self.visit_node(value.as_ref().clone())); + ( + NodeKind::Def { + pattern, + value, + info, + }, + node.ty.clone(), + ) + } + NodeKind::Assign { target, value, info } => { + let value = Rc::new(self.visit_node(value.as_ref().clone())); + (NodeKind::Assign { target, value, info }, node.ty.clone()) + } + NodeKind::Tuple { elements } => { + let elements = elements.into_iter().map(|e| Rc::new(self.visit_node(e.as_ref().clone()))).collect(); + (NodeKind::Tuple { elements }, node.ty.clone()) + } + NodeKind::Record { fields, layout } => { + let fields = fields.into_iter().map(|(k, v)| (k, Rc::new(self.visit_node(v.as_ref().clone())))).collect(); + (NodeKind::Record { fields, layout }, node.ty.clone()) + } + NodeKind::Expansion { + original_call, + expanded, + } => { + let expanded = Rc::new(self.visit_node(expanded.as_ref().clone())); + ( + NodeKind::Expansion { + original_call, + expanded, + }, + node.ty.clone(), + ) + } + k => (k, node.ty.clone()), + }; + + Node { + identity: node.identity, + kind: new_kind, + ty: metrics, + } + } + + fn specialize_call_logic( + &self, + callee: Rc, + args: Rc, + original_ty: StaticType, + ) -> (AnalyzedNode, AnalyzedNode, StaticType) { + let new_callee = self.visit_node(callee.as_ref().clone()); + let new_args = self.visit_node(args.as_ref().clone()); + + let address = if let NodeKind::Identifier { + binding: IdentifierBinding::Reference(addr), + .. + } = &new_callee.kind + { + *addr + } else { + return (new_callee, new_args, original_ty); + }; + + let arg_types: Vec = + if let StaticType::Tuple(elements) = &new_args.ty.original.ty { + elements.clone() + } else { + vec![new_args.ty.original.ty.clone()] + }; + + if arg_types.iter().any(|t| matches!(t, StaticType::Any)) { + return (new_callee, new_args, original_ty); + } + + let key = MonoCacheKey { + address, + arg_types: arg_types.clone(), + }; + + if let Some((val, ret_ty)) = self.cache.borrow().get(&key) { + let specialized_callee = self.make_constant_node( + val.clone(), + StaticType::Function(Box::new(Signature { + params: StaticType::Tuple(arg_types), + ret: ret_ty.clone(), + })), + &new_callee, + ); + return (specialized_callee, new_args, ret_ty.clone()); + } + + if let Some(rtl_lookup) = &self.rtl_lookup + && let NodeKind::Identifier { symbol, .. } = &new_callee.kind + && let Some((val, ret_ty)) = rtl_lookup(&symbol.name, &arg_types) + { + self.cache + .borrow_mut() + .insert(key.clone(), (val.clone(), ret_ty.clone())); + let specialized_callee = self.make_constant_node( + val.clone(), + StaticType::Function(Box::new(Signature { + params: StaticType::Tuple(arg_types), + ret: ret_ty.clone(), + })), + &new_callee, + ); + return (specialized_callee, new_args, ret_ty); + } + + if let Some(registry) = &self.registry + && let Some(func_node) = registry.resolve_analyzed(address) + && func_node.ty.is_recursive + { + return (new_callee, new_args, original_ty); + } + + if let Some(compiler) = &self.compiler + && let Some(func_node) = self.registry.as_ref().and_then(|r| r.resolve(address)) + && let Ok((compiled_val, ret_ty)) = compiler(func_node, &arg_types) + { + self.cache + .borrow_mut() + .insert(key, (compiled_val.clone(), ret_ty.clone())); + + // Only replace the callee if the compiled value is actually a function/object. + // If it's a scalar (like 30 from folding), we DON'T fold here. + // We keep the Call but update the callee to the specialized version if it's an object. + if let Value::Object(_) | Value::Function(_) = &compiled_val { + let specialized_callee = self.make_constant_node( + compiled_val, + StaticType::Function(Box::new(Signature { + params: StaticType::Tuple(arg_types), + ret: ret_ty.clone(), + })), + &new_callee, + ); + return (specialized_callee, new_args, ret_ty); + } + } + + (new_callee, new_args, original_ty) + } + + fn make_constant_node( + &self, + val: Value, + ty: StaticType, + template: &AnalyzedNode, + ) -> AnalyzedNode { + let typed_original = Rc::new(Node { + identity: template.identity.clone(), + kind: NodeKind::Constant(val.clone()), + ty: ty.clone(), + }); + Node { + identity: template.identity.clone(), + kind: NodeKind::Constant(val), + ty: NodeMetrics { + original: typed_original, + purity: Purity::Pure, + is_recursive: false, + }, + } + } +} diff --git a/src/ast/compiler/type_checker.rs b/src/ast/compiler/type_checker.rs index a0f23b3..dc09af2 100644 --- a/src/ast/compiler/type_checker.rs +++ b/src/ast/compiler/type_checker.rs @@ -1,918 +1,1068 @@ -use crate::ast::compiler::bound_nodes::{Address, BoundKind, BoundPhase, CompilerPhase, Node, TypedNode, VirtualId}; -use crate::ast::diagnostics::Diagnostics; -use crate::ast::types::StaticType; -use std::collections::HashMap; -use std::rc::Rc; - -/// Manages the types of locals and upvalues during a single type-checking pass. -struct TypeContext<'a> { - _parent: Option<&'a TypeContext<'a>>, - /// Maps slot index -> Inferred Type - slots: HashMap, - /// Types of captured variables (passed from outer scope) - upvalue_types: Vec, - /// Access to root types for unified resolution - root_types: &'a std::cell::RefCell>, - /// The expected parameters of the current function (for 'again' validation) - current_params_ty: Option, -} - -impl<'a> TypeContext<'a> { - fn new( - _slot_count: u32, - upvalue_types: Vec, - root_types: &'a std::cell::RefCell>, - parent: Option<&'a TypeContext<'a>>, - ) -> Self { - Self { - _parent: parent, - slots: HashMap::new(), - upvalue_types, - root_types, - current_params_ty: None, - } - } - - fn get_type(&self, addr: Address) -> StaticType { - match addr { - Address::Local(slot) => self - .slots - .get(&slot.0) - .cloned() - .unwrap_or(StaticType::Any), - Address::Global(idx) => self - .root_types - .borrow() - .get(idx.0 as usize) - .cloned() - .unwrap_or(StaticType::Any), - Address::Upvalue(idx) => self - .upvalue_types - .get(idx.0 as usize) - .cloned() - .unwrap_or(StaticType::Any), - } - } - - fn set_type(&mut self, addr: Address, ty: StaticType) { - match addr { - Address::Local(slot) => { - self.slots.insert(slot.0, ty); - } - Address::Global(idx) => { - let mut rt = self.root_types.borrow_mut(); - if (idx.0 as usize) < rt.len() { - rt[idx.0 as usize] = ty; - } - } - _ => {} - } - } -} - -pub struct TypeChecker { - root_types: Rc>>, -} - -impl TypeChecker { - pub fn new(root_types: Rc>>) -> Self { - Self { root_types } - } - - pub fn check( - &self, - node: &Node, - arg_types: &[StaticType], - diag: &mut Diagnostics, - ) -> TypedNode { - self.check_node_as_bound(node, arg_types, diag) - } - - /// Allows re-checking a node from any phase as if it were a bound node. - /// This is useful for specialization where we re-type a TypedNode with more specific info. - pub fn check_node_as_bound>( - &self, - node: &Node

, - arg_types: &[StaticType], - diag: &mut Diagnostics, - ) -> TypedNode { - match &node.kind { - BoundKind::Lambda { - params, - upvalues, - body, - positional_count, - } => { - let mut upvalue_types = Vec::with_capacity(upvalues.len()); - for &_addr in upvalues { - upvalue_types.push(StaticType::Any); - } - - let root_ctx = TypeContext::new(0, vec![], &self.root_types, None); - let mut lambda_ctx = - TypeContext::new(64, upvalue_types, &self.root_types, Some(&root_ctx)); - - let arg_tuple_ty = if arg_types.is_empty() { - StaticType::Any - } else { - StaticType::Tuple(arg_types.to_vec()) - }; - - let params_typed = self.check_params( - params.as_ref(), - &arg_tuple_ty, - &mut lambda_ctx, - diag, - ); - - let body_typed = self.check_node(body, &mut lambda_ctx, diag); - let ret_ty = body_typed.ty.clone(); - let final_params_ty = params_typed.ty.clone(); - - let fn_ty = StaticType::Function(Box::new(crate::ast::types::Signature { - params: final_params_ty, - ret: ret_ty, - })); - - Node { - identity: node.identity.clone(), - kind: BoundKind::Lambda { - params: Rc::new(params_typed), - upvalues: upvalues.clone(), - body: Rc::new(body_typed), - positional_count: *positional_count, - }, - ty: fn_ty, - } - } - _ => { - let mut root_ctx = TypeContext::new(0, vec![], &self.root_types, None); - self.check_node(node, &mut root_ctx, diag) - } - } - } - - fn check_params>( - &self, - node: &Node

, - specialized_ty: &StaticType, - ctx: &mut TypeContext, - diag: &mut Diagnostics, - ) -> TypedNode { - let (kind, ty) = match &node.kind { - BoundKind::Define { - name, - addr, - kind: decl_kind, - captured_by, - .. - } => { - ctx.set_type(*addr, specialized_ty.clone()); - ( - BoundKind::Define { - name: name.clone(), - addr: *addr, - kind: *decl_kind, - value: Rc::new(Node { - identity: node.identity.clone(), - kind: BoundKind::Nop, - ty: specialized_ty.clone(), - }), - captured_by: captured_by.clone(), - }, - specialized_ty.clone(), - ) - } - BoundKind::Set { - addr, - value: _value, - } => { - ( - BoundKind::Set { - addr: *addr, - value: Rc::new(Node { - identity: node.identity.clone(), - kind: BoundKind::Nop, - ty: specialized_ty.clone(), - }), - }, - specialized_ty.clone(), - ) - } - BoundKind::Tuple { elements } => { - match specialized_ty { - StaticType::Any - | StaticType::Tuple(_) - | StaticType::Vector(_, _) - | StaticType::Matrix(_, _) - | StaticType::List(_) - | StaticType::Record(_) - | StaticType::Error => {} - _ => { - diag.push_error( - format!( - "Cannot destructure type {} as a tuple/vector", - specialized_ty - ), - Some(node.identity.clone()), - ); - return Node { - identity: node.identity.clone(), - kind: BoundKind::Error, - ty: StaticType::Error, - }; - } - } - - let mut typed_elements = Vec::new(); - let mut elem_types = Vec::new(); - - for (i, el) in elements.iter().enumerate() { - let sub_ty = match specialized_ty { - StaticType::Tuple(t) => t.get(i).cloned().unwrap_or(StaticType::Any), - StaticType::Vector(inner, _) => (**inner).clone(), - StaticType::Matrix(inner, _) => (**inner).clone(), - StaticType::List(inner) => (**inner).clone(), - StaticType::Record(layout) => layout - .fields - .get(i) - .map(|(_, ty): &(crate::ast::types::Keyword, StaticType)| ty.clone()) - .unwrap_or(StaticType::Any), - StaticType::Error => StaticType::Error, - _ => StaticType::Any, - }; - let t = self.check_params(el.as_ref(), &sub_ty, ctx, diag); - elem_types.push(t.ty.clone()); - typed_elements.push(Rc::new(t)); - } - ( - BoundKind::Tuple { - elements: typed_elements, - }, - StaticType::Tuple(elem_types), - ) - } - BoundKind::Error => (BoundKind::Error, StaticType::Error), - _ => { - diag.push_error( - "Invalid node in parameter list", - Some(node.identity.clone()), - ); - (BoundKind::Error, StaticType::Error) - } - }; - - Node { - identity: node.identity.clone(), - kind, - ty, - } - } - - fn check_node>( - &self, - node: &Node

, - ctx: &mut TypeContext, - diag: &mut Diagnostics, - ) -> TypedNode { - let (kind, ty) = match &node.kind { - BoundKind::Nop => (BoundKind::Nop, StaticType::Void), - - BoundKind::Constant(v) => { - let ty = v.static_type(); - (BoundKind::Constant(v.clone()), ty) - } - - BoundKind::Define { - name, - addr, - kind: decl_kind, - value, - captured_by, - } => { - let val_typed = self.check_node(value, ctx, diag); - let ty = val_typed.ty.clone(); - ctx.set_type(*addr, ty.clone()); - - ( - BoundKind::Define { - name: name.clone(), - addr: *addr, - kind: *decl_kind, - value: Rc::new(val_typed), - captured_by: captured_by.clone(), - }, - ty, - ) - } - - BoundKind::Get { addr, name } => { - let ty = ctx.get_type(*addr); - ( - BoundKind::Get { - addr: *addr, - name: name.clone(), - }, - ty, - ) - } - - BoundKind::FieldAccessor(k) => { - (BoundKind::FieldAccessor(*k), StaticType::FieldAccessor(*k)) - } - - BoundKind::GetField { rec, field } => { - let rec_typed = self.check_node(rec, ctx, diag); - let field_ty = match &rec_typed.ty { - StaticType::Record(layout) => { - if let Some(idx) = layout.index_of(*field) { - layout.fields[idx].1.clone() - } else { - diag.push_error( - format!("Record does not have field :{}", field.name()), - Some(rec_typed.identity.clone()), - ); - StaticType::Error - } - } - StaticType::Any => StaticType::Any, - StaticType::Error => StaticType::Error, - _ => { - diag.push_error( - format!( - "Cannot access field :{} on non-record type {}", - field.name(), - rec_typed.ty - ), - Some(rec_typed.identity.clone()), - ); - StaticType::Error - } - }; - ( - BoundKind::GetField { - rec: Rc::new(rec_typed), - field: *field, - }, - field_ty, - ) - } - - BoundKind::Set { addr, value } => { - let val_typed = self.check_node(value, ctx, diag); - let ty = val_typed.ty.clone(); - ctx.set_type(*addr, ty.clone()); - - ( - BoundKind::Set { - addr: *addr, - value: Rc::new(val_typed), - }, - ty, - ) - } - - BoundKind::Destructure { pattern, value } => { - let val_typed = self.check_node(value, ctx, diag); - let pat_typed = self.check_params(pattern.as_ref(), &val_typed.ty, ctx, diag); - let ty = val_typed.ty.clone(); - ( - BoundKind::Destructure { - pattern: Rc::new(pat_typed), - value: Rc::new(val_typed), - }, - ty, - ) - } - - BoundKind::If { - cond, - then_br, - else_br, - } => { - let cond_typed = self.check_node(cond, ctx, diag); - let then_typed = self.check_node(then_br, ctx, diag); - - let mut else_typed = None; - let mut final_ty = then_typed.ty.clone(); - - if let Some(e) = else_br { - let et = self.check_node(e, ctx, diag); - if et.ty != final_ty { - final_ty = StaticType::Any; - } - else_typed = Some(Rc::new(et)); - } else { - final_ty = StaticType::Optional(Box::new(then_typed.ty.clone())); - } - - ( - BoundKind::If { - cond: Rc::new(cond_typed), - then_br: Rc::new(then_typed), - else_br: else_typed, - }, - final_ty, - ) - } - - BoundKind::Pipe { inputs, lambda, .. } => { - let mut typed_inputs = Vec::with_capacity(inputs.len()); - let mut arg_types = Vec::with_capacity(inputs.len()); - for input in inputs { - let typed_input = self.check_node(input, ctx, diag); - let arg_ty = if let StaticType::Series(inner) = &typed_input.ty { - *inner.clone() - } else if let StaticType::Stream(inner) = &typed_input.ty { - *inner.clone() - } else { - StaticType::Any - }; - arg_types.push(arg_ty); - typed_inputs.push(Rc::new(typed_input)); - } - - let typed_lambda = self.check_node_as_bound(lambda.as_ref(), &arg_types, diag); - - let ret_ty = if let StaticType::Function(sig) = &typed_lambda.ty { - if let StaticType::Optional(inner) = &sig.ret { - *inner.clone() - } else { - sig.ret.clone() - } - } else { - StaticType::Any - }; - ( - BoundKind::Pipe { - inputs: typed_inputs, - lambda: Rc::new(typed_lambda), - out_type: ret_ty.clone(), - }, - StaticType::Stream(Box::new(ret_ty)), - ) - } - - BoundKind::Block { exprs } => { - let mut typed_exprs = Vec::new(); - let mut last_ty = StaticType::Void; - - for e in exprs { - let t = self.check_node(e, ctx, diag); - last_ty = t.ty.clone(); - typed_exprs.push(Rc::new(t)); - } - - (BoundKind::Block { exprs: typed_exprs }, last_ty) - } - - BoundKind::Lambda { - params, - upvalues, - body, - positional_count, - } => { - let mut upvalue_types = Vec::with_capacity(upvalues.len()); - for &addr in upvalues { - upvalue_types.push(ctx.get_type(addr)); - } - - let mut lambda_ctx = - TypeContext::new(64, upvalue_types, ctx.root_types, Some(ctx)); - - let params_typed = self.check_params( - params.as_ref(), - &StaticType::Any, - &mut lambda_ctx, - diag, - ); - - lambda_ctx.current_params_ty = Some(params_typed.ty.clone()); - - let body_typed = self.check_node(body, &mut lambda_ctx, diag); - let ret_ty = body_typed.ty.clone(); - - let fn_ty = StaticType::Function(Box::new(crate::ast::types::Signature { - params: params_typed.ty.clone(), - ret: ret_ty, - })); - - ( - BoundKind::Lambda { - params: Rc::new(params_typed), - upvalues: upvalues.clone(), - body: Rc::new(body_typed), - positional_count: *positional_count, - }, - fn_ty, - ) - } - - BoundKind::Call { callee, args } => { - let callee_typed = self.check_node(callee, ctx, diag); - - let args_typed = if let BoundKind::Tuple { elements } = &args.kind { - let mut typed_elements = Vec::new(); - let mut elem_types = Vec::new(); - for e in elements { - let t = self.check_node(e, ctx, diag); - elem_types.push(t.ty.clone()); - typed_elements.push(Rc::new(t)); - } - Node { - identity: args.identity.clone(), - kind: BoundKind::Tuple { - elements: typed_elements, - }, - ty: StaticType::Tuple(elem_types), - } - } else { - self.check_node(args, ctx, diag) - }; - - let ret_ty = match callee_typed.ty.resolve_call(&args_typed.ty) { - Some(ty) => ty, - None => { - diag.push_error( - format!( - "Invalid arguments for function call. Expected {}, got {}", - callee_typed.ty, args_typed.ty - ), - Some(node.identity.clone()), - ); - StaticType::Error - } - }; - - ( - BoundKind::Call { - callee: Rc::new(callee_typed), - args: Rc::new(args_typed), - }, - ret_ty, - ) - } - - BoundKind::Again { args } => { - let args_typed = if let BoundKind::Tuple { elements } = &args.kind { - let mut typed_elements = Vec::new(); - let mut elem_types = Vec::new(); - for e in elements { - let t = self.check_node(e, ctx, diag); - elem_types.push(t.ty.clone()); - typed_elements.push(Rc::new(t)); - } - Node { - identity: args.identity.clone(), - kind: BoundKind::Tuple { - elements: typed_elements, - }, - ty: StaticType::Tuple(elem_types), - } - } else { - self.check_node(args, ctx, diag) - }; - - if let Some(expected_ty) = &ctx.current_params_ty - && !expected_ty.is_assignable_from(&args_typed.ty) - { - diag.push_error( - format!( - "Type mismatch in 'again' call: expected {}, but got {}", - expected_ty, args_typed.ty - ), - Some(node.identity.clone()), - ); - } - - ( - BoundKind::Again { - args: Rc::new(args_typed), - }, - StaticType::Any, - ) - } - - BoundKind::Tuple { elements } => { - let mut typed_elements = Vec::new(); - for e in elements { - typed_elements.push(Rc::new(self.check_node(e, ctx, diag))); - } - - let ty = if typed_elements.is_empty() { - StaticType::Vector(Box::new(StaticType::Any), 0) - } else { - let first_ty = &typed_elements[0].ty; - let all_same = typed_elements.iter().all(|e| e.ty == *first_ty); - - if all_same { - match first_ty { - StaticType::Vector(inner, len) => { - StaticType::Matrix(inner.clone(), vec![typed_elements.len(), *len]) - } - StaticType::Matrix(inner, shape) => { - let mut new_shape = vec![typed_elements.len()]; - new_shape.extend(shape); - StaticType::Matrix(inner.clone(), new_shape) - } - _ => { - StaticType::Vector(Box::new(first_ty.clone()), typed_elements.len()) - } - } - } else { - StaticType::Tuple(typed_elements.iter().map(|e| e.ty.clone()).collect()) - } - }; - - ( - BoundKind::Tuple { - elements: typed_elements, - }, - ty, - ) - } - - BoundKind::Record { layout, values } => { - let mut typed_values = Vec::with_capacity(values.len()); - let mut fields_ty = Vec::with_capacity(values.len()); - - for (i, v) in values.iter().enumerate() { - let vt = self.check_node(v, ctx, diag); - fields_ty.push((layout.fields[i].0, vt.ty.clone())); - typed_values.push(Rc::new(vt)); - } - - let new_layout = crate::ast::types::RecordLayout::get_or_create(fields_ty); - ( - BoundKind::Record { - layout: new_layout.clone(), - values: typed_values, - }, - StaticType::Record(new_layout), - ) - } - - BoundKind::Expansion { - original_call, - bound_expanded, - } => { - let expanded_typed = self.check_node(bound_expanded, ctx, diag); - let ty = expanded_typed.ty.clone(); - ( - BoundKind::Expansion { - original_call: original_call.clone(), - bound_expanded: Rc::new(expanded_typed), - }, - ty, - ) - } - - BoundKind::Extension(_ext) => { - diag.push_error( - format!( - "TypeChecking for extension '{}' not implemented", - _ext.display_name() - ), - Some(node.identity.clone()), - ); - (BoundKind::Error, StaticType::Error) - } - BoundKind::Error => (BoundKind::Error, StaticType::Error), - }; - - Node { - identity: node.identity.clone(), - kind, - ty, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::ast::types::StaticType; - - fn check_source(source: &str) -> TypedNode { - let env = crate::ast::environment::Environment::new(); - env.compile(source).into_result().unwrap() - } - - #[test] - fn test_destructuring_scalar_error() { - let env = crate::ast::environment::Environment::new(); - let result = env.compile("(def [x] 5)").into_result(); - assert!(result.is_err()); - assert_eq!( - result.unwrap_err(), - "Statement 'def' cannot be used as an expression.\nCannot destructure type int as a tuple/vector" - ); - } - - #[test] - fn test_call_argument_mismatch() { - let env = crate::ast::environment::Environment::new(); - // fn takes 1 arg, called with 0 - let result = env.compile("(do (def f (fn [x] x)) (f))").into_result(); - assert!(result.is_err()); - assert!( - result - .unwrap_err() - .contains("Invalid arguments for function call") - ); - } - - #[test] - fn test_destructuring_vector_args() { - let env = crate::ast::environment::Environment::new(); - // ((fn [[x y]] (+ x y)) [10 20]) -> 30 - let result = env.compile("((fn [[x y]] (+ x y)) [10 20])").into_result(); - assert!(result.is_ok()); - assert_eq!(get_ret_type(&result.unwrap()), StaticType::Int); - } - - fn get_ret_type(node: &TypedNode) -> StaticType { - if let StaticType::Function(sig) = &node.ty { - sig.ret.clone() - } else { - node.ty.clone() - } - } - - #[test] - fn test_inference_constants() { - assert_eq!(get_ret_type(&check_source("10")), StaticType::Int); - assert_eq!(get_ret_type(&check_source("10.5")), StaticType::Float); - assert_eq!(get_ret_type(&check_source("true")), StaticType::Bool); - assert_eq!(get_ret_type(&check_source("\"hello\"")), StaticType::Text); - } - - #[test] - fn test_inference_variable_propagation() { - // (do (def x 10) x) -> The last 'x' must be Int - let typed = check_source("(do (def x 10) x)"); - // Outer is Lambda, Body is Block - if let BoundKind::Lambda { body, .. } = &typed.kind { - if let BoundKind::Block { exprs } = &body.kind { - let last_expr = exprs.last().unwrap(); - assert_eq!( - last_expr.ty, - StaticType::Int, - "Variable 'x' should be inferred as Int" - ); - } else { - panic!("Expected block in lambda body"); - } - } else { - panic!("Expected Lambda wrapper"); - } - } - - #[test] - fn test_inference_block_type() { - // Block type = last expression type - assert_eq!(get_ret_type(&check_source("(do 1 2.5)")), StaticType::Float); - assert_eq!( - get_ret_type(&check_source("(do 1.5 \"test\")")), - StaticType::Text - ); - } - - #[test] - fn test_inference_lambda_return() { - // (fn [a] 10) -> fn(any) -> Int - // Since it's already a Lambda, it's NOT wrapped further. - let typed = check_source("(fn [a] 10)"); - if let StaticType::Function(sig) = &typed.ty { - assert_eq!(sig.ret, StaticType::Int); - } else { - panic!("Expected function type, got {:?}", typed.ty); - } - - // Nested: (fn [] (do 1 2.5)) -> fn() -> Float - let typed_nested = check_source("(fn [] (do 1 2.5))"); - if let StaticType::Function(sig) = &typed_nested.ty { - assert_eq!(sig.ret, StaticType::Float); - } else { - panic!("Expected function type"); - } - } - - #[test] - fn test_inference_assignment_updates_type() { - // (do (def x 10) (assign x 20.5) x) -> x becomes Float after assignment - let typed = check_source("(do (def x 10) (assign x 20.5) x)"); - if let BoundKind::Lambda { body, .. } = &typed.kind { - if let BoundKind::Block { exprs } = &body.kind { - let last_expr = exprs.last().unwrap(); - assert_eq!( - last_expr.ty, - StaticType::Float, - "Variable 'x' should be specialized to Float after assignment" - ); - } else { - panic!("Expected block"); - } - } else { - panic!("Expected Lambda"); - } - } - - #[test] - fn test_operator_overloading_inference() { - assert_eq!(get_ret_type(&check_source("(+ 1 2)")), StaticType::Int); - assert_eq!( - get_ret_type(&check_source("(+ 1.0 2.0)")), - StaticType::Float - ); - assert_eq!( - get_ret_type(&check_source("(+ \"a\" \"b\")")), - StaticType::Text - ); - assert_eq!(get_ret_type(&check_source("(/ 1 2)")), StaticType::Float); - } - - #[test] - fn test_datetime_inference() { - // date("2023-01-01") -> DateTime - assert_eq!( - get_ret_type(&check_source("(date \"2023-01-01\")")), - StaticType::DateTime - ); - // DateTime + Int -> DateTime - assert_eq!( - get_ret_type(&check_source("(+ (date \"2023-01-01\") 86400000)")), - StaticType::DateTime - ); - // DateTime - DateTime -> Int (Duration) - assert_eq!( - get_ret_type(&check_source( - "(- (date \"2023-01-02\") (date \"2023-01-01\"))" - )), - StaticType::Int - ); - // DateTime comparison -> Bool - assert_eq!( - get_ret_type(&check_source( - "(> (date \"2023-01-02\") (date \"2023-01-01\"))" - )), - StaticType::Bool - ); - } - - #[test] - fn test_inference_tuple_vector_matrix() { - // Heterogeneous -> Tuple - assert_eq!( - get_ret_type(&check_source("[1 3.14 \"text\"]")), - StaticType::Tuple(vec![StaticType::Int, StaticType::Float, StaticType::Text]) - ); - - // Homogeneous -> Vector - assert_eq!( - get_ret_type(&check_source("[10 20 30]")), - StaticType::Vector(Box::new(StaticType::Int), 3) - ); - - // Nested Homogeneous -> Matrix - assert_eq!( - get_ret_type(&check_source("[[1 2] [3 4]]")), - StaticType::Matrix(Box::new(StaticType::Int), vec![2, 2]) - ); - - // Deep Matrix - assert_eq!( - get_ret_type(&check_source("[[[1 2]] [[3 4]]]")), - StaticType::Matrix(Box::new(StaticType::Int), vec![2, 1, 2]) - ); - - // Shape mismatch -> Tuple of Vectors - let mixed = get_ret_type(&check_source("[[1 2] [3 4 5]]")); - if let StaticType::Tuple(elements) = mixed { - assert_eq!( - elements[0], - StaticType::Vector(Box::new(StaticType::Int), 2) - ); - assert_eq!( - elements[1], - StaticType::Vector(Box::new(StaticType::Int), 3) - ); - } else { - panic!("Expected Tuple for shape mismatch, got {:?}", mixed); - } - } - - #[test] - fn test_inference_record() { - use crate::ast::types::Keyword; - let typed = check_source("{:x 1 :y 0.3}"); - let ty = get_ret_type(&typed); - if let StaticType::Record(layout) = ty { - assert_eq!(layout.fields.len(), 2); - assert_eq!(layout.fields[0], (Keyword::intern("x"), StaticType::Int)); - assert_eq!(layout.fields[1], (Keyword::intern("y"), StaticType::Float)); - } else { - panic!("Expected Record, got {:?}", ty); - } - } -} +use crate::ast::compiler::bound_nodes::{ + Address, AssignBinding, BoundLike, DefBinding, IdentifierBinding, LambdaBinding, Node, + NodeKind, TypedNode, TypedPhase, VirtualId, +}; +use crate::ast::diagnostics::Diagnostics; +use crate::ast::types::StaticType; +use std::collections::HashMap; +use std::rc::Rc; + +/// Manages the types of locals and upvalues during a single type-checking pass. +struct TypeContext<'a> { + _parent: Option<&'a TypeContext<'a>>, + /// Maps slot index -> Inferred Type + slots: HashMap, + /// Types of captured variables (passed from outer scope) + upvalue_types: Vec, + /// Access to root types for unified resolution + root_types: &'a std::cell::RefCell>, + /// The expected parameters of the current function (for 'again' validation) + current_params_ty: Option, +} + +impl<'a> TypeContext<'a> { + fn new( + _slot_count: u32, + upvalue_types: Vec, + root_types: &'a std::cell::RefCell>, + parent: Option<&'a TypeContext<'a>>, + ) -> Self { + Self { + _parent: parent, + slots: HashMap::new(), + upvalue_types, + root_types, + current_params_ty: None, + } + } + + fn get_type(&self, addr: Address) -> StaticType { + match addr { + Address::Local(slot) => self + .slots + .get(&slot.0) + .cloned() + .unwrap_or(StaticType::Any), + Address::Global(idx) => self + .root_types + .borrow() + .get(idx.0 as usize) + .cloned() + .unwrap_or(StaticType::Any), + Address::Upvalue(idx) => self + .upvalue_types + .get(idx.0 as usize) + .cloned() + .unwrap_or(StaticType::Any), + } + } + + fn set_type(&mut self, addr: Address, ty: StaticType) { + match addr { + Address::Local(slot) => { + self.slots.insert(slot.0, ty); + } + Address::Global(idx) => { + let mut rt = self.root_types.borrow_mut(); + if (idx.0 as usize) < rt.len() { + rt[idx.0 as usize] = ty; + } + } + _ => {} + } + } +} + +pub struct TypeChecker { + root_types: Rc>>, +} + +impl TypeChecker { + pub fn new(root_types: Rc>>) -> Self { + Self { root_types } + } + + pub fn check( + &self, + node: &Node, + arg_types: &[StaticType], + diag: &mut Diagnostics, + ) -> TypedNode { + self.check_node_as_bound(node, arg_types, diag) + } + + /// Allows re-checking a node from any phase as if it were a bound node. + /// This is useful for specialization where we re-type a TypedNode with more specific info. + pub fn check_node_as_bound( + &self, + node: &Node

, + arg_types: &[StaticType], + diag: &mut Diagnostics, + ) -> TypedNode { + match &node.kind { + NodeKind::Lambda { + params, + body, + info, + } => { + let upvalues = &info.upvalues; + let positional_count = info.positional_count; + + let mut upvalue_types = Vec::with_capacity(upvalues.len()); + for &_addr in upvalues { + upvalue_types.push(StaticType::Any); + } + + let root_ctx = TypeContext::new(0, vec![], &self.root_types, None); + let mut lambda_ctx = + TypeContext::new(64, upvalue_types, &self.root_types, Some(&root_ctx)); + + let arg_tuple_ty = if arg_types.is_empty() { + StaticType::Any + } else { + StaticType::Tuple(arg_types.to_vec()) + }; + + let params_typed = self.check_params( + params.as_ref(), + &arg_tuple_ty, + &mut lambda_ctx, + diag, + ); + + let body_typed = self.check_node(body, &mut lambda_ctx, diag); + let ret_ty = body_typed.ty.clone(); + let final_params_ty = params_typed.ty.clone(); + + let fn_ty = StaticType::Function(Box::new(crate::ast::types::Signature { + params: final_params_ty, + ret: ret_ty, + })); + + Node { + identity: node.identity.clone(), + kind: NodeKind::Lambda { + params: Rc::new(params_typed), + body: Rc::new(body_typed), + info: LambdaBinding { + upvalues: upvalues.clone(), + positional_count, + }, + }, + ty: fn_ty, + } + } + _ => { + let mut root_ctx = TypeContext::new(0, vec![], &self.root_types, None); + self.check_node(node, &mut root_ctx, diag) + } + } + } + + fn check_params( + &self, + node: &Node

, + specialized_ty: &StaticType, + ctx: &mut TypeContext, + diag: &mut Diagnostics, + ) -> TypedNode { + let (kind, ty): (NodeKind, StaticType) = match &node.kind { + NodeKind::Def { + pattern, + info, + .. + } => { + if let NodeKind::Identifier { + symbol, + binding: IdentifierBinding::Declaration { addr, kind: decl_kind }, + } = &pattern.kind + { + ctx.set_type(*addr, specialized_ty.clone()); + ( + NodeKind::Def { + pattern: Rc::new(Node { + identity: pattern.identity.clone(), + kind: NodeKind::Identifier { + symbol: symbol.clone(), + binding: IdentifierBinding::Declaration { + addr: *addr, + kind: *decl_kind, + }, + }, + ty: specialized_ty.clone(), + }), + value: Rc::new(Node { + identity: node.identity.clone(), + kind: NodeKind::Nop, + ty: specialized_ty.clone(), + }), + info: DefBinding { + captured_by: info.captured_by.clone(), + }, + }, + specialized_ty.clone(), + ) + } else { + // Destructuring def in params — fall through to tuple handling + // if pattern is a Tuple, handle elements + if let NodeKind::Tuple { elements } = &pattern.kind { + return self.check_params_tuple(node, elements, specialized_ty, ctx, diag); + } + diag.push_error( + "Invalid pattern in parameter definition", + Some(node.identity.clone()), + ); + (NodeKind::Error, StaticType::Error) + } + } + NodeKind::Assign { + info, + .. + } => { + ( + NodeKind::Assign { + target: Rc::new(Node { + identity: node.identity.clone(), + kind: NodeKind::Nop, + ty: specialized_ty.clone(), + }), + value: Rc::new(Node { + identity: node.identity.clone(), + kind: NodeKind::Nop, + ty: specialized_ty.clone(), + }), + info: AssignBinding { + addr: info.addr, + }, + }, + specialized_ty.clone(), + ) + } + NodeKind::Identifier { + symbol, + binding: IdentifierBinding::Declaration { addr, kind: decl_kind }, + } => { + ctx.set_type(*addr, specialized_ty.clone()); + ( + NodeKind::Identifier { + symbol: symbol.clone(), + binding: IdentifierBinding::Declaration { + addr: *addr, + kind: *decl_kind, + }, + }, + specialized_ty.clone(), + ) + } + NodeKind::Tuple { elements } => { + return self.check_params_tuple(node, elements, specialized_ty, ctx, diag); + } + NodeKind::Nop => (NodeKind::Nop, StaticType::Void), + NodeKind::Error => (NodeKind::Error, StaticType::Error), + _ => { + diag.push_error( + "Invalid node in parameter list", + Some(node.identity.clone()), + ); + (NodeKind::Error, StaticType::Error) + } + }; + + Node { + identity: node.identity.clone(), + kind, + ty, + } + } + + fn check_params_tuple( + &self, + node: &Node

, + elements: &[Rc>], + specialized_ty: &StaticType, + ctx: &mut TypeContext, + diag: &mut Diagnostics, + ) -> TypedNode { + match specialized_ty { + StaticType::Any + | StaticType::Tuple(_) + | StaticType::Vector(_, _) + | StaticType::Matrix(_, _) + | StaticType::List(_) + | StaticType::Record(_) + | StaticType::Error => {} + _ => { + diag.push_error( + format!( + "Cannot destructure type {} as a tuple/vector", + specialized_ty + ), + Some(node.identity.clone()), + ); + return Node { + identity: node.identity.clone(), + kind: NodeKind::Error, + ty: StaticType::Error, + }; + } + } + + let mut typed_elements = Vec::new(); + let mut elem_types = Vec::new(); + + for (i, el) in elements.iter().enumerate() { + let sub_ty = match specialized_ty { + StaticType::Tuple(t) => t.get(i).cloned().unwrap_or(StaticType::Any), + StaticType::Vector(inner, _) => (**inner).clone(), + StaticType::Matrix(inner, _) => (**inner).clone(), + StaticType::List(inner) => (**inner).clone(), + StaticType::Record(layout) => layout + .fields + .get(i) + .map(|(_, ty): &(crate::ast::types::Keyword, StaticType)| ty.clone()) + .unwrap_or(StaticType::Any), + StaticType::Error => StaticType::Error, + _ => StaticType::Any, + }; + let t = self.check_params(el.as_ref(), &sub_ty, ctx, diag); + elem_types.push(t.ty.clone()); + typed_elements.push(Rc::new(t)); + } + Node { + identity: node.identity.clone(), + kind: NodeKind::Tuple { + elements: typed_elements, + }, + ty: StaticType::Tuple(elem_types), + } + } + + fn check_node( + &self, + node: &Node

, + ctx: &mut TypeContext, + diag: &mut Diagnostics, + ) -> TypedNode { + let (kind, ty): (NodeKind, StaticType) = match &node.kind { + NodeKind::Nop => (NodeKind::Nop, StaticType::Void), + + NodeKind::Constant(v) => { + let ty = v.static_type(); + (NodeKind::Constant(v.clone()), ty) + } + + NodeKind::Def { + pattern, + value, + info, + } => { + let val_typed = self.check_node(value, ctx, diag); + let ty = val_typed.ty.clone(); + + // Extract addr from pattern to register the type + if let NodeKind::Identifier { + binding: IdentifierBinding::Declaration { addr, .. }, + .. + } = &pattern.kind + { + ctx.set_type(*addr, ty.clone()); + } + + // For destructuring defs, check params on the pattern + if let NodeKind::Tuple { .. } = &pattern.kind { + let pat_typed = self.check_params(pattern.as_ref(), &val_typed.ty, ctx, diag); + let ty = val_typed.ty.clone(); + return Node { + identity: node.identity.clone(), + kind: NodeKind::Def { + pattern: Rc::new(pat_typed), + value: Rc::new(val_typed), + info: DefBinding { + captured_by: info.captured_by.clone(), + }, + }, + ty, + }; + } + + // Simple def — reconstruct the pattern node with the new type + let new_pattern: TypedNode = Node { + identity: pattern.identity.clone(), + kind: match &pattern.kind { + NodeKind::Identifier { symbol, binding } => NodeKind::Identifier { + symbol: symbol.clone(), + binding: match binding { + IdentifierBinding::Declaration { addr, kind: decl_kind } => { + IdentifierBinding::Declaration { + addr: *addr, + kind: *decl_kind, + } + } + IdentifierBinding::Reference(addr) => { + IdentifierBinding::Reference(*addr) + } + }, + }, + _ => NodeKind::Error, + }, + ty: ty.clone(), + }; + + ( + NodeKind::Def { + pattern: Rc::new(new_pattern), + value: Rc::new(val_typed), + info: DefBinding { + captured_by: info.captured_by.clone(), + }, + }, + ty, + ) + } + + NodeKind::Identifier { symbol, binding } => { + if let IdentifierBinding::Reference(addr) = binding { + let ty = ctx.get_type(*addr); + ( + NodeKind::Identifier { + symbol: symbol.clone(), + binding: IdentifierBinding::Reference(*addr), + }, + ty, + ) + } else if let IdentifierBinding::Declaration { addr, kind: decl_kind } = binding { + let ty = ctx.get_type(*addr); + ( + NodeKind::Identifier { + symbol: symbol.clone(), + binding: IdentifierBinding::Declaration { + addr: *addr, + kind: *decl_kind, + }, + }, + ty, + ) + } else { + (NodeKind::Error, StaticType::Error) + } + } + + NodeKind::FieldAccessor(k) => { + (NodeKind::FieldAccessor(*k), StaticType::FieldAccessor(*k)) + } + + NodeKind::GetField { rec, field } => { + let rec_typed = self.check_node(rec, ctx, diag); + let field_ty = match &rec_typed.ty { + StaticType::Record(layout) => { + if let Some(idx) = layout.index_of(*field) { + layout.fields[idx].1.clone() + } else { + diag.push_error( + format!("Record does not have field :{}", field.name()), + Some(rec_typed.identity.clone()), + ); + StaticType::Error + } + } + StaticType::Any => StaticType::Any, + StaticType::Error => StaticType::Error, + _ => { + diag.push_error( + format!( + "Cannot access field :{} on non-record type {}", + field.name(), + rec_typed.ty + ), + Some(rec_typed.identity.clone()), + ); + StaticType::Error + } + }; + ( + NodeKind::GetField { + rec: Rc::new(rec_typed), + field: *field, + }, + field_ty, + ) + } + + NodeKind::Assign { target, value, info } => { + let val_typed = self.check_node(value, ctx, diag); + let ty = val_typed.ty.clone(); + if let Some(addr) = info.addr { + ctx.set_type(addr, ty.clone()); + } + + // For destructuring assigns (addr = None), preserve the target pattern + // so the VM can unpack values. For simple assigns, the target is unused. + let target_typed = if info.addr.is_none() { + Rc::new(self.check_node(target, ctx, diag)) + } else { + Rc::new(Node { + identity: node.identity.clone(), + kind: NodeKind::Nop, + ty: ty.clone(), + }) + }; + + ( + NodeKind::Assign { + target: target_typed, + value: Rc::new(val_typed), + info: AssignBinding { + addr: info.addr, + }, + }, + ty, + ) + } + + NodeKind::If { + cond, + then_br, + else_br, + } => { + let cond_typed = self.check_node(cond, ctx, diag); + let then_typed = self.check_node(then_br, ctx, diag); + + let mut else_typed = None; + let mut final_ty = then_typed.ty.clone(); + + if let Some(e) = else_br { + let et = self.check_node(e, ctx, diag); + if et.ty != final_ty { + final_ty = StaticType::Any; + } + else_typed = Some(Rc::new(et)); + } else { + final_ty = StaticType::Optional(Box::new(then_typed.ty.clone())); + } + + ( + NodeKind::If { + cond: Rc::new(cond_typed), + then_br: Rc::new(then_typed), + else_br: else_typed, + }, + final_ty, + ) + } + + NodeKind::Pipe { inputs, lambda } => { + let mut typed_inputs = Vec::with_capacity(inputs.len()); + let mut arg_types = Vec::with_capacity(inputs.len()); + for input in inputs { + let typed_input = self.check_node(input, ctx, diag); + let arg_ty = if let StaticType::Series(inner) = &typed_input.ty { + *inner.clone() + } else if let StaticType::Stream(inner) = &typed_input.ty { + *inner.clone() + } else { + StaticType::Any + }; + arg_types.push(arg_ty); + typed_inputs.push(Rc::new(typed_input)); + } + + let typed_lambda = self.check_node_as_bound(lambda.as_ref(), &arg_types, diag); + + let ret_ty = if let StaticType::Function(sig) = &typed_lambda.ty { + if let StaticType::Optional(inner) = &sig.ret { + *inner.clone() + } else { + sig.ret.clone() + } + } else { + StaticType::Any + }; + ( + NodeKind::Pipe { + inputs: typed_inputs, + lambda: Rc::new(typed_lambda), + }, + StaticType::Stream(Box::new(ret_ty)), + ) + } + + NodeKind::Block { exprs } => { + let mut typed_exprs = Vec::new(); + let mut last_ty = StaticType::Void; + + for e in exprs { + let t = self.check_node(e, ctx, diag); + last_ty = t.ty.clone(); + typed_exprs.push(Rc::new(t)); + } + + (NodeKind::Block { exprs: typed_exprs }, last_ty) + } + + NodeKind::Lambda { + params, + body, + info, + } => { + let upvalues = &info.upvalues; + let positional_count = info.positional_count; + + let mut upvalue_types = Vec::with_capacity(upvalues.len()); + for &addr in upvalues { + upvalue_types.push(ctx.get_type(addr)); + } + + let mut lambda_ctx = + TypeContext::new(64, upvalue_types, ctx.root_types, Some(ctx)); + + let params_typed = self.check_params( + params.as_ref(), + &StaticType::Any, + &mut lambda_ctx, + diag, + ); + + lambda_ctx.current_params_ty = Some(params_typed.ty.clone()); + + let body_typed = self.check_node(body, &mut lambda_ctx, diag); + let ret_ty = body_typed.ty.clone(); + + let fn_ty = StaticType::Function(Box::new(crate::ast::types::Signature { + params: params_typed.ty.clone(), + ret: ret_ty, + })); + + ( + NodeKind::Lambda { + params: Rc::new(params_typed), + body: Rc::new(body_typed), + info: LambdaBinding { + upvalues: upvalues.clone(), + positional_count, + }, + }, + fn_ty, + ) + } + + NodeKind::Call { callee, args } => { + let callee_typed = self.check_node(callee, ctx, diag); + + let args_typed = if let NodeKind::Tuple { elements } = &args.kind { + let mut typed_elements = Vec::new(); + let mut elem_types = Vec::new(); + for e in elements { + let t = self.check_node(e, ctx, diag); + elem_types.push(t.ty.clone()); + typed_elements.push(Rc::new(t)); + } + Node { + identity: args.identity.clone(), + kind: NodeKind::Tuple { + elements: typed_elements, + }, + ty: StaticType::Tuple(elem_types), + } + } else { + self.check_node(args, ctx, diag) + }; + + let ret_ty = match callee_typed.ty.resolve_call(&args_typed.ty) { + Some(ty) => ty, + None => { + diag.push_error( + format!( + "Invalid arguments for function call. Expected {}, got {}", + callee_typed.ty, args_typed.ty + ), + Some(node.identity.clone()), + ); + StaticType::Error + } + }; + + ( + NodeKind::Call { + callee: Rc::new(callee_typed), + args: Rc::new(args_typed), + }, + ret_ty, + ) + } + + NodeKind::Again { args } => { + let args_typed = if let NodeKind::Tuple { elements } = &args.kind { + let mut typed_elements = Vec::new(); + let mut elem_types = Vec::new(); + for e in elements { + let t = self.check_node(e, ctx, diag); + elem_types.push(t.ty.clone()); + typed_elements.push(Rc::new(t)); + } + Node { + identity: args.identity.clone(), + kind: NodeKind::Tuple { + elements: typed_elements, + }, + ty: StaticType::Tuple(elem_types), + } + } else { + self.check_node(args, ctx, diag) + }; + + if let Some(expected_ty) = &ctx.current_params_ty + && !expected_ty.is_assignable_from(&args_typed.ty) + { + diag.push_error( + format!( + "Type mismatch in 'again' call: expected {}, but got {}", + expected_ty, args_typed.ty + ), + Some(node.identity.clone()), + ); + } + + ( + NodeKind::Again { + args: Rc::new(args_typed), + }, + StaticType::Any, + ) + } + + NodeKind::Tuple { elements } => { + let mut typed_elements = Vec::new(); + for e in elements { + typed_elements.push(Rc::new(self.check_node(e, ctx, diag))); + } + + let ty = if typed_elements.is_empty() { + StaticType::Vector(Box::new(StaticType::Any), 0) + } else { + let first_ty = &typed_elements[0].ty; + let all_same = typed_elements.iter().all(|e| e.ty == *first_ty); + + if all_same { + match first_ty { + StaticType::Vector(inner, len) => { + StaticType::Matrix(inner.clone(), vec![typed_elements.len(), *len]) + } + StaticType::Matrix(inner, shape) => { + let mut new_shape = vec![typed_elements.len()]; + new_shape.extend(shape); + StaticType::Matrix(inner.clone(), new_shape) + } + _ => { + StaticType::Vector(Box::new(first_ty.clone()), typed_elements.len()) + } + } + } else { + StaticType::Tuple(typed_elements.iter().map(|e| e.ty.clone()).collect()) + } + }; + + ( + NodeKind::Tuple { + elements: typed_elements, + }, + ty, + ) + } + + NodeKind::Record { fields, layout } => { + let mut typed_fields = Vec::with_capacity(fields.len()); + let mut fields_ty = Vec::with_capacity(fields.len()); + + for (i, (key_node, val_node)) in fields.iter().enumerate() { + let kt = self.check_node(key_node, ctx, diag); + let vt = self.check_node(val_node, ctx, diag); + fields_ty.push((layout.fields[i].0, vt.ty.clone())); + typed_fields.push((Rc::new(kt), Rc::new(vt))); + } + + let new_layout = crate::ast::types::RecordLayout::get_or_create(fields_ty); + ( + NodeKind::Record { + fields: typed_fields, + layout: new_layout.clone(), + }, + StaticType::Record(new_layout), + ) + } + + NodeKind::Expansion { + original_call, + expanded, + } => { + let expanded_typed = self.check_node(expanded, ctx, diag); + let ty = expanded_typed.ty.clone(); + ( + NodeKind::Expansion { + original_call: original_call.clone(), + expanded: Rc::new(expanded_typed), + }, + ty, + ) + } + + NodeKind::Extension(_ext) => { + diag.push_error( + format!( + "TypeChecking for extension '{}' not implemented", + _ext.display_name() + ), + Some(node.identity.clone()), + ); + (NodeKind::Error, StaticType::Error) + } + NodeKind::Error => (NodeKind::Error, StaticType::Error), + + // Syntax-only variants should not appear in bound phases + NodeKind::MacroDecl { .. } + | NodeKind::Template(_) + | NodeKind::Placeholder(_) + | NodeKind::Splice(_) => { + diag.push_error( + "Unexpected syntax-only node in type checking", + Some(node.identity.clone()), + ); + (NodeKind::Error, StaticType::Error) + } + }; + + Node { + identity: node.identity.clone(), + kind, + ty, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ast::types::StaticType; + + fn check_source(source: &str) -> TypedNode { + let env = crate::ast::environment::Environment::new(); + env.compile(source).into_result().unwrap() + } + + #[test] + fn test_destructuring_scalar_error() { + let env = crate::ast::environment::Environment::new(); + let result = env.compile("(def [x] 5)").into_result(); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err(), + "Statement 'def' cannot be used as an expression.\nCannot destructure type int as a tuple/vector" + ); + } + + #[test] + fn test_call_argument_mismatch() { + let env = crate::ast::environment::Environment::new(); + // fn takes 1 arg, called with 0 + let result = env.compile("(do (def f (fn [x] x)) (f))").into_result(); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .contains("Invalid arguments for function call") + ); + } + + #[test] + fn test_destructuring_vector_args() { + let env = crate::ast::environment::Environment::new(); + // ((fn [[x y]] (+ x y)) [10 20]) -> 30 + let result = env.compile("((fn [[x y]] (+ x y)) [10 20])").into_result(); + assert!(result.is_ok()); + assert_eq!(get_ret_type(&result.unwrap()), StaticType::Int); + } + + fn get_ret_type(node: &TypedNode) -> StaticType { + if let StaticType::Function(sig) = &node.ty { + sig.ret.clone() + } else { + node.ty.clone() + } + } + + #[test] + fn test_inference_constants() { + assert_eq!(get_ret_type(&check_source("10")), StaticType::Int); + assert_eq!(get_ret_type(&check_source("10.5")), StaticType::Float); + assert_eq!(get_ret_type(&check_source("true")), StaticType::Bool); + assert_eq!(get_ret_type(&check_source("\"hello\"")), StaticType::Text); + } + + #[test] + fn test_inference_variable_propagation() { + // (do (def x 10) x) -> The last 'x' must be Int + let typed = check_source("(do (def x 10) x)"); + // Outer is Lambda, Body is Block + if let NodeKind::Lambda { body, .. } = &typed.kind { + if let NodeKind::Block { exprs } = &body.kind { + let last_expr = exprs.last().unwrap(); + assert_eq!( + last_expr.ty, + StaticType::Int, + "Variable 'x' should be inferred as Int" + ); + } else { + panic!("Expected block in lambda body"); + } + } else { + panic!("Expected Lambda wrapper"); + } + } + + #[test] + fn test_inference_block_type() { + // Block type = last expression type + assert_eq!(get_ret_type(&check_source("(do 1 2.5)")), StaticType::Float); + assert_eq!( + get_ret_type(&check_source("(do 1.5 \"test\")")), + StaticType::Text + ); + } + + #[test] + fn test_inference_lambda_return() { + // (fn [a] 10) -> fn(any) -> Int + // Since it's already a Lambda, it's NOT wrapped further. + let typed = check_source("(fn [a] 10)"); + if let StaticType::Function(sig) = &typed.ty { + assert_eq!(sig.ret, StaticType::Int); + } else { + panic!("Expected function type, got {:?}", typed.ty); + } + + // Nested: (fn [] (do 1 2.5)) -> fn() -> Float + let typed_nested = check_source("(fn [] (do 1 2.5))"); + if let StaticType::Function(sig) = &typed_nested.ty { + assert_eq!(sig.ret, StaticType::Float); + } else { + panic!("Expected function type"); + } + } + + #[test] + fn test_inference_assignment_updates_type() { + // (do (def x 10) (assign x 20.5) x) -> x becomes Float after assignment + let typed = check_source("(do (def x 10) (assign x 20.5) x)"); + if let NodeKind::Lambda { body, .. } = &typed.kind { + if let NodeKind::Block { exprs } = &body.kind { + let last_expr = exprs.last().unwrap(); + assert_eq!( + last_expr.ty, + StaticType::Float, + "Variable 'x' should be specialized to Float after assignment" + ); + } else { + panic!("Expected block"); + } + } else { + panic!("Expected Lambda"); + } + } + + #[test] + fn test_operator_overloading_inference() { + assert_eq!(get_ret_type(&check_source("(+ 1 2)")), StaticType::Int); + assert_eq!( + get_ret_type(&check_source("(+ 1.0 2.0)")), + StaticType::Float + ); + assert_eq!( + get_ret_type(&check_source("(+ \"a\" \"b\")")), + StaticType::Text + ); + assert_eq!(get_ret_type(&check_source("(/ 1 2)")), StaticType::Float); + } + + #[test] + fn test_datetime_inference() { + // date("2023-01-01") -> DateTime + assert_eq!( + get_ret_type(&check_source("(date \"2023-01-01\")")), + StaticType::DateTime + ); + // DateTime + Int -> DateTime + assert_eq!( + get_ret_type(&check_source("(+ (date \"2023-01-01\") 86400000)")), + StaticType::DateTime + ); + // DateTime - DateTime -> Int (Duration) + assert_eq!( + get_ret_type(&check_source( + "(- (date \"2023-01-02\") (date \"2023-01-01\"))" + )), + StaticType::Int + ); + // DateTime comparison -> Bool + assert_eq!( + get_ret_type(&check_source( + "(> (date \"2023-01-02\") (date \"2023-01-01\"))" + )), + StaticType::Bool + ); + } + + #[test] + fn test_inference_tuple_vector_matrix() { + // Heterogeneous -> Tuple + assert_eq!( + get_ret_type(&check_source("[1 3.14 \"text\"]")), + StaticType::Tuple(vec![StaticType::Int, StaticType::Float, StaticType::Text]) + ); + + // Homogeneous -> Vector + assert_eq!( + get_ret_type(&check_source("[10 20 30]")), + StaticType::Vector(Box::new(StaticType::Int), 3) + ); + + // Nested Homogeneous -> Matrix + assert_eq!( + get_ret_type(&check_source("[[1 2] [3 4]]")), + StaticType::Matrix(Box::new(StaticType::Int), vec![2, 2]) + ); + + // Deep Matrix + assert_eq!( + get_ret_type(&check_source("[[[1 2]] [[3 4]]]")), + StaticType::Matrix(Box::new(StaticType::Int), vec![2, 1, 2]) + ); + + // Shape mismatch -> Tuple of Vectors + let mixed = get_ret_type(&check_source("[[1 2] [3 4 5]]")); + if let StaticType::Tuple(elements) = mixed { + assert_eq!( + elements[0], + StaticType::Vector(Box::new(StaticType::Int), 2) + ); + assert_eq!( + elements[1], + StaticType::Vector(Box::new(StaticType::Int), 3) + ); + } else { + panic!("Expected Tuple for shape mismatch, got {:?}", mixed); + } + } + + #[test] + fn test_inference_record() { + use crate::ast::types::Keyword; + let typed = check_source("{:x 1 :y 0.3}"); + let ty = get_ret_type(&typed); + if let StaticType::Record(layout) = ty { + assert_eq!(layout.fields.len(), 2); + assert_eq!(layout.fields[0], (Keyword::intern("x"), StaticType::Int)); + assert_eq!(layout.fields[1], (Keyword::intern("y"), StaticType::Float)); + } else { + panic!("Expected Record, got {:?}", ty); + } + } +} diff --git a/src/ast/environment.rs b/src/ast/environment.rs index 7af98ec..5139f4a 100644 --- a/src/ast/environment.rs +++ b/src/ast/environment.rs @@ -10,8 +10,8 @@ use std::path::{Path, PathBuf}; use std::rc::Rc; use crate::ast::compiler::bound_nodes::{ - Address, AnalyzedNode, BoundKind, ExecNode, GlobalAnalyzedRegistry, GlobalFunctionRegistry, GlobalIdx, - Node, VirtualId, + Address, AnalyzedNode, ExecNode, GlobalAnalyzedRegistry, GlobalFunctionRegistry, GlobalIdx, + LambdaBinding, Node, NodeKind, VirtualId, }; use crate::ast::compiler::dumper::Dumper; use crate::ast::compiler::lambda_collector::LambdaCollector; @@ -117,7 +117,7 @@ impl MacroEvaluator for RuntimeMacroEvaluator { node: &SyntaxNode, bindings: &HashMap, SyntaxNode>, ) -> Result { - if let SyntaxKind::Identifier(sym) = &node.kind + if let SyntaxKind::Identifier { symbol: sym, .. } = &node.kind && let Some(arg_node) = bindings.get(&sym.name) { return Ok(Value::Object(Rc::new(arg_node.clone()) as Rc)); @@ -390,8 +390,8 @@ impl Environment { fn discover_globals(&self, node: &SyntaxNode) { match &node.kind { - SyntaxKind::Def { target, .. } => { - if let SyntaxKind::Identifier(sym) = &target.kind { + SyntaxKind::Def { pattern, .. } => { + if let SyntaxKind::Identifier { symbol: sym, .. } = &pattern.kind { let mut root_scopes = self.root_scopes.borrow_mut(); let last_idx = root_scopes.len() - 1; let current_scope = &mut root_scopes[last_idx]; @@ -417,9 +417,9 @@ impl Environment { fn extract_names(node: &SyntaxNode) -> Vec> { match &node.kind { - SyntaxKind::Identifier(sym) => vec![sym.name.clone()], + SyntaxKind::Identifier { symbol: sym, .. } => vec![sym.name.clone()], SyntaxKind::Tuple { elements } => { - elements.iter().flat_map(extract_names).collect() + elements.iter().flat_map(|e| extract_names(e)).collect() } _ => vec![], } @@ -476,20 +476,22 @@ impl Environment { LambdaCollector::collect(&bound_ast, &mut self.function_registry.borrow_mut()); let checker = TypeChecker::new(self.root_types.clone()); - let wrapped_ast = if let BoundKind::Lambda { .. } = bound_ast.kind { + let wrapped_ast = if let NodeKind::Lambda { .. } = bound_ast.kind { bound_ast } else { Node { identity: bound_ast.identity.clone(), - kind: BoundKind::Lambda { + kind: NodeKind::Lambda { params: std::rc::Rc::new(Node { identity: bound_ast.identity.clone(), - kind: BoundKind::Tuple { elements: vec![] }, + kind: NodeKind::Tuple { elements: vec![] }, ty: (), }), - upvalues: vec![], body: std::rc::Rc::new(bound_ast), - positional_count: Some(0), + info: LambdaBinding { + upvalues: vec![], + positional_count: Some(0), + }, }, ty: (), } @@ -650,20 +652,19 @@ impl Environment { pub fn instantiate(&self, node: ExecNode) -> Rc { let root_values = self.root_values.clone(); - if let BoundKind::Lambda { + if let NodeKind::Lambda { params, - upvalues, body, - positional_count, + info, } = &node.kind - && upvalues.is_empty() + && info.upvalues.is_empty() { let closure = Rc::new(crate::ast::vm::Closure::new( params.clone(), - body.ty.original.clone(), + body.ty.original.clone(), body.clone(), Vec::new(), - *positional_count, + info.positional_count, node.ty.stack_size, )); let closure_obj: Rc = closure; diff --git a/src/ast/nodes.rs b/src/ast/nodes.rs index f82921f..f6eac8d 100644 --- a/src/ast/nodes.rs +++ b/src/ast/nodes.rs @@ -1,134 +1,65 @@ -use crate::ast::types::{Identity, Object, Value}; -use std::any::Any; -use std::fmt::Debug; -use std::rc::Rc; - -/// A name with an optional context for macro hygiene. -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct Symbol { - pub name: Rc, - /// Points to the identity of the Expansion node if this symbol - /// was created/referenced inside a macro expansion. - pub context: Option, -} - -impl From> for Symbol { - fn from(name: Rc) -> Self { - Self { - name, - context: None, - } - } -} - -impl From<&str> for Symbol { - fn from(name: &str) -> Self { - Self { - name: Rc::from(name), - context: None, - } - } -} - -/// A parser AST Node wrapper to preserve identity and metadata -#[derive(Debug, Clone, PartialEq)] -pub struct SyntaxNode { - pub identity: Identity, - pub kind: SyntaxKind, -} - -impl Object for SyntaxNode { - fn type_name(&self) -> &'static str { - "ast-node" - } - fn as_any(&self) -> &dyn Any { - self - } -} - -/// The base for custom node types (extensions) -pub trait CustomNode: Debug { - fn display_name(&self) -> &'static str; - fn clone_box(&self) -> Box; -} - -impl Clone for Box { - fn clone(&self) -> Self { - self.clone_box() - } -} - -#[derive(Debug, Clone, PartialEq)] -pub enum SyntaxKind { - Nop, - Constant(Value), - /// A general identifier (used for both references and declarations in the syntax AST). - Identifier(Symbol), - /// A first-class field accessor (e.g. .name) - FieldAccessor(crate::ast::types::Keyword), - If { - cond: Box, - then_br: Box, - else_br: Option>, - }, - Def { - target: Box, - value: Box, - }, - Assign { - target: Box, - value: Box, - }, - Lambda { - params: Box, - body: Rc, - }, - Call { - callee: Box, - args: Box, - }, - Again { - args: Box, - }, - Pipe { - inputs: Vec, - lambda: Box, - }, - Block { - exprs: Vec, - }, - Tuple { - elements: Vec, - }, - Record { - fields: Vec<(SyntaxNode, SyntaxNode)>, - }, - /// A macro declaration that can be expanded at compile time. - MacroDecl { - name: Symbol, - params: Box, - body: Box, - }, - /// A template for AST nodes, allowing for substitutions. (Quasiquote) - Template(Box), - /// A placeholder inside a template to be replaced by a node or value. (Unquote) - Placeholder(Box), - /// A placeholder that flattens a sequence or merges a record into its surroundings. (Splice) - Splice(Box), - /// Represents an expanded macro call, preserving the original call for debugging. - Expansion { - /// The original call from the source AST. - call: Box, - /// The resulting AST after macro expansion. - expanded: Box, - }, - /// A diagnostic poison node, allowing compilation to continue after an error. - Error, - Extension(Box), -} - -impl PartialEq for Box { - fn eq(&self, _other: &Self) -> bool { - false - } -} +use crate::ast::compiler::bound_nodes::{Node, NodeKind, SyntaxPhase}; +use crate::ast::types::{Identity, Object}; +use std::any::Any; +use std::fmt::Debug; +use std::rc::Rc; + +/// A name with an optional context for macro hygiene. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Symbol { + pub name: Rc, + /// Points to the identity of the Expansion node if this symbol + /// was created/referenced inside a macro expansion. + pub context: Option, +} + +impl From> for Symbol { + fn from(name: Rc) -> Self { + Self { + name, + context: None, + } + } +} + +impl From<&str> for Symbol { + fn from(name: &str) -> Self { + Self { + name: Rc::from(name), + context: None, + } + } +} + +/// Type alias: the parser AST is now `Node`. +pub type SyntaxNode = Node; + +/// Type alias: the parser AST kind is now `NodeKind`. +pub type SyntaxKind = NodeKind; + +impl Object for Node { + fn type_name(&self) -> &'static str { + "ast-node" + } + fn as_any(&self) -> &dyn Any { + self + } +} + +/// The base for custom node types (extensions) +pub trait CustomNode: Debug { + fn display_name(&self) -> &'static str; + fn clone_box(&self) -> Box; +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_box() + } +} + +impl PartialEq for Box { + fn eq(&self, _other: &Self) -> bool { + false + } +} diff --git a/src/ast/parser.rs b/src/ast/parser.rs index 8e776eb..0d844b8 100644 --- a/src/ast/parser.rs +++ b/src/ast/parser.rs @@ -1,522 +1,544 @@ -use crate::ast::diagnostics::Diagnostics; -use crate::ast::lexer::{Lexer, Token, TokenKind}; -use crate::ast::nodes::{Symbol, SyntaxKind, SyntaxNode}; -use crate::ast::types::{Identity, Keyword, NodeIdentity, Value}; -use std::rc::Rc; - -pub struct Parser<'a> { - lexer: Lexer<'a>, - current_token: Token, - pub diagnostics: Diagnostics, -} - -impl<'a> Parser<'a> { - pub fn new(input: &'a str) -> Self { - let mut lexer = Lexer::new(input); - let mut diagnostics = Diagnostics::new(); - let current_token = match lexer.next_token() { - Ok(t) => t, - Err(e) => { - diagnostics.push_error(e, None); - Token { - kind: TokenKind::EOF, - location: crate::ast::types::SourceLocation { line: 1, col: 1 }, - } - } - }; - Self { - lexer, - current_token, - diagnostics, - } - } - - fn advance(&mut self) -> Token { - let next = match self.lexer.next_token() { - Ok(t) => t, - Err(e) => { - self.diagnostics - .push_error(e, Some(NodeIdentity::new(self.current_token.location))); - Token { - kind: TokenKind::EOF, - location: self.current_token.location, - } - } - }; - std::mem::replace(&mut self.current_token, next) - } - - fn peek(&self) -> &TokenKind { - &self.current_token.kind - } - - pub fn parse_expression(&mut self) -> SyntaxNode { - let token_loc = self.current_token.location; - let identity = NodeIdentity::new(token_loc); - - match self.peek() { - TokenKind::LeftParen => self.parse_list(), - TokenKind::LeftBracket => self.parse_vector_literal(), - TokenKind::LeftBrace => self.parse_record_literal(), - TokenKind::Quote => { - self.advance(); // consume ' - let expr = self.parse_expression(); - SyntaxNode { - identity: identity.clone(), - kind: SyntaxKind::Call { - callee: Box::new(self.make_id_node("quote", identity.clone())), - args: Box::new(SyntaxNode { - identity, - kind: SyntaxKind::Tuple { - elements: vec![expr], - }, - - }), - }, - - } - } - TokenKind::Backtick => { - self.advance(); // consume ` - let expr = self.parse_expression(); - SyntaxNode { - identity, - kind: SyntaxKind::Template(Box::new(expr)), - - } - } - TokenKind::Tilde => { - self.advance(); // consume ~ - if *self.peek() == TokenKind::At { - self.advance(); // consume @ - let expr = self.parse_expression(); - SyntaxNode { - identity, - kind: SyntaxKind::Splice(Box::new(expr)), - - } - } else { - let expr = self.parse_expression(); - SyntaxNode { - identity, - kind: SyntaxKind::Placeholder(Box::new(expr)), - - } - } - } - _ => self.parse_atom(), - } - } - - pub fn at_eof(&self) -> bool { - matches!(self.current_token.kind, TokenKind::EOF) - } - - fn synchronize(&mut self) { - while !self.at_eof() { - match self.peek() { - TokenKind::RightParen | TokenKind::RightBracket | TokenKind::RightBrace => { - self.advance(); - return; - } - _ => { - self.advance(); - } - } - } - } - - fn parse_atom(&mut self) -> SyntaxNode { - let token = self.advance(); - let identity = NodeIdentity::new(token.location); - - let kind = match token.kind { - TokenKind::Integer(n) => SyntaxKind::Constant(Value::Int(n)), - TokenKind::Float(n) => SyntaxKind::Constant(Value::Float(n)), - TokenKind::String(s) => SyntaxKind::Constant(Value::Text(s)), - TokenKind::Keyword(k) => SyntaxKind::Constant(Value::Keyword(Keyword::intern(&k))), - TokenKind::Identifier(id) => match id.as_ref() { - "..." => SyntaxKind::Nop, - s if s.starts_with('.') && s.len() > 1 => { - SyntaxKind::FieldAccessor(Keyword::intern(&s[1..])) - } - _ => SyntaxKind::Identifier(id.into()), - }, - TokenKind::EOF => SyntaxKind::Error, // Error already logged by advance - _ => { - self.diagnostics.push_error( - format!("Unexpected token in atom: {:?}", token.kind), - Some(identity.clone()), - ); - SyntaxKind::Error - } - }; - - SyntaxNode { - identity, - kind, - } - } - - fn parse_list(&mut self) -> SyntaxNode { - let start_loc = self.advance().location; // consume '(' - let identity = NodeIdentity::new(start_loc); - - if *self.peek() == TokenKind::RightParen { - self.diagnostics.push_error( - "Empty list () is not a valid expression", - Some(identity.clone()), - ); - self.advance(); // consume ) - return SyntaxNode { - identity, - kind: SyntaxKind::Error, - - }; - } - - let head = self.parse_expression(); - - let node = if let SyntaxKind::Identifier(ref sym) = head.kind { - match sym.name.as_ref() { - "if" => self.parse_if(identity), - "fn" => self.parse_fn(identity), - "pipe" => self.parse_pipe(identity), - "again" => self.parse_again(identity), - "def" => self.parse_def(identity), - "assign" => self.parse_assign(identity), - "do" => self.parse_do(identity), - "macro" => self.parse_macro_decl(identity), - _ => self.parse_call(head, identity), - } - } else { - self.parse_call(head, identity) - }; - - self.expect(TokenKind::RightParen); - node - } - - fn parse_again(&mut self, identity: Identity) -> SyntaxNode { - let mut elements = Vec::new(); - - while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF { - elements.push(self.parse_expression()); - } - - let args_node = SyntaxNode { - identity: identity.clone(), - kind: SyntaxKind::Tuple { elements }, - - }; - - SyntaxNode { - identity, - kind: SyntaxKind::Again { - args: Box::new(args_node), - }, - - } - } - - fn parse_if(&mut self, identity: Identity) -> SyntaxNode { - let cond = Box::new(self.parse_expression()); - let then_br = Box::new(self.parse_expression()); - let mut else_br = None; - - if *self.peek() != TokenKind::RightParen { - else_br = Some(Box::new(self.parse_expression())); - } - - SyntaxNode { - identity, - kind: SyntaxKind::If { - cond, - then_br, - else_br, - }, - - } - } - - fn parse_def(&mut self, identity: Identity) -> SyntaxNode { - let target = Box::new(self.parse_pattern()); - let value = Box::new(self.parse_expression()); - - SyntaxNode { - identity, - kind: SyntaxKind::Def { target, value }, - - } - } - - fn parse_assign(&mut self, identity: Identity) -> SyntaxNode { - // (assign target value) - let target = Box::new(self.parse_expression()); - let value = Box::new(self.parse_expression()); - - SyntaxNode { - identity, - kind: SyntaxKind::Assign { target, value }, - - } - } - - fn parse_do(&mut self, identity: Identity) -> SyntaxNode { - let mut exprs = Vec::new(); - while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF { - exprs.push(self.parse_expression()); - } - SyntaxNode { - identity, - kind: SyntaxKind::Block { exprs }, - - } - } - - fn parse_pipe(&mut self, identity: Identity) -> SyntaxNode { - let inputs_node = self.parse_expression(); - let inputs = match inputs_node.kind { - SyntaxKind::Tuple { elements } => elements, - _ => vec![inputs_node], - }; - let lambda = Box::new(self.parse_expression()); - - SyntaxNode { - identity, - kind: SyntaxKind::Pipe { inputs, lambda }, - - } - } - - fn parse_fn(&mut self, identity: Identity) -> SyntaxNode { - let params = Box::new(self.parse_param_vector()); - let body = self.parse_expression(); - - SyntaxNode { - identity, - kind: SyntaxKind::Lambda { - params, - body: Rc::new(body), - }, - - } - } - - fn parse_macro_decl(&mut self, identity: Identity) -> SyntaxNode { - let name_node = self.parse_expression(); - let name = match name_node.kind { - SyntaxKind::Identifier(sym) => sym, - _ => { - self.diagnostics.push_error( - "Expected identifier for macro name", - Some(name_node.identity.clone()), - ); - Symbol::from("error") - } - }; - let params = Box::new(self.parse_param_vector()); - let body = self.parse_expression(); - SyntaxNode { - identity, - kind: SyntaxKind::MacroDecl { - name, - params, - body: Box::new(body), - }, - - } - } - - fn parse_param_vector(&mut self) -> SyntaxNode { - if *self.peek() != TokenKind::LeftBracket { - self.diagnostics.push_error( - format!( - "Expected parameter vector [...] for fn, found {:?}", - self.peek() - ), - Some(NodeIdentity::new(self.current_token.location)), - ); - return SyntaxNode { - identity: NodeIdentity::new(self.current_token.location), - kind: SyntaxKind::Error, - - }; - } - self.parse_pattern() - } - - fn parse_pattern(&mut self) -> SyntaxNode { - let next = self.peek(); - match next { - TokenKind::Identifier(_) => { - let token = self.advance(); - let sym: Symbol = match token.kind { - TokenKind::Identifier(s) => s.into(), - _ => unreachable!(), - }; - SyntaxNode { - identity: NodeIdentity::new(token.location), - kind: SyntaxKind::Identifier(sym), - - } - } - TokenKind::LeftBracket => { - let token = self.advance(); - let identity = NodeIdentity::new(token.location); - - let mut elements = Vec::new(); - while *self.peek() != TokenKind::RightBracket && *self.peek() != TokenKind::EOF { - elements.push(self.parse_pattern()); - } - self.expect(TokenKind::RightBracket); - - SyntaxNode { - identity, - kind: SyntaxKind::Tuple { elements }, - - } - } - TokenKind::Tilde => { - let token = self.advance(); // consume ~ - if *self.peek() == TokenKind::At { - self.advance(); // consume @ - let expr = self.parse_expression(); - SyntaxNode { - identity: NodeIdentity::new(token.location), - kind: SyntaxKind::Splice(Box::new(expr)), - - } - } else { - let expr = self.parse_expression(); - SyntaxNode { - identity: NodeIdentity::new(token.location), - kind: SyntaxKind::Placeholder(Box::new(expr)), - - } - } - } - _ => { - self.diagnostics.push_error( - format!( - "Expected identifier or pattern vector [...] for definition, found {:?}", - next - ), - Some(NodeIdentity::new(self.current_token.location)), - ); - self.advance(); - SyntaxNode { - identity: NodeIdentity::new(self.current_token.location), - kind: SyntaxKind::Error, - - } - } - } - } - - fn parse_call(&mut self, callee: SyntaxNode, identity: Identity) -> SyntaxNode { - let mut elements = Vec::new(); - - while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF { - elements.push(self.parse_expression()); - } - - // The arguments are wrapped in a Tuple node, reusing the call's identity/location. - let args_node = SyntaxNode { - identity: identity.clone(), - kind: SyntaxKind::Tuple { elements }, - - }; - - SyntaxNode { - identity, - kind: SyntaxKind::Call { - callee: Box::new(callee), - args: Box::new(args_node), - }, - - } - } - - fn parse_vector_literal(&mut self) -> SyntaxNode { - let token = self.advance(); - let mut elements = Vec::new(); - - while *self.peek() != TokenKind::RightBracket && *self.peek() != TokenKind::EOF { - let expr = self.parse_expression(); - elements.push(expr); - } - - self.expect(TokenKind::RightBracket); - - SyntaxNode { - identity: NodeIdentity::new(token.location), - kind: SyntaxKind::Tuple { elements }, - - } - } - - fn parse_record_literal(&mut self) -> SyntaxNode { - let token = self.advance(); - let mut fields = Vec::new(); - - while *self.peek() != TokenKind::RightBrace && *self.peek() != TokenKind::EOF { - let key_node = self.parse_expression(); - // We check for keyword kind here (syntactically) to avoid ambiguity, but - // strictly we could allow any expression and check at runtime. - // Delphi enforces keywords. We can do minimal check here. - match &key_node.kind { - SyntaxKind::Constant(Value::Keyword(_)) => {} - _ => { - self.diagnostics.push_error( - "Record keys must be keywords (syntactically)", - Some(key_node.identity.clone()), - ); - } - } - - if *self.peek() == TokenKind::RightBrace || *self.peek() == TokenKind::EOF { - self.diagnostics.push_error( - "Record literal must have even number of forms", - Some(NodeIdentity::new(self.current_token.location)), - ); - break; - } - let val_node = self.parse_expression(); - - fields.push((key_node, val_node)); - } - self.expect(TokenKind::RightBrace); - - SyntaxNode { - identity: NodeIdentity::new(token.location), - kind: SyntaxKind::Record { fields }, - - } - } - - fn expect(&mut self, kind: TokenKind) -> Token { - if self.peek() == &kind { - self.advance() - } else { - self.diagnostics.push_error( - format!("Expected {:?}, but found {:?}", kind, self.peek()), - Some(NodeIdentity::new(self.current_token.location)), - ); - // Recovery: skip until we find what we expected or a synchronization point - self.synchronize(); - Token { - kind, - location: self.current_token.location, - } - } - } - - fn make_id_node(&self, name: &str, identity: Identity) -> SyntaxNode { - SyntaxNode { - identity, - kind: SyntaxKind::Identifier(Symbol::from(name)), - - } - } -} +use crate::ast::diagnostics::Diagnostics; +use crate::ast::lexer::{Lexer, Token, TokenKind}; +use crate::ast::nodes::{Symbol, SyntaxKind, SyntaxNode}; +use crate::ast::types::{Identity, Keyword, NodeIdentity, Value}; +use std::rc::Rc; + +pub struct Parser<'a> { + lexer: Lexer<'a>, + current_token: Token, + pub diagnostics: Diagnostics, +} + +impl<'a> Parser<'a> { + pub fn new(input: &'a str) -> Self { + let mut lexer = Lexer::new(input); + let mut diagnostics = Diagnostics::new(); + let current_token = match lexer.next_token() { + Ok(t) => t, + Err(e) => { + diagnostics.push_error(e, None); + Token { + kind: TokenKind::EOF, + location: crate::ast::types::SourceLocation { line: 1, col: 1 }, + } + } + }; + Self { + lexer, + current_token, + diagnostics, + } + } + + fn advance(&mut self) -> Token { + let next = match self.lexer.next_token() { + Ok(t) => t, + Err(e) => { + self.diagnostics + .push_error(e, Some(NodeIdentity::new(self.current_token.location))); + Token { + kind: TokenKind::EOF, + location: self.current_token.location, + } + } + }; + std::mem::replace(&mut self.current_token, next) + } + + fn peek(&self) -> &TokenKind { + &self.current_token.kind + } + + pub fn parse_expression(&mut self) -> SyntaxNode { + let token_loc = self.current_token.location; + let identity = NodeIdentity::new(token_loc); + + match self.peek() { + TokenKind::LeftParen => self.parse_list(), + TokenKind::LeftBracket => self.parse_vector_literal(), + TokenKind::LeftBrace => self.parse_record_literal(), + TokenKind::Quote => { + self.advance(); // consume ' + let expr = self.parse_expression(); + SyntaxNode { + identity: identity.clone(), + kind: SyntaxKind::Call { + callee: Rc::new(self.make_id_node("quote", identity.clone())), + args: Rc::new(SyntaxNode { + identity, + kind: SyntaxKind::Tuple { + elements: vec![Rc::new(expr)], + }, + ty: (), + }), + }, + ty: (), + } + } + TokenKind::Backtick => { + self.advance(); // consume ` + let expr = self.parse_expression(); + SyntaxNode { + identity, + kind: SyntaxKind::Template(Rc::new(expr)), + ty: (), + } + } + TokenKind::Tilde => { + self.advance(); // consume ~ + if *self.peek() == TokenKind::At { + self.advance(); // consume @ + let expr = self.parse_expression(); + SyntaxNode { + identity, + kind: SyntaxKind::Splice(Rc::new(expr)), + ty: (), + } + } else { + let expr = self.parse_expression(); + SyntaxNode { + identity, + kind: SyntaxKind::Placeholder(Rc::new(expr)), + ty: (), + } + } + } + _ => self.parse_atom(), + } + } + + pub fn at_eof(&self) -> bool { + matches!(self.current_token.kind, TokenKind::EOF) + } + + fn synchronize(&mut self) { + while !self.at_eof() { + match self.peek() { + TokenKind::RightParen | TokenKind::RightBracket | TokenKind::RightBrace => { + self.advance(); + return; + } + _ => { + self.advance(); + } + } + } + } + + fn parse_atom(&mut self) -> SyntaxNode { + let token = self.advance(); + let identity = NodeIdentity::new(token.location); + + let kind = match token.kind { + TokenKind::Integer(n) => SyntaxKind::Constant(Value::Int(n)), + TokenKind::Float(n) => SyntaxKind::Constant(Value::Float(n)), + TokenKind::String(s) => SyntaxKind::Constant(Value::Text(s)), + TokenKind::Keyword(k) => SyntaxKind::Constant(Value::Keyword(Keyword::intern(&k))), + TokenKind::Identifier(id) => match id.as_ref() { + "..." => SyntaxKind::Nop, + s if s.starts_with('.') && s.len() > 1 => { + SyntaxKind::FieldAccessor(Keyword::intern(&s[1..])) + } + _ => SyntaxKind::Identifier { + symbol: id.into(), + binding: (), + }, + }, + TokenKind::EOF => SyntaxKind::Error, // Error already logged by advance + _ => { + self.diagnostics.push_error( + format!("Unexpected token in atom: {:?}", token.kind), + Some(identity.clone()), + ); + SyntaxKind::Error + } + }; + + SyntaxNode { + identity, + kind, + ty: (), + } + } + + fn parse_list(&mut self) -> SyntaxNode { + let start_loc = self.advance().location; // consume '(' + let identity = NodeIdentity::new(start_loc); + + if *self.peek() == TokenKind::RightParen { + self.diagnostics.push_error( + "Empty list () is not a valid expression", + Some(identity.clone()), + ); + self.advance(); // consume ) + return SyntaxNode { + identity, + kind: SyntaxKind::Error, + ty: (), + }; + } + + let head = self.parse_expression(); + + let node = if let SyntaxKind::Identifier { ref symbol, .. } = head.kind { + match symbol.name.as_ref() { + "if" => self.parse_if(identity), + "fn" => self.parse_fn(identity), + "pipe" => self.parse_pipe(identity), + "again" => self.parse_again(identity), + "def" => self.parse_def(identity), + "assign" => self.parse_assign(identity), + "do" => self.parse_do(identity), + "macro" => self.parse_macro_decl(identity), + _ => self.parse_call(head, identity), + } + } else { + self.parse_call(head, identity) + }; + + self.expect(TokenKind::RightParen); + node + } + + fn parse_again(&mut self, identity: Identity) -> SyntaxNode { + let mut elements = Vec::new(); + + while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF { + elements.push(Rc::new(self.parse_expression())); + } + + let args_node = SyntaxNode { + identity: identity.clone(), + kind: SyntaxKind::Tuple { elements }, + ty: (), + }; + + SyntaxNode { + identity, + kind: SyntaxKind::Again { + args: Rc::new(args_node), + }, + ty: (), + } + } + + fn parse_if(&mut self, identity: Identity) -> SyntaxNode { + let cond = Rc::new(self.parse_expression()); + let then_br = Rc::new(self.parse_expression()); + let mut else_br = None; + + if *self.peek() != TokenKind::RightParen { + else_br = Some(Rc::new(self.parse_expression())); + } + + SyntaxNode { + identity, + kind: SyntaxKind::If { + cond, + then_br, + else_br, + }, + ty: (), + } + } + + fn parse_def(&mut self, identity: Identity) -> SyntaxNode { + let pattern = Rc::new(self.parse_pattern()); + let value = Rc::new(self.parse_expression()); + + SyntaxNode { + identity, + kind: SyntaxKind::Def { + pattern, + value, + info: (), + }, + ty: (), + } + } + + fn parse_assign(&mut self, identity: Identity) -> SyntaxNode { + // (assign target value) + let target = Rc::new(self.parse_expression()); + let value = Rc::new(self.parse_expression()); + + SyntaxNode { + identity, + kind: SyntaxKind::Assign { + target, + value, + info: (), + }, + ty: (), + } + } + + fn parse_do(&mut self, identity: Identity) -> SyntaxNode { + let mut exprs = Vec::new(); + while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF { + exprs.push(Rc::new(self.parse_expression())); + } + SyntaxNode { + identity, + kind: SyntaxKind::Block { exprs }, + ty: (), + } + } + + fn parse_pipe(&mut self, identity: Identity) -> SyntaxNode { + let inputs_node = self.parse_expression(); + let inputs = match inputs_node.kind { + SyntaxKind::Tuple { elements } => elements, + _ => vec![Rc::new(inputs_node)], + }; + let lambda = Rc::new(self.parse_expression()); + + SyntaxNode { + identity, + kind: SyntaxKind::Pipe { inputs, lambda }, + ty: (), + } + } + + fn parse_fn(&mut self, identity: Identity) -> SyntaxNode { + let params = Rc::new(self.parse_param_vector()); + let body = self.parse_expression(); + + SyntaxNode { + identity, + kind: SyntaxKind::Lambda { + params, + body: Rc::new(body), + info: (), + }, + ty: (), + } + } + + fn parse_macro_decl(&mut self, identity: Identity) -> SyntaxNode { + let name_node = self.parse_expression(); + let name = match name_node.kind { + SyntaxKind::Identifier { symbol, .. } => symbol, + _ => { + self.diagnostics.push_error( + "Expected identifier for macro name", + Some(name_node.identity.clone()), + ); + Symbol::from("error") + } + }; + let params = Rc::new(self.parse_param_vector()); + let body = self.parse_expression(); + SyntaxNode { + identity, + kind: SyntaxKind::MacroDecl { + name, + params, + body: Rc::new(body), + }, + ty: (), + } + } + + fn parse_param_vector(&mut self) -> SyntaxNode { + if *self.peek() != TokenKind::LeftBracket { + self.diagnostics.push_error( + format!( + "Expected parameter vector [...] for fn, found {:?}", + self.peek() + ), + Some(NodeIdentity::new(self.current_token.location)), + ); + return SyntaxNode { + identity: NodeIdentity::new(self.current_token.location), + kind: SyntaxKind::Error, + ty: (), + }; + } + self.parse_pattern() + } + + fn parse_pattern(&mut self) -> SyntaxNode { + let next = self.peek(); + match next { + TokenKind::Identifier(_) => { + let token = self.advance(); + let sym: Symbol = match token.kind { + TokenKind::Identifier(s) => s.into(), + _ => unreachable!(), + }; + SyntaxNode { + identity: NodeIdentity::new(token.location), + kind: SyntaxKind::Identifier { + symbol: sym, + binding: (), + }, + ty: (), + } + } + TokenKind::LeftBracket => { + let token = self.advance(); + let identity = NodeIdentity::new(token.location); + + let mut elements = Vec::new(); + while *self.peek() != TokenKind::RightBracket && *self.peek() != TokenKind::EOF { + elements.push(Rc::new(self.parse_pattern())); + } + self.expect(TokenKind::RightBracket); + + SyntaxNode { + identity, + kind: SyntaxKind::Tuple { elements }, + ty: (), + } + } + TokenKind::Tilde => { + let token = self.advance(); // consume ~ + if *self.peek() == TokenKind::At { + self.advance(); // consume @ + let expr = self.parse_expression(); + SyntaxNode { + identity: NodeIdentity::new(token.location), + kind: SyntaxKind::Splice(Rc::new(expr)), + ty: (), + } + } else { + let expr = self.parse_expression(); + SyntaxNode { + identity: NodeIdentity::new(token.location), + kind: SyntaxKind::Placeholder(Rc::new(expr)), + ty: (), + } + } + } + _ => { + self.diagnostics.push_error( + format!( + "Expected identifier or pattern vector [...] for definition, found {:?}", + next + ), + Some(NodeIdentity::new(self.current_token.location)), + ); + self.advance(); + SyntaxNode { + identity: NodeIdentity::new(self.current_token.location), + kind: SyntaxKind::Error, + ty: (), + } + } + } + } + + fn parse_call(&mut self, callee: SyntaxNode, identity: Identity) -> SyntaxNode { + let mut elements = Vec::new(); + + while *self.peek() != TokenKind::RightParen && *self.peek() != TokenKind::EOF { + elements.push(Rc::new(self.parse_expression())); + } + + // The arguments are wrapped in a Tuple node, reusing the call's identity/location. + let args_node = SyntaxNode { + identity: identity.clone(), + kind: SyntaxKind::Tuple { elements }, + ty: (), + }; + + SyntaxNode { + identity, + kind: SyntaxKind::Call { + callee: Rc::new(callee), + args: Rc::new(args_node), + }, + ty: (), + } + } + + fn parse_vector_literal(&mut self) -> SyntaxNode { + let token = self.advance(); + let mut elements = Vec::new(); + + while *self.peek() != TokenKind::RightBracket && *self.peek() != TokenKind::EOF { + let expr = self.parse_expression(); + elements.push(Rc::new(expr)); + } + + self.expect(TokenKind::RightBracket); + + SyntaxNode { + identity: NodeIdentity::new(token.location), + kind: SyntaxKind::Tuple { elements }, + ty: (), + } + } + + fn parse_record_literal(&mut self) -> SyntaxNode { + let token = self.advance(); + let mut fields = Vec::new(); + + while *self.peek() != TokenKind::RightBrace && *self.peek() != TokenKind::EOF { + let key_node = self.parse_expression(); + // We check for keyword kind here (syntactically) to avoid ambiguity, but + // strictly we could allow any expression and check at runtime. + // Delphi enforces keywords. We can do minimal check here. + match &key_node.kind { + SyntaxKind::Constant(Value::Keyword(_)) => {} + _ => { + self.diagnostics.push_error( + "Record keys must be keywords (syntactically)", + Some(key_node.identity.clone()), + ); + } + } + + if *self.peek() == TokenKind::RightBrace || *self.peek() == TokenKind::EOF { + self.diagnostics.push_error( + "Record literal must have even number of forms", + Some(NodeIdentity::new(self.current_token.location)), + ); + break; + } + let val_node = self.parse_expression(); + + fields.push((Rc::new(key_node), Rc::new(val_node))); + } + self.expect(TokenKind::RightBrace); + + SyntaxNode { + identity: NodeIdentity::new(token.location), + kind: SyntaxKind::Record { + fields, + layout: (), + }, + ty: (), + } + } + + fn expect(&mut self, kind: TokenKind) -> Token { + if self.peek() == &kind { + self.advance() + } else { + self.diagnostics.push_error( + format!("Expected {:?}, but found {:?}", kind, self.peek()), + Some(NodeIdentity::new(self.current_token.location)), + ); + // Recovery: skip until we find what we expected or a synchronization point + self.synchronize(); + Token { + kind, + location: self.current_token.location, + } + } + } + + fn make_id_node(&self, name: &str, identity: Identity) -> SyntaxNode { + SyntaxNode { + identity, + kind: SyntaxKind::Identifier { + symbol: Symbol::from(name), + binding: (), + }, + ty: (), + } + } +} diff --git a/src/ast/vm.rs b/src/ast/vm.rs index f7b112d..1093d39 100644 --- a/src/ast/vm.rs +++ b/src/ast/vm.rs @@ -1,4 +1,4 @@ -use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, ExecNode, StackOffset}; +use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, ExecNode, IdentifierBinding, NodeKind, StackOffset}; use crate::ast::types::{Object, Value}; use std::any::Any; use std::cell::RefCell; @@ -105,7 +105,7 @@ impl VMObserver for TracingObserver { self.indent = self.indent.saturating_sub(1); let pad = self.pad(); match &node.kind { - BoundKind::Define { .. } | BoundKind::Set { .. } => { + NodeKind::Def { .. } | NodeKind::Assign { .. } => { let s_pad = format!("{}| ", pad); self.logs.push(format!("{}--- Scope Status ---", s_pad)); self.logs.push(format!( @@ -303,58 +303,76 @@ impl VM { #[inline(always)] fn eval_core(&mut self, obs: &mut O, node: &ExecNode) -> Result { match &node.kind { - BoundKind::Nop => Ok(Value::Void), - BoundKind::Constant(v) => Ok(v.clone()), - BoundKind::Define { - addr, - value, - captured_by, - .. - } => { + NodeKind::Nop => Ok(Value::Void), + NodeKind::Constant(v) => Ok(v.clone()), + NodeKind::Def { pattern, value, info } => { let val = self.eval_internal(obs, value)?; - - let mut needs_cell_wrap = false; - if !captured_by.is_empty() - && let Address::Local(slot) = addr - { - let frame = self.frames.last().unwrap(); - let abs_index = frame.stack_base + (slot.0 as usize); - - // Robustness Fix: If the slot is already a Cell (due to forward capture - // in a recursive scenario or complex pre-allocation), don't wrap it again. - // This prevents Cell(Cell(Value)) nesting which causes type errors. - if abs_index >= self.stack.len() || !matches!(self.stack[abs_index], Value::Cell(_)) { - needs_cell_wrap = true; + match &pattern.kind { + NodeKind::Identifier { binding, .. } => { + let addr = match binding { + IdentifierBinding::Declaration { addr, .. } | IdentifierBinding::Reference(addr) => *addr, + }; + + let mut needs_cell_wrap = false; + if !info.captured_by.is_empty() + && let Address::Local(slot) = addr + { + let frame = self.frames.last().unwrap(); + let abs_index = frame.stack_base + (slot.0 as usize); + + // Robustness Fix: If the slot is already a Cell (due to forward capture + // in a recursive scenario or complex pre-allocation), don't wrap it again. + // This prevents Cell(Cell(Value)) nesting which causes type errors. + if abs_index >= self.stack.len() || !matches!(self.stack[abs_index], Value::Cell(_)) { + needs_cell_wrap = true; + } + } + + let store_val = if needs_cell_wrap { + Value::Cell(Rc::new(RefCell::new(val.clone()))) + } else { + val.clone() + }; + + self.set_value(addr, store_val)?; + } + _ => { + // Destructuring (was Destructure variant) + let mut offset = 0; + if let Some(vals) = val.as_slice() { + self.unpack(pattern.as_ref(), vals, &mut offset)?; + } else { + self.unpack(pattern.as_ref(), std::slice::from_ref(&val), &mut offset)?; + } } } - - let store_val = if needs_cell_wrap { - Value::Cell(Rc::new(RefCell::new(val.clone()))) - } else { - val.clone() - }; - - self.set_value(*addr, store_val)?; - // Define always evaluates to the unwrapped value for immediate use. + // Def always evaluates to the unwrapped value for immediate use. Ok(val) } - BoundKind::Destructure { pattern, value } => { + NodeKind::Assign { target, value, info } => { let val = self.eval_internal(obs, value)?; - let mut offset = 0; - - // Destructuring works on tuples/vectors, or single values wrapped in a slice - if let Some(vals) = val.as_slice() { - self.unpack(pattern.as_ref(), vals, &mut offset)?; + if let Some(addr) = info.addr { + self.set_value(addr, val.clone())?; } else { - self.unpack(pattern.as_ref(), std::slice::from_ref(&val), &mut offset)?; + // Destructuring assign + let mut offset = 0; + if let Some(vals) = val.as_slice() { + self.unpack(target.as_ref(), vals, &mut offset)?; + } else { + self.unpack(target.as_ref(), std::slice::from_ref(&val), &mut offset)?; + } } Ok(val) } - BoundKind::Get { addr, .. } => self.get_value(*addr), + NodeKind::Identifier { binding, .. } => { + match binding { + IdentifierBinding::Reference(addr) | IdentifierBinding::Declaration { addr, .. } => self.get_value(*addr), + } + } - BoundKind::FieldAccessor(k) => Ok(Value::FieldAccessor(*k)), + NodeKind::FieldAccessor(k) => Ok(Value::FieldAccessor(*k)), - BoundKind::GetField { rec, field } => { + NodeKind::GetField { rec, field } => { let rec_val = self.eval_internal(obs, rec)?; match rec_val { @@ -395,12 +413,7 @@ impl VM { } } - BoundKind::Set { addr, value } => { - let val = self.eval_internal(obs, value)?; - self.set_value(*addr, val.clone())?; - Ok(val) - } - BoundKind::If { + NodeKind::If { cond, then_br, else_br, @@ -414,10 +427,9 @@ impl VM { Ok(Value::Void) } } - BoundKind::Pipe { + NodeKind::Pipe { inputs, lambda, - out_type, } => { use crate::ast::rtl::streams::StreamNode; @@ -464,25 +476,24 @@ impl VM { // Delegate to the RTL Factory for specialized buffer instantiation let node = - crate::ast::rtl::streams::build_pipeline_node(obs_streams, executor, out_type); + crate::ast::rtl::streams::build_pipeline_node(obs_streams, executor, &node.ty.ty); Ok(Value::Object(node)) } - BoundKind::Block { exprs } => { + NodeKind::Block { exprs } => { let mut last = Value::Void; for e in exprs { last = self.eval_internal(obs, e)?; } Ok(last) } - BoundKind::Lambda { + NodeKind::Lambda { params, - upvalues, body, - positional_count, + info, } => { - let mut captured = Vec::with_capacity(upvalues.len()); - for addr in upvalues { + let mut captured = Vec::with_capacity(info.upvalues.len()); + for addr in &info.upvalues { captured.push(self.capture_upvalue(*addr)?); } let stack_size = node.ty.stack_size; @@ -491,12 +502,12 @@ impl VM { body.ty.original.clone(), body.clone(), captured, - *positional_count, + info.positional_count, stack_size, ); Ok(Value::Object(Rc::new(closure))) } - BoundKind::Call { callee, args } => { + NodeKind::Call { callee, args } => { let func_val = self.eval_internal(obs, callee)?; let base = self.stack.len(); @@ -654,7 +665,7 @@ impl VM { } else { let args_for_unpack = self.stack[base..].to_vec(); self.stack.truncate(base); - if let BoundKind::Tuple { elements } = + if let NodeKind::Tuple { elements } = &closure.parameter_node.kind { let mut offset = 0; @@ -755,7 +766,7 @@ impl VM { } } } - BoundKind::Again { args } => { + NodeKind::Again { args } => { let base = self.stack.len(); if let Err(e) = self.eval_args_to_stack(obs, args) { self.stack.truncate(base); @@ -774,16 +785,16 @@ impl VM { Err("'again' called outside of a closure".to_string()) } } - BoundKind::Tuple { elements } => { + NodeKind::Tuple { elements } => { let mut vals = Vec::with_capacity(elements.len()); for e in elements { vals.push(self.eval_internal(obs, e)?); } Ok(Value::make_tuple(vals)) } - BoundKind::Record { layout, values } => { - let mut evaluated_values = Vec::with_capacity(values.len()); - for v in values { + NodeKind::Record { fields, layout } => { + let mut evaluated_values = Vec::with_capacity(fields.len()); + for (_, v) in fields { evaluated_values.push(self.eval_internal(obs, v)?); } Ok(Value::Record( @@ -791,11 +802,11 @@ impl VM { std::rc::Rc::new(evaluated_values), )) } - BoundKind::Expansion { bound_expanded, .. } => { - let mut curr = bound_expanded; + NodeKind::Expansion { expanded, .. } => { + let mut curr = expanded; if !O::ACTIVE { - while let BoundKind::Expansion { - bound_expanded: next, + while let NodeKind::Expansion { + expanded: next, .. } = &curr.kind { @@ -804,11 +815,15 @@ impl VM { } self.eval_internal(obs, curr) } - BoundKind::Extension(ext) => Err(format!( + NodeKind::Extension(ext) => Err(format!( "Execution of extension '{}' not implemented yet", ext.display_name() )), - BoundKind::Error => Err("Cannot execute a poisoned AST node".to_string()), + NodeKind::Error => Err("Cannot execute a poisoned AST node".to_string()), + // Syntax-only variants that should never appear in runtime phase + NodeKind::MacroDecl { .. } | NodeKind::Template(_) | NodeKind::Placeholder(_) | NodeKind::Splice(_) => { + Err("Syntax-only node reached the VM".to_string()) + } } } @@ -818,12 +833,12 @@ impl VM { args: &ExecNode, ) -> Result<(), String> { match &args.kind { - BoundKind::Tuple { elements } => { + NodeKind::Tuple { elements } => { for e in elements { let mut curr = e.as_ref(); if !O::ACTIVE { - while let BoundKind::Expansion { - bound_expanded: next, + while let NodeKind::Expansion { + expanded: next, .. } = &curr.kind { @@ -832,7 +847,7 @@ impl VM { } match &curr.kind { - BoundKind::Constant(v) if !O::ACTIVE => self.stack.push(v.clone()), + NodeKind::Constant(v) if !O::ACTIVE => self.stack.push(v.clone()), _ => { let val = self.eval_internal(obs, curr)?; self.stack.push(val); @@ -841,7 +856,7 @@ impl VM { } Ok(()) } - BoundKind::Constant(v) => { + NodeKind::Constant(v) => { if let Some(slice) = v.as_slice() { self.stack.extend_from_slice(slice); } else { @@ -849,11 +864,11 @@ impl VM { } Ok(()) } - BoundKind::Expansion { bound_expanded, .. } => { - let mut curr = bound_expanded; + NodeKind::Expansion { expanded, .. } => { + let mut curr = expanded; if !O::ACTIVE { - while let BoundKind::Expansion { - bound_expanded: next, + while let NodeKind::Expansion { + expanded: next, .. } = &curr.kind { @@ -1008,17 +1023,15 @@ impl VM { offset: &mut usize, ) -> Result<(), String> { match &pattern.kind { - BoundKind::Define { addr, .. } => { + NodeKind::Identifier { binding, .. } => { + let addr = match binding { + IdentifierBinding::Declaration { addr, .. } | IdentifierBinding::Reference(addr) => *addr, + }; let val = values.get(*offset).cloned().unwrap_or(Value::Void); *offset += 1; - self.set_value(*addr, val) + self.set_value(addr, val) } - BoundKind::Set { addr, .. } => { - let val = values.get(*offset).cloned().unwrap_or(Value::Void); - *offset += 1; - self.set_value(*addr, val) - } - BoundKind::Tuple { elements } => { + NodeKind::Tuple { elements } => { if let Some(sub_values) = values.get(*offset).and_then(|v| v.as_slice()) { *offset += 1; let mut sub_offset = 0;