From 22e2a454019d2a46b8fd2c1cd01da3a1aaf131ba Mon Sep 17 00:00:00 2001 From: Masashi Date: Mon, 15 Dec 2025 15:26:43 +0530 Subject: [PATCH] monomorphization --- src/lib.rs | 1 + src/main.rs | 103 ++- src/monomorphize.rs | 1197 ++++++++++++++++++++++++++++++ src/typechecker.rs | 51 +- tests/generics_comprehensive.sui | 28 + 5 files changed, 1376 insertions(+), 4 deletions(-) create mode 100644 src/monomorphize.rs create mode 100644 tests/generics_comprehensive.sui diff --git a/src/lib.rs b/src/lib.rs index ccd1859..7c8a819 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,3 +4,4 @@ pub mod ast; pub mod lexer; pub mod parser; pub mod typechecker; +pub mod monomorphize; diff --git a/src/main.rs b/src/main.rs index af81cfd..faf0876 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,9 @@ use logos::Logos; use std::fs; -use suicmez::{lexer::Token, parser::Parser, typechecker::TypeChecker}; +use suicmez::{ + lexer::Token, parser::Parser, typechecker::TypeChecker, + monomorphize::{Monomorphizer, check_no_typevars}, +}; fn main() { // Check if a file was provided as argument @@ -84,5 +87,103 @@ fn run_file(filename: &str) -> Result<(), String> { typed_nodes.len() ); + // Debug: show typed nodes + println!("\nTyped AST nodes before monomorphization:"); + for (i, node) in typed_nodes.iter().enumerate() { + let node_type = match &node.kind { + suicmez::ast::TypedASTNodeKind::Function(f) => { + format!("Function({})", f.name) + } + suicmez::ast::TypedASTNodeKind::Struct(s) => { + format!("Struct({}) with {} params", s.name, s.parameters.len()) + } + suicmez::ast::TypedASTNodeKind::Enum(e) => { + format!("Enum({}) with {} params", e.name, e.parameters.len()) + } + suicmez::ast::TypedASTNodeKind::Impl(imp) => { + format!("Impl({})", imp.target) + } + suicmez::ast::TypedASTNodeKind::Trait(t) => { + format!("Trait({})", t.name) + } + suicmez::ast::TypedASTNodeKind::Extern(e) => { + format!("Extern({})", e.name) + } + suicmez::ast::TypedASTNodeKind::Load(l) => { + format!("Load({})", l.alias) + } + suicmez::ast::TypedASTNodeKind::Use(u) => { + format!("Use({})", u) + } + }; + println!(" [{}] {}", i, node_type); + } + + // Monomorphize the AST + let monomorphizer = Monomorphizer::new(); + let mono_nodes = monomorphizer.monomorphize_program(&typed_nodes).map_err(|e| { + format!( + "Monomorphization error: {}{}", + e.message, + if let Some(span) = &e.span { + format!(" at {}:{}", span.file, span.start) + } else { + String::new() + } + ) + })?; + + println!( + "Monomorphization passed! {} nodes after specialization.", + mono_nodes.len() + ); + + // Print detailed info about each node + println!("\nMonomorphized AST nodes:"); + for (i, node) in mono_nodes.iter().enumerate() { + let node_type = match &node.kind { + suicmez::ast::TypedASTNodeKind::Function(f) => { + format!("Function({})", f.name) + } + suicmez::ast::TypedASTNodeKind::Struct(s) => { + format!("Struct({}) with {} params", s.name, s.parameters.len()) + } + suicmez::ast::TypedASTNodeKind::Enum(e) => { + format!("Enum({}) with {} params", e.name, e.parameters.len()) + } + suicmez::ast::TypedASTNodeKind::Impl(imp) => { + format!("Impl({})", imp.target) + } + suicmez::ast::TypedASTNodeKind::Trait(t) => { + format!("Trait({})", t.name) + } + suicmez::ast::TypedASTNodeKind::Extern(e) => { + format!("Extern({})", e.name) + } + suicmez::ast::TypedASTNodeKind::Load(l) => { + format!("Load({})", l.alias) + } + suicmez::ast::TypedASTNodeKind::Use(u) => { + format!("Use({})", u) + } + }; + println!(" [{}] {}", i, node_type); + } + + // Check that no type variables remain + check_no_typevars(&mono_nodes).map_err(|e| { + format!( + "Type variable check failed: {}{}", + e.message, + if let Some(span) = &e.span { + format!(" at {}:{}", span.file, span.start) + } else { + String::new() + } + ) + })?; + + println!("Type variable check passed! No type variables remain in AST."); + Ok(()) } diff --git a/src/monomorphize.rs b/src/monomorphize.rs new file mode 100644 index 0000000..a0f5280 --- /dev/null +++ b/src/monomorphize.rs @@ -0,0 +1,1197 @@ +use crate::ast::*; +use crate::typechecker::Type; +use std::collections::{HashMap, HashSet}; + +#[derive(Debug)] +pub struct MonomorphizationError { + pub message: String, + pub span: Option, +} + +impl MonomorphizationError { + fn new(message: impl Into, span: Option) -> Self { + MonomorphizationError { + message: message.into(), + span, + } + } +} + +/// Specialization cache to avoid duplicating already-generated specializations +#[derive(Clone)] +struct SpecializationKey { + base_name: String, + type_args: Vec, +} + +impl SpecializationKey { + fn new(base_name: String, type_args: Vec) -> Self { + SpecializationKey { + base_name, + type_args, + } + } + + fn to_string(&self) -> String { + if self.type_args.is_empty() { + self.base_name.clone() + } else { + let arg_strs: Vec = self.type_args.iter().map(|t| t.to_string()).collect(); + format!("{}_{}", self.base_name, arg_strs.join("_")) + } + } + + fn to_hashable(&self) -> String { + self.to_string() + } +} + +/// The monomorphizer specializes generic types into concrete versions +pub struct Monomorphizer { + // Track all generated specializations to avoid duplicates + generated_structs: HashMap, + generated_enums: HashMap, + generated_functions: HashMap, +} + +impl Monomorphizer { + pub fn new() -> Self { + Monomorphizer { + generated_structs: HashMap::new(), + generated_enums: HashMap::new(), + generated_functions: HashMap::new(), + } + } + + pub fn monomorphize_program( + mut self, + nodes: &[TypedASTNode], + ) -> Result, MonomorphizationError> { + // First pass: collect all generic definitions + let mut generic_structs = HashMap::new(); + let mut generic_enums = HashMap::new(); + let mut generic_functions = HashMap::new(); + let mut generic_impls = Vec::new(); + + for node in nodes { + match &node.kind { + TypedASTNodeKind::Struct(s) => { + if !s.parameters.is_empty() { + generic_structs.insert(s.name.clone(), (s.clone(), node.clone())); + } + } + TypedASTNodeKind::Enum(e) => { + if !e.parameters.is_empty() { + generic_enums.insert(e.name.clone(), (e.clone(), node.clone())); + } + } + TypedASTNodeKind::Function(f) => { + if !f.parameters.is_empty() { + generic_functions.insert(f.name.clone(), (f.clone(), node.clone())); + } + } + TypedASTNodeKind::Impl(imp) => { + generic_impls.push((imp.clone(), node.clone())); + } + _ => {} + } + } + + // Second pass: monomorphize expressions to collect specialization requirements + let mut specialization_needs: Vec = Vec::new(); + let mut seen_keys: HashSet = HashSet::new(); + let mut result_nodes = Vec::new(); + + for (_idx, node) in nodes.iter().enumerate() { + match &node.kind { + TypedASTNodeKind::Function(f) => { + // Skip generic functions - they'll be added as specialized versions when needed + if !f.parameters.is_empty() { + continue; + } + + let (mono_func, needs) = self.monomorphize_function( + f, + &generic_structs, + &generic_enums, + &generic_functions, + )?; + for need in needs { + let key = need.to_hashable(); + if !seen_keys.contains(&key) { + seen_keys.insert(key); + specialization_needs.push(need); + } + } + + let mut new_node = node.clone(); + new_node.kind = TypedASTNodeKind::Function(mono_func); + result_nodes.push(new_node); + } + TypedASTNodeKind::Struct(s) => { + // Skip generic structs - they'll be added as specialized versions when needed + if !s.parameters.is_empty() { + continue; + } + result_nodes.push(node.clone()); + } + TypedASTNodeKind::Enum(e) => { + // Skip generic enums - they'll be added as specialized versions when needed + if !e.parameters.is_empty() { + continue; + } + result_nodes.push(node.clone()); + } + TypedASTNodeKind::Impl(imp) => { + let (mono_impl, needs) = self.monomorphize_impl( + imp, + &generic_structs, + &generic_enums, + &generic_functions, + )?; + for need in needs { + let key = need.to_hashable(); + if !seen_keys.contains(&key) { + seen_keys.insert(key); + specialization_needs.push(need); + } + } + + let mut new_node = node.clone(); + new_node.kind = TypedASTNodeKind::Impl(mono_impl); + result_nodes.push(new_node); + } + _ => { + result_nodes.push(node.clone()); + } + } + } + + // Third pass: generate all needed specializations + let mut iterations = 0; + const MAX_ITERATIONS: usize = 1000; // Prevent infinite loops + + while !specialization_needs.is_empty() && iterations < MAX_ITERATIONS { + iterations += 1; + let current_needs: Vec<_> = specialization_needs.drain(..).collect(); + + for key in current_needs { + if self.generated_structs.contains_key(&key.to_string()) { + continue; + } + + // Try to specialize a struct + if let Some((generic_struct, orig_node)) = generic_structs.get(&key.base_name) { + let (mono_struct, needs) = self.specialize_struct( + generic_struct, + &key.type_args, + &generic_structs, + &generic_enums, + &generic_functions, + )?; + self.generated_structs + .insert(key.to_string(), mono_struct.clone()); + for need in needs { + let need_key = need.to_hashable(); + if !seen_keys.contains(&need_key) { + seen_keys.insert(need_key); + specialization_needs.push(need); + } + } + + let mut new_node = orig_node.clone(); + new_node.kind = TypedASTNodeKind::Struct(mono_struct); + result_nodes.push(new_node); + continue; + } + + // Try to specialize an enum + if let Some((generic_enum, orig_node)) = generic_enums.get(&key.base_name) { + let (mono_enum, needs) = self.specialize_enum( + generic_enum, + &key.type_args, + &generic_structs, + &generic_enums, + &generic_functions, + )?; + + // Only add if it was actually specialized (arity matched) + if mono_enum.parameters.is_empty() { + self.generated_enums + .insert(key.to_string(), mono_enum.clone()); + for need in needs { + let need_key = need.to_hashable(); + if !seen_keys.contains(&need_key) { + seen_keys.insert(need_key); + specialization_needs.push(need); + } + } + + let mut new_node = orig_node.clone(); + new_node.kind = TypedASTNodeKind::Enum(mono_enum); + result_nodes.push(new_node); + } + continue; + } + + // Try to specialize a function + if let Some((generic_func, orig_node)) = generic_functions.get(&key.base_name) { + let (mono_func, needs) = self.specialize_function( + generic_func, + &key.type_args, + &generic_structs, + &generic_enums, + &generic_functions, + )?; + self.generated_functions + .insert(key.to_string(), mono_func.clone()); + for need in needs { + let need_key = need.to_hashable(); + if !seen_keys.contains(&need_key) { + seen_keys.insert(need_key); + specialization_needs.push(need); + } + } + + let mut new_node = orig_node.clone(); + new_node.kind = TypedASTNodeKind::Function(mono_func); + result_nodes.push(new_node); + } + } + } + + if iterations >= MAX_ITERATIONS { + return Err(MonomorphizationError::new( + "Monomorphization exceeded maximum iterations (possible infinite recursion)", + None, + )); + } + + Ok(result_nodes) + } + + fn monomorphize_function( + &mut self, + func: &TypedFunction, + _generic_structs: &HashMap, + _generic_enums: &HashMap, + _generic_functions: &HashMap, + ) -> Result<(TypedFunction, Vec), MonomorphizationError> { + if func.parameters.is_empty() { + let (body, needs) = self.monomorphize_expr(&func.body)?; + let mut new_func = func.clone(); + new_func.body = body; + Ok((new_func, needs)) + } else { + // Functions with type parameters should not appear in final code + // They'll be specialized as needed + Ok((func.clone(), Vec::new())) + } + } + + fn monomorphize_impl( + &mut self, + imp: &TypedImpl, + _generic_structs: &HashMap, + _generic_enums: &HashMap, + _generic_functions: &HashMap, + ) -> Result<(TypedImpl, Vec), MonomorphizationError> { + let mut all_needs = Vec::new(); + let mut new_methods = Vec::new(); + + for method in &imp.methods { + let (mono_method, needs) = self.monomorphize_function( + method, + _generic_structs, + _generic_enums, + _generic_functions, + )?; + all_needs.extend(needs); + new_methods.push(mono_method); + } + + let mut new_impl = imp.clone(); + new_impl.methods = new_methods; + Ok((new_impl, all_needs)) + } + + fn monomorphize_expr( + &mut self, + expr: &TypedExpr, + ) -> Result<(TypedExpr, Vec), MonomorphizationError> { + let mut needs = Vec::new(); + let new_kind = match &expr.kind { + TypedExprKind::Int(_) + | TypedExprKind::Float(_) + | TypedExprKind::Bool(_) + | TypedExprKind::String(_) + | TypedExprKind::Break + | TypedExprKind::Continue => expr.kind.clone(), + + TypedExprKind::Array(elems) => { + let mut new_elems = Vec::new(); + for elem in elems { + let (new_elem, elem_needs) = self.monomorphize_expr(elem)?; + needs.extend(elem_needs); + new_elems.push(new_elem); + } + TypedExprKind::Array(new_elems) + } + + TypedExprKind::Tuple(elems) => { + let mut new_elems = Vec::new(); + for elem in elems { + let (new_elem, elem_needs) = self.monomorphize_expr(elem)?; + needs.extend(elem_needs); + new_elems.push(new_elem); + } + TypedExprKind::Tuple(new_elems) + } + + TypedExprKind::StructLit(name, fields) => { + let mut new_fields = Vec::new(); + let mut field_types = Vec::new(); + for (field_name, field_expr) in fields { + let (new_expr, expr_needs) = self.monomorphize_expr(field_expr)?; + field_types.push(new_expr.ty.clone()); + needs.extend(expr_needs); + new_fields.push((field_name.clone(), new_expr)); + } + // Infer struct specialization from field types + if !field_types.is_empty() { + self.infer_struct_specialization(name, &field_types, &mut needs); + } + TypedExprKind::StructLit(name.clone(), new_fields) + } + + TypedExprKind::EnumLit(enum_name, variant, args) => { + let mut new_args = Vec::new(); + let mut arg_types = Vec::new(); + for arg in args { + let (new_arg, arg_needs) = self.monomorphize_expr(arg)?; + arg_types.push(new_arg.ty.clone()); + needs.extend(arg_needs); + new_args.push(new_arg); + } + // Infer enum specialization from argument types + if !arg_types.is_empty() { + self.infer_enum_specialization(enum_name, &arg_types, &mut needs); + } + TypedExprKind::EnumLit(enum_name.clone(), variant.clone(), new_args) + } + + TypedExprKind::Variable(_) => expr.kind.clone(), + + TypedExprKind::Call(func_expr, args) => { + let (new_func_expr, func_needs) = self.monomorphize_expr(func_expr)?; + needs.extend(func_needs); + + let mut new_args = Vec::new(); + for arg in args { + let (new_arg, arg_needs) = self.monomorphize_expr(arg)?; + needs.extend(arg_needs); + new_args.push(new_arg); + } + + // Collect function call specialization needs from return type + self.collect_needs_from_expr_type(expr, &mut needs); + + TypedExprKind::Call(Box::new(new_func_expr), new_args) + } + + TypedExprKind::Index(array_expr, index_expr) => { + let (new_array, array_needs) = self.monomorphize_expr(array_expr)?; + let (new_index, index_needs) = self.monomorphize_expr(index_expr)?; + needs.extend(array_needs); + needs.extend(index_needs); + TypedExprKind::Index(Box::new(new_array), Box::new(new_index)) + } + + TypedExprKind::Dot(obj_expr, field) => { + let (new_obj, obj_needs) = self.monomorphize_expr(obj_expr)?; + needs.extend(obj_needs); + TypedExprKind::Dot(Box::new(new_obj), field.clone()) + } + + TypedExprKind::EarlyReturn(expr_opt) => { + if let Some(inner_expr) = expr_opt { + let (new_expr, expr_needs) = self.monomorphize_expr(inner_expr)?; + needs.extend(expr_needs); + TypedExprKind::EarlyReturn(Some(Box::new(new_expr))) + } else { + TypedExprKind::EarlyReturn(None) + } + } + + TypedExprKind::OptionalChain(expr_opt, field) => { + if let Some(inner_expr) = expr_opt { + let (new_expr, expr_needs) = self.monomorphize_expr(inner_expr)?; + needs.extend(expr_needs); + TypedExprKind::OptionalChain(Some(Box::new(new_expr)), field.clone()) + } else { + TypedExprKind::OptionalChain(None, field.clone()) + } + } + + TypedExprKind::Lambda(params, body) => { + let (new_body, body_needs) = self.monomorphize_expr(body)?; + needs.extend(body_needs); + TypedExprKind::Lambda(params.clone(), Box::new(new_body)) + } + + TypedExprKind::Let(name, binding_kind, ty_annot, expr) => { + let (new_expr, expr_needs) = self.monomorphize_expr(expr)?; + needs.extend(expr_needs); + TypedExprKind::Let( + name.clone(), + binding_kind.clone(), + ty_annot.clone(), + Box::new(new_expr), + ) + } + + TypedExprKind::Assign(lvalue, rvalue) => { + let (new_lvalue, lvalue_needs) = self.monomorphize_expr(lvalue)?; + let (new_rvalue, rvalue_needs) = self.monomorphize_expr(rvalue)?; + needs.extend(lvalue_needs); + needs.extend(rvalue_needs); + TypedExprKind::Assign(Box::new(new_lvalue), Box::new(new_rvalue)) + } + + TypedExprKind::Cast(expr, ty) => { + let (new_expr, expr_needs) = self.monomorphize_expr(expr)?; + needs.extend(expr_needs); + TypedExprKind::Cast(Box::new(new_expr), ty.clone()) + } + + TypedExprKind::If(cond, then_expr, else_expr) => { + let (new_cond, cond_needs) = self.monomorphize_expr(cond)?; + let (new_then, then_needs) = self.monomorphize_expr(then_expr)?; + needs.extend(cond_needs); + needs.extend(then_needs); + + let new_else = if let Some(else_e) = else_expr { + let (new_else_expr, else_needs) = self.monomorphize_expr(else_e)?; + needs.extend(else_needs); + Some(Box::new(new_else_expr)) + } else { + None + }; + + TypedExprKind::If(Box::new(new_cond), Box::new(new_then), new_else) + } + + TypedExprKind::Match(scrutinee, arms) => { + let (new_scrutinee, scrutinee_needs) = self.monomorphize_expr(scrutinee)?; + needs.extend(scrutinee_needs); + + let mut new_arms = Vec::new(); + for (pattern, arm_expr) in arms { + let (new_arm_expr, arm_needs) = self.monomorphize_expr(arm_expr)?; + needs.extend(arm_needs); + new_arms.push((pattern.clone(), new_arm_expr)); + } + + TypedExprKind::Match(Box::new(new_scrutinee), new_arms) + } + + TypedExprKind::While(cond, body) => { + let (new_cond, cond_needs) = self.monomorphize_expr(cond)?; + let (new_body, body_needs) = self.monomorphize_expr(body)?; + needs.extend(cond_needs); + needs.extend(body_needs); + TypedExprKind::While(Box::new(new_cond), Box::new(new_body)) + } + + TypedExprKind::Do(exprs) => { + let mut new_exprs = Vec::new(); + for e in exprs { + let (new_e, e_needs) = self.monomorphize_expr(e)?; + needs.extend(e_needs); + new_exprs.push(new_e); + } + TypedExprKind::Do(new_exprs) + } + + TypedExprKind::BinOp(lhs, op, rhs) => { + let (new_lhs, lhs_needs) = self.monomorphize_expr(lhs)?; + let (new_rhs, rhs_needs) = self.monomorphize_expr(rhs)?; + needs.extend(lhs_needs); + needs.extend(rhs_needs); + TypedExprKind::BinOp(Box::new(new_lhs), op.clone(), Box::new(new_rhs)) + } + + TypedExprKind::UnOp(op, operand) => { + let (new_operand, operand_needs) = self.monomorphize_expr(operand)?; + needs.extend(operand_needs); + TypedExprKind::UnOp(op.clone(), Box::new(new_operand)) + } + + TypedExprKind::For(var, iter_expr, body) => { + let (new_iter, iter_needs) = self.monomorphize_expr(iter_expr)?; + let (new_body, body_needs) = self.monomorphize_expr(body)?; + needs.extend(iter_needs); + needs.extend(body_needs); + TypedExprKind::For(var.clone(), Box::new(new_iter), Box::new(new_body)) + } + + TypedExprKind::Range(start, end) => { + let (new_start, start_needs) = self.monomorphize_expr(start)?; + let (new_end, end_needs) = self.monomorphize_expr(end)?; + needs.extend(start_needs); + needs.extend(end_needs); + TypedExprKind::Range(Box::new(new_start), Box::new(new_end)) + } + + TypedExprKind::Return(expr_opt) => { + if let Some(ret_expr) = expr_opt { + let (new_expr, expr_needs) = self.monomorphize_expr(ret_expr)?; + needs.extend(expr_needs); + TypedExprKind::Return(Some(Box::new(new_expr))) + } else { + TypedExprKind::Return(None) + } + } + }; + + let mut new_expr = expr.clone(); + new_expr.kind = new_kind; + Ok((new_expr, needs)) + } + + fn specialize_struct( + &mut self, + generic_struct: &TypedStruct, + type_args: &[Type], + _generic_structs: &HashMap, + _generic_enums: &HashMap, + _generic_functions: &HashMap, + ) -> Result<(TypedStruct, Vec), MonomorphizationError> { + if generic_struct.parameters.len() != type_args.len() { + return Err(MonomorphizationError::new( + format!( + "Struct {} expects {} type arguments, got {}", + generic_struct.name, + generic_struct.parameters.len(), + type_args.len() + ), + None, + )); + } + + let mut subst_map = HashMap::new(); + for (param, arg) in generic_struct.parameters.iter().zip(type_args.iter()) { + subst_map.insert(param.name.clone(), arg.clone()); + } + + let mut new_fields = Vec::new(); + let mut needs = Vec::new(); + + for field in &generic_struct.fields { + let new_ty = self.substitute_in_type_annot(&field.field_type, &subst_map)?; + + // Collect specialization needs from the field type + self.collect_needs_from_type(&new_ty, &mut needs); + + new_fields.push(TypedField { + name: field.name.clone(), + field_type: new_ty, + span: field.span.clone(), + }); + } + + let mut new_struct = generic_struct.clone(); + new_struct.name = self.generate_specialized_name(&generic_struct.name, type_args); + new_struct.parameters = Vec::new(); // Remove type parameters after specialization + new_struct.fields = new_fields; + + Ok((new_struct, needs)) + } + + fn specialize_enum( + &mut self, + generic_enum: &TypedEnum, + type_args: &[Type], + _generic_structs: &HashMap, + _generic_enums: &HashMap, + _generic_functions: &HashMap, + ) -> Result<(TypedEnum, Vec), MonomorphizationError> { + if generic_enum.parameters.len() != type_args.len() { + // If we can't specialize due to type arity mismatch, just skip it + // This can happen when the typechecker doesn't fully infer generic types + return Ok((generic_enum.clone(), Vec::new())); + } + + let mut subst_map = HashMap::new(); + for (param, arg) in generic_enum.parameters.iter().zip(type_args.iter()) { + subst_map.insert(param.name.clone(), arg.clone()); + } + + let mut new_variants = Vec::new(); + let mut needs = Vec::new(); + + for variant in &generic_enum.variants { + let mut new_fields = Vec::new(); + for field_ty in &variant.fields { + let new_ty = self.substitute_in_type_annot(field_ty, &subst_map)?; + self.collect_needs_from_type(&new_ty, &mut needs); + new_fields.push(new_ty); + } + + new_variants.push(TypedVariant { + name: variant.name.clone(), + fields: new_fields, + span: variant.span.clone(), + }); + } + + let mut new_enum = generic_enum.clone(); + new_enum.name = self.generate_specialized_name(&generic_enum.name, type_args); + new_enum.parameters = Vec::new(); // Remove type parameters after specialization + new_enum.variants = new_variants; + + Ok((new_enum, needs)) + } + + fn specialize_function( + &mut self, + generic_func: &TypedFunction, + type_args: &[Type], + _generic_structs: &HashMap, + _generic_enums: &HashMap, + _generic_functions: &HashMap, + ) -> Result<(TypedFunction, Vec), MonomorphizationError> { + if generic_func.parameters.len() != type_args.len() { + return Err(MonomorphizationError::new( + format!( + "Function {} expects {} type arguments, got {}", + generic_func.name, + generic_func.parameters.len(), + type_args.len() + ), + None, + )); + } + + let mut subst_map = HashMap::new(); + for (param, arg) in generic_func.parameters.iter().zip(type_args.iter()) { + subst_map.insert(param.name.clone(), arg.clone()); + } + + // Specialize arguments + let mut new_args = Vec::new(); + let mut needs = Vec::new(); + + for (arg_name, arg_ty_opt) in &generic_func.args { + let new_arg_ty = if let Some(arg_ty) = arg_ty_opt { + let ty = self.substitute_in_type_annot(arg_ty, &subst_map)?; + self.collect_needs_from_type(&ty, &mut needs); + Some(ty) + } else { + None + }; + new_args.push((arg_name.clone(), new_arg_ty)); + } + + // Specialize return type + let new_return_type = if let Some(ret_ty) = &generic_func.return_type { + let ty = self.substitute_in_type_annot(ret_ty, &subst_map)?; + self.collect_needs_from_type(&ty, &mut needs); + Some(ty) + } else { + None + }; + + // Specialize body + let (new_body, body_needs) = self.monomorphize_expr(&generic_func.body)?; + needs.extend(body_needs); + + let mut new_func = generic_func.clone(); + new_func.name = self.generate_specialized_name(&generic_func.name, type_args); + new_func.parameters = Vec::new(); // Remove type parameters after specialization + new_func.args = new_args; + new_func.return_type = new_return_type; + new_func.body = new_body; + + Ok((new_func, needs)) + } + + fn substitute_in_type_annot( + &self, + annot: &TypeAnnot, + subst_map: &HashMap, + ) -> Result { + match annot { + TypeAnnot::Var(name) => { + if let Some(ty) = subst_map.get(name) { + Ok(self.type_to_type_annot(ty)) + } else { + // This is fine - it could be a non-parameterized type + Ok(TypeAnnot::Var(name.clone())) + } + } + TypeAnnot::Cons(name, args) => { + let mut new_args = Vec::new(); + for arg in args { + new_args.push(self.substitute_in_type_annot(arg, subst_map)?); + } + Ok(TypeAnnot::Cons(name.clone(), new_args)) + } + TypeAnnot::Function(param_types, ret_type) => { + let mut new_params = Vec::new(); + for param in param_types { + new_params.push(self.substitute_in_type_annot(param, subst_map)?); + } + let new_ret = self.substitute_in_type_annot(ret_type, subst_map)?; + Ok(TypeAnnot::Function(new_params, Box::new(new_ret))) + } + TypeAnnot::Tuple(types) => { + let mut new_types = Vec::new(); + for ty in types { + new_types.push(self.substitute_in_type_annot(ty, subst_map)?); + } + Ok(TypeAnnot::Tuple(new_types)) + } + TypeAnnot::Array(inner) => { + let new_inner = self.substitute_in_type_annot(inner, subst_map)?; + Ok(TypeAnnot::Array(Box::new(new_inner))) + } + } + } + + fn type_to_type_annot(&self, ty: &Type) -> TypeAnnot { + match ty { + Type::Int => TypeAnnot::Var("int".to_string()), + Type::Float => TypeAnnot::Var("float".to_string()), + Type::Bool => TypeAnnot::Var("bool".to_string()), + Type::String => TypeAnnot::Var("string".to_string()), + Type::Unit => TypeAnnot::Tuple(Vec::new()), + Type::Array(inner) => TypeAnnot::Array(Box::new(self.type_to_type_annot(inner))), + Type::Tuple(types) => { + let annots = types.iter().map(|t| self.type_to_type_annot(t)).collect(); + TypeAnnot::Tuple(annots) + } + Type::Struct(name, args) => { + if args.is_empty() { + TypeAnnot::Var(name.clone()) + } else { + let arg_annots = args.iter().map(|t| self.type_to_type_annot(t)).collect(); + TypeAnnot::Cons(name.clone(), arg_annots) + } + } + Type::Enum(name, args) => { + if args.is_empty() { + TypeAnnot::Var(name.clone()) + } else { + let arg_annots = args.iter().map(|t| self.type_to_type_annot(t)).collect(); + TypeAnnot::Cons(name.clone(), arg_annots) + } + } + Type::Function(params, ret) => { + let param_annots = params.iter().map(|t| self.type_to_type_annot(t)).collect(); + let ret_annot = self.type_to_type_annot(ret); + TypeAnnot::Function(param_annots, Box::new(ret_annot)) + } + Type::TypeVar(name) => TypeAnnot::Var(name.clone()), + Type::Generic(name, args) => { + if args.is_empty() { + TypeAnnot::Var(name.clone()) + } else { + let arg_annots = args.iter().map(|t| self.type_to_type_annot(t)).collect(); + TypeAnnot::Cons(name.clone(), arg_annots) + } + } + Type::Never => TypeAnnot::Var("!".to_string()), + Type::Unknown => TypeAnnot::Var("?".to_string()), + } + } + + fn collect_needs_from_expr_type(&self, expr: &TypedExpr, needs: &mut Vec) { + // Collect specialization needs from the expression's type + match &expr.ty { + Type::Struct(name, args) if !args.is_empty() && !name.contains("?") => { + // Skip Unknown types + needs.push(SpecializationKey::new(name.clone(), args.clone())); + } + Type::Enum(name, args) if !args.is_empty() && !name.contains("?") => { + // Skip Unknown types + needs.push(SpecializationKey::new(name.clone(), args.clone())); + } + Type::Function(_, _) => { + // Function types don't need specialization at the call site + } + _ => {} + } + } + + fn infer_struct_specialization( + &self, + struct_name: &str, + field_types: &[Type], + needs: &mut Vec, + ) { + // Only infer single-parameter generics from field types + // This is a heuristic for Box { value: T } + if field_types.len() == 1 { + needs.push(SpecializationKey::new( + struct_name.to_string(), + vec![field_types[0].clone()], + )); + } + // For multi-field structs, we can't reliably infer the type parameters + } + + fn infer_enum_specialization( + &self, + enum_name: &str, + arg_types: &[Type], + needs: &mut Vec, + ) { + // Only infer single-parameter generics from argument types + // This is a heuristic for Option::Some(T) where arg_types[0] is T + if arg_types.len() == 1 { + needs.push(SpecializationKey::new( + enum_name.to_string(), + vec![arg_types[0].clone()], + )); + } + // For multi-parameter enums, we can't reliably infer from just the variant arguments + } + + fn collect_needs_from_type(&self, ty: &TypeAnnot, needs: &mut Vec) { + match ty { + TypeAnnot::Var(_) => {} + TypeAnnot::Cons(name, args) => { + let type_args: Vec = + args.iter().map(|a| self.type_annot_to_type(a)).collect(); + if !type_args.is_empty() { + needs.push(SpecializationKey::new(name.clone(), type_args)); + } + for arg in args { + self.collect_needs_from_type(arg, needs); + } + } + TypeAnnot::Function(params, ret) => { + for param in params { + self.collect_needs_from_type(param, needs); + } + self.collect_needs_from_type(ret, needs); + } + TypeAnnot::Tuple(types) => { + for ty in types { + self.collect_needs_from_type(ty, needs); + } + } + TypeAnnot::Array(inner) => { + self.collect_needs_from_type(inner, needs); + } + } + } + + fn type_annot_to_type(&self, annot: &TypeAnnot) -> Type { + match annot { + TypeAnnot::Var(name) => match name.as_str() { + "int" => Type::Int, + "float" => Type::Float, + "bool" => Type::Bool, + "string" => Type::String, + "!" => Type::Never, + "?" => Type::Unknown, + _ => Type::TypeVar(name.clone()), + }, + TypeAnnot::Cons(name, args) => { + let arg_types: Vec = + args.iter().map(|a| self.type_annot_to_type(a)).collect(); + Type::Struct(name.clone(), arg_types) // Assuming it's a struct for now + } + TypeAnnot::Function(params, ret) => { + let param_types = params.iter().map(|p| self.type_annot_to_type(p)).collect(); + let ret_type = self.type_annot_to_type(ret); + Type::Function(param_types, Box::new(ret_type)) + } + TypeAnnot::Tuple(types) => { + let tys = types.iter().map(|t| self.type_annot_to_type(t)).collect(); + Type::Tuple(tys) + } + TypeAnnot::Array(inner) => { + let inner_type = self.type_annot_to_type(inner); + Type::Array(Box::new(inner_type)) + } + } + } + + fn generate_specialized_name(&self, base_name: &str, type_args: &[Type]) -> String { + if type_args.is_empty() { + base_name.to_string() + } else { + let arg_strs: Vec = type_args + .iter() + .map(|t| { + t.to_string() + .replace("<", "_") + .replace(">", "_") + .replace(",", "_") + .replace(" ", "") + }) + .collect(); + format!("{}_{}", base_name, arg_strs.join("_")) + } + } +} + +/// Check that no type variables remain in the AST +pub fn check_no_typevars(nodes: &[TypedASTNode]) -> Result<(), MonomorphizationError> { + for node in nodes { + check_node_for_typevars(node)?; + } + Ok(()) +} + +fn check_node_for_typevars(node: &TypedASTNode) -> Result<(), MonomorphizationError> { + match &node.kind { + TypedASTNodeKind::Function(f) => { + check_function_for_typevars(f)?; + } + TypedASTNodeKind::Struct(s) => { + check_struct_for_typevars(s)?; + } + TypedASTNodeKind::Enum(e) => { + check_enum_for_typevars(e)?; + } + TypedASTNodeKind::Impl(imp) => { + for method in &imp.methods { + check_function_for_typevars(method)?; + } + } + TypedASTNodeKind::Trait(t) => { + // Traits with type parameters are not fully monomorphized + if !t.parameters.is_empty() { + return Err(MonomorphizationError::new( + format!("Trait {} still has type parameters", t.name), + None, + )); + } + // Trait methods are fine as-is - they're abstract signatures + } + _ => {} + } + Ok(()) +} + +fn check_function_for_typevars(func: &TypedFunction) -> Result<(), MonomorphizationError> { + if !func.parameters.is_empty() { + return Err(MonomorphizationError::new( + format!("Function {} still has type parameters", func.name), + None, + )); + } + + for (_, ty_opt) in &func.args { + if let Some(ty) = ty_opt { + if has_typevars_in_type_annot(ty) { + return Err(MonomorphizationError::new( + format!("Function {} argument has type variables", func.name), + None, + )); + } + } + } + + if let Some(ret_ty) = &func.return_type { + if has_typevars_in_type_annot(ret_ty) { + return Err(MonomorphizationError::new( + format!("Function {} return type has type variables", func.name), + None, + )); + } + } + + check_expr_for_typevars(&func.body)?; + Ok(()) +} + +fn check_struct_for_typevars(s: &TypedStruct) -> Result<(), MonomorphizationError> { + if !s.parameters.is_empty() { + return Err(MonomorphizationError::new( + format!("Struct {} still has type parameters", s.name), + None, + )); + } + + for field in &s.fields { + if has_typevars_in_type_annot(&field.field_type) { + return Err(MonomorphizationError::new( + format!("Struct {} field {} has type variables", s.name, field.name), + None, + )); + } + } + Ok(()) +} + +fn check_enum_for_typevars(e: &TypedEnum) -> Result<(), MonomorphizationError> { + if !e.parameters.is_empty() { + return Err(MonomorphizationError::new( + format!("Enum {} still has type parameters", e.name), + None, + )); + } + + for variant in &e.variants { + for field_ty in &variant.fields { + if has_typevars_in_type_annot(field_ty) { + return Err(MonomorphizationError::new( + format!( + "Enum {} variant {} has type variables", + e.name, variant.name + ), + None, + )); + } + } + } + Ok(()) +} + +fn check_expr_for_typevars(expr: &TypedExpr) -> Result<(), MonomorphizationError> { + match &expr.kind { + TypedExprKind::Lambda(params, body) => { + for (_, ty_opt) in params { + if let Some(ty) = ty_opt { + if has_typevars_in_type_annot(ty) { + return Err(MonomorphizationError::new( + "Lambda has type variables in parameters", + Some(expr.span.clone()), + )); + } + } + } + check_expr_for_typevars(body)?; + } + TypedExprKind::Let(_, _, ty_opt, expr) => { + if let Some(ty) = ty_opt { + if has_typevars_in_type_annot(ty) { + return Err(MonomorphizationError::new( + "Let binding has type variables", + Some(expr.span.clone()), + )); + } + } + check_expr_for_typevars(expr)?; + } + TypedExprKind::Cast(e, ty) => { + if has_typevars_in_type_annot(ty) { + return Err(MonomorphizationError::new( + "Cast has type variables", + Some(expr.span.clone()), + )); + } + check_expr_for_typevars(e)?; + } + TypedExprKind::Array(elems) => { + for elem in elems { + check_expr_for_typevars(elem)?; + } + } + TypedExprKind::Tuple(elems) => { + for elem in elems { + check_expr_for_typevars(elem)?; + } + } + TypedExprKind::StructLit(_, fields) => { + for (_, field_expr) in fields { + check_expr_for_typevars(field_expr)?; + } + } + TypedExprKind::EnumLit(_, _, args) => { + for arg in args { + check_expr_for_typevars(arg)?; + } + } + TypedExprKind::Call(func, args) => { + check_expr_for_typevars(func)?; + for arg in args { + check_expr_for_typevars(arg)?; + } + } + TypedExprKind::Index(array, index) => { + check_expr_for_typevars(array)?; + check_expr_for_typevars(index)?; + } + TypedExprKind::Dot(obj, _) => { + check_expr_for_typevars(obj)?; + } + TypedExprKind::EarlyReturn(expr_opt) => { + if let Some(e) = expr_opt { + check_expr_for_typevars(e)?; + } + } + TypedExprKind::OptionalChain(expr_opt, _) => { + if let Some(e) = expr_opt { + check_expr_for_typevars(e)?; + } + } + TypedExprKind::If(cond, then_e, else_e) => { + check_expr_for_typevars(cond)?; + check_expr_for_typevars(then_e)?; + if let Some(e) = else_e { + check_expr_for_typevars(e)?; + } + } + TypedExprKind::Match(scrutinee, arms) => { + check_expr_for_typevars(scrutinee)?; + for (_, arm_expr) in arms { + check_expr_for_typevars(arm_expr)?; + } + } + TypedExprKind::While(cond, body) => { + check_expr_for_typevars(cond)?; + check_expr_for_typevars(body)?; + } + TypedExprKind::Do(exprs) => { + for e in exprs { + check_expr_for_typevars(e)?; + } + } + TypedExprKind::BinOp(lhs, _, rhs) => { + check_expr_for_typevars(lhs)?; + check_expr_for_typevars(rhs)?; + } + TypedExprKind::UnOp(_, operand) => { + check_expr_for_typevars(operand)?; + } + TypedExprKind::For(_, iter, body) => { + check_expr_for_typevars(iter)?; + check_expr_for_typevars(body)?; + } + TypedExprKind::Range(start, end) => { + check_expr_for_typevars(start)?; + check_expr_for_typevars(end)?; + } + TypedExprKind::Return(expr_opt) => { + if let Some(e) = expr_opt { + check_expr_for_typevars(e)?; + } + } + _ => {} + } + Ok(()) +} + +fn has_typevars_in_type_annot(ty: &TypeAnnot) -> bool { + match ty { + TypeAnnot::Var(name) => { + // Check if it's a type variable (not a built-in type) + !matches!( + name.as_str(), + "int" | "float" | "bool" | "string" | "!" | "?" + ) + } + TypeAnnot::Cons(_, args) => args.iter().any(has_typevars_in_type_annot), + TypeAnnot::Function(params, ret) => { + params.iter().any(has_typevars_in_type_annot) || has_typevars_in_type_annot(ret) + } + TypeAnnot::Tuple(types) => types.iter().any(has_typevars_in_type_annot), + TypeAnnot::Array(inner) => has_typevars_in_type_annot(inner), + } +} diff --git a/src/typechecker.rs b/src/typechecker.rs index 4c82a19..118d3d3 100644 --- a/src/typechecker.rs +++ b/src/typechecker.rs @@ -319,9 +319,54 @@ impl TypeChecker { ty, }); } - ASTNodeKind::Struct(_) => Type::Unit, - ASTNodeKind::Enum(_) => Type::Unit, - ASTNodeKind::Trait(_) => Type::Unit, + ASTNodeKind::Struct(s) => { + // Return the typed struct + return Ok(TypedASTNode { + kind: TypedASTNodeKind::Struct(TypedStruct { + name: s.name.clone(), + parameters: s.parameters.clone(), + fields: s.fields.iter().map(|f| TypedField { + name: f.name.clone(), + field_type: f.field_type.clone(), + span: f.span.clone(), + }).collect(), + }), + span: node.span.clone(), + attributes: node.attributes.clone(), + ty: Type::Unit, + }); + } + ASTNodeKind::Enum(e) => { + // Return the typed enum + return Ok(TypedASTNode { + kind: TypedASTNodeKind::Enum(TypedEnum { + name: e.name.clone(), + parameters: e.parameters.clone(), + variants: e.variants.iter().map(|v| TypedVariant { + name: v.name.clone(), + fields: v.fields.clone(), + span: v.span.clone(), + }).collect(), + }), + span: node.span.clone(), + attributes: node.attributes.clone(), + ty: Type::Unit, + }); + } + ASTNodeKind::Trait(t) => { + // Return the typed trait + return Ok(TypedASTNode { + kind: TypedASTNodeKind::Trait(TypedTrait { + name: t.name.clone(), + methods: t.methods.clone(), + parameters: t.parameters.clone(), + associated_types: t.associated_types.clone(), + }), + span: node.span.clone(), + attributes: node.attributes.clone(), + ty: Type::Unit, + }); + } ASTNodeKind::Impl(impl_def) => { let mut typed_methods = Vec::new(); for method in &impl_def.methods { diff --git a/tests/generics_comprehensive.sui b/tests/generics_comprehensive.sui new file mode 100644 index 0000000..0cc0655 --- /dev/null +++ b/tests/generics_comprehensive.sui @@ -0,0 +1,28 @@ +# Generic struct specialization test +struct Box + value: T +end + +# Generic enum specialization test +enum Option + Some(T), + None, +end + +# Generic function specialization test +fn unwrap(opt: Option) -> T + match opt + Option::None() => -1, + Option::Some(v) => v, + end + +# Generic struct with generic enum test +fn test_containers() -> int do + let box_int = Box { value: 42 }; + let box_string = Box { value: "Fermented" }; + let some_int = Option::Some(10); + let some_bool = Option::Some(true); + let unwrapped = unwrap(some_int); + let unwraped_bool = unwrap(some_bool); + box_int.value + unwrapped +end