diff --git a/book/build.py b/book/build.py index 40303899..23473e30 100755 --- a/book/build.py +++ b/book/build.py @@ -6,12 +6,11 @@ 4. Runs zensical build """ -import subprocess -from pathlib import Path -import urllib.parse import os +import subprocess import sys - +import urllib.parse +from pathlib import Path SCRIPT_DIR = Path(__file__).parent.resolve() diff --git a/book/src/basics/structs.md b/book/src/basics/structs.md index 4e2e6446..488ff154 100644 --- a/book/src/basics/structs.md +++ b/book/src/basics/structs.md @@ -33,6 +33,31 @@ let side_length = cbrt(mass / tungsten.density) -> cm print("A tungsten cube with a mass of {mass} has a side length of {side_length:.2}.") ``` +Field access is often enough for simple data, but structs can also define behavior with methods: + +```nbt +struct Element { + name: String, + atomic_number: Scalar, + density: MassDensity, + + fn tungsten() -> Self = Self { + name: "Tungsten", + atomic_number: 74, + density: 19.25 g/cm³, + } + + fn cube_side_length(self, mass: Mass) -> Length = + cbrt(mass / self.density) +} + +let tungsten = Element::tungsten() +let side_length = tungsten.cube_side_length(1 kg) -> cm +print("A 1 kg tungsten cube has side length {side_length:.2}.") +``` + +Note that constructor methods can return `Self` or `Element` in this case, but the former is more concise and is preferred by convention. + ## Generic structs Structs can be generic over type parameters. Type parameters are declared in angle brackets after the struct name: @@ -57,3 +82,26 @@ struct Vec { let position = Vec { x: 1 m, y: 2 m } let velocity: Vec = Vec { x: 1 m/s, y: 2 m/s } ``` + +Structs with generic type parameters can also have methods that use those type parameters, and methods can introduce additional type parameters of their own: + +```nbt +struct Vec { + x: X, + y: X, + + fn scale(self, factor: Scalar) -> Self = + Vec { x: self.x * factor, y: self.y * factor } + + fn dot_product(self, other: Vec) -> X * Y = + self.x * other.x + self.y * other.y +} + +let v1 = Vec { x: 1 m, y: 2 m } +let v2 = Vec { x: 3 m, y: 4 m } +let v2_cm = Vec { x: 300 cm, y: 400 cm } + +let v3 = v1.scale(2) # Vec { x: 2 m, y: 4 m } +let dp_m = v1.dot_product(v2) # 11 m² +let dp_cm = v1.dot_product(v2_cm) # 110_000 cm² +``` diff --git a/book/src/examples/example-numbat_syntax.md b/book/src/examples/example-numbat_syntax.md index 35611e6d..37ef306a 100644 --- a/book/src/examples/example-numbat_syntax.md +++ b/book/src/examples/example-numbat_syntax.md @@ -136,6 +136,9 @@ struct Element { # Define a struct name: String, atomic_number: Scalar, density: MassDensity, + + fn neptunium() -> Self = + Self { name: "Neptunium", atomic_number: 93, density: 20.45 g/cm³ } } let hydrogen = Element { # Instantiate it @@ -146,8 +149,23 @@ let hydrogen = Element { # Instantiate it hydrogen.density # Access the field of a struct +let neptunium = # Instantiate using a constructor method + Element::neptunium() +print(neptunium) + struct Vec2 { # A generic struct with type parameter x: D, y: D, + + # Methods can have their own type generics and inherit the struct's generics + fn new(x: D, y: D) -> Self = Self { x: x, y: y } + fn scale(self, factor: Scalar) -> Self = Self { x: self.x * factor, y: self.y * factor } + fn dot(self, other: Vec2) -> D * E = self.x * other.x + self.y * other.y } + +let v = Vec2::new(3 m, 4 m) # Constructor call +let w = Vec2::new(300 cm, 400 cm) +let scaled_v = v.scale(2) # Returning `Self` creates a new value (no in-place mutation) + +assert_eq(v.dot(w) -> m², 25 m²) # Method generic + unit conversion ``` diff --git a/book/src/examples/example-paper_size.md b/book/src/examples/example-paper_size.md index 5495dc61..bccab837 100644 --- a/book/src/examples/example-paper_size.md +++ b/book/src/examples/example-paper_size.md @@ -2,7 +2,7 @@ # Paper sizes -[:material-play-circle: Run this example](https://numbat.dev/?q=%23+Compute+ISO+216+paper+sizes+for+the+A+series%0A%23%0A%23+https%3A%2F%2Fen.wikipedia.org%2Fwiki%2FISO_216%0A%0Astruct+PaperSize+%7B%0A++++width%3A+Length%2C%0A++++height%3A+Length%2C%0A%7D%0A%0Afn+paper_size_A%28n%3A+Scalar%29+-%3E+PaperSize+%3D%0A++if+n+%3D%3D+0%0A++++then%0A++++++PaperSize+%7B%0A++++++++width%3A+841+mm%2C%0A++++++++height%3A+1189+mm%0A++++++%7D%0A++++else%0A++++++PaperSize+%7B%0A++++++++width%3A+floor_in%28mm%2C+paper_size_A%28n+-+1%29.height+%2F+2%29%2C%0A++++++++height%3A+paper_size_A%28n+-+1%29.width%2C%0A++++++%7D%0A%0A%0Afn+paper_area%28size%3A+PaperSize%29+-%3E+Area+%3D%0A++++size.width+%2A+size.height%0A%0A%0Afn+size_as_string%28size%3A+PaperSize%29+%3D+%22%7Bsize.width%3A%3E4%7D+%C3%97+%7Bsize.height%3A%3E5%7D+++%7Bpaper_area%28size%29+-%3E+cm%C2%B2%3A%3E6.1f%7D%22%0Afn+row%28n%29+%3D+%22A%7Bn%3A%3C3%7D+++%7Bsize_as_string%28paper_size_A%28n%29%29%7D%22%0A%0Aprint%28%22Name++++Width+++++Height++++++++Area++%22%29%0Aprint%28%22----+++-------+++--------+++----------%22%29%0Aprint%28join%28map%28row%2C+range%280%2C+10%29%29%2C+%22%5Cn%22%29%29%0A){ .md-button .md-button--primary } +[:material-play-circle: Run this example](https://numbat.dev/?q=%23+Compute+ISO+216+paper+sizes+for+the+A+series%0A%23%0A%23+https%3A%2F%2Fen.wikipedia.org%2Fwiki%2FISO_216%0A%0Astruct+PaperSize+%7B%0A++++width%3A+Length%2C%0A++++height%3A+Length%2C%0A%0A++++fn+a%28n%3A+Scalar%29+-%3E+Self+%3D%0A++++++if+n+%3D%3D+0%0A++++++++then%0A++++++++++PaperSize+%7B%0A++++++++++++width%3A+841+mm%2C%0A++++++++++++height%3A+1189+mm%0A++++++++++%7D%0A++++++++else%0A++++++++++PaperSize+%7B%0A++++++++++++width%3A+floor_in%28mm%2C+PaperSize%3A%3Aa%28n+-+1%29.height+%2F+2%29%2C%0A++++++++++++height%3A+PaperSize%3A%3Aa%28n+-+1%29.width%2C%0A++++++++++%7D%0A%0A++++fn+area%28self%29+-%3E+Area+%3D%0A++++++++self.width+%2A+self.height%0A%0A++++fn+as_string%28self%29+-%3E+String+%3D%0A++++++++%22%7Bself.width%3A%3E4%7D+%C3%97+%7Bself.height%3A%3E5%7D+++%7Bself.area%28%29+-%3E+cm%C2%B2%3A%3E6.1f%7D%22%0A%7D%0A%0A%0Afn+row%28n%29+%3D+%22A%7Bn%3A%3C3%7D+++%7Bsize.as_string%28%29%7D%22%0A++where+size+%3D+PaperSize%3A%3Aa%28n%29%0A%0Aprint%28%22Name++++Width+++++Height++++++++Area++%22%29%0Aprint%28%22----+++-------+++--------+++----------%22%29%0Aprint%28join%28map%28row%2C+range%280%2C+10%29%29%2C+%22%5Cn%22%29%29%0A){ .md-button .md-button--primary } ```numbat # Compute ISO 216 paper sizes for the A series @@ -12,28 +12,30 @@ struct PaperSize { width: Length, height: Length, -} - -fn paper_size_A(n: Scalar) -> PaperSize = - if n == 0 - then - PaperSize { - width: 841 mm, - height: 1189 mm - } - else - PaperSize { - width: floor_in(mm, paper_size_A(n - 1).height / 2), - height: paper_size_A(n - 1).width, - } - -fn paper_area(size: PaperSize) -> Area = - size.width * size.height + fn a(n: Scalar) -> Self = + if n == 0 + then + PaperSize { + width: 841 mm, + height: 1189 mm + } + else + PaperSize { + width: floor_in(mm, PaperSize::a(n - 1).height / 2), + height: PaperSize::a(n - 1).width, + } + + fn area(self) -> Area = + self.width * self.height + + fn as_string(self) -> String = + "{self.width:>4} × {self.height:>5} {self.area() -> cm²:>6.1f}" +} -fn size_as_string(size: PaperSize) = "{size.width:>4} × {size.height:>5} {paper_area(size) -> cm²:>6.1f}" -fn row(n) = "A{n:<3} {size_as_string(paper_size_A(n))}" +fn row(n) = "A{n:<3} {size.as_string()}" + where size = PaperSize::a(n) print("Name Width Height Area ") print("---- ------- -------- ----------") diff --git a/examples/3d_printing.nbt b/examples/3d_printing.nbt index 86ec661b..b1b0ba69 100644 --- a/examples/3d_printing.nbt +++ b/examples/3d_printing.nbt @@ -2,6 +2,13 @@ struct Material { diameter: Length, density: MassDensity, price: Money / Mass, + + fn print_cost(self, mass: Mass) -> Money = self.price × mass + + fn filament_length(self, mass: Mass) -> Length = volume / cross_section -> meter + where r = self.diameter / 2 + and cross_section: Area = π r² + and volume: Volume = mass / self.density } let PLA = Material { @@ -10,18 +17,11 @@ let PLA = Material { price: 16.99 €/kg, } -fn print_cost(material: Material, mass: Mass) -> Money = material.price × mass - -fn filament_length(material: Material, mass: Mass) -> Length = volume / cross_section -> meter - where r = material.diameter / 2 - and cross_section: Area = π r² - and volume: Volume = mass / material.density - # Print parameters let mass_model = 80 g let material = PLA print("Mass of model: {mass_model}") -print("Filament length: {filament_length(material, mass_model):.2}") -print("Cost of model: {print_cost(material, mass_model):.2}") +print("Filament length: {material.filament_length(mass_model):.2}") +print("Cost of model: {material.print_cost(mass_model):.2}") diff --git a/examples/interactive/tidal_chart.nbt b/examples/interactive/tidal_chart.nbt index 9c69fc08..d41829ed 100644 --- a/examples/interactive/tidal_chart.nbt +++ b/examples/interactive/tidal_chart.nbt @@ -9,10 +9,10 @@ struct Constituent { period: Time, amplitude: Length, phase: Angle, -} -fn height(c: Constituent, t: Time) -> Length = - c.amplitude cos(2π t / c.period + c.phase) + fn height(self, t: Time) -> Length = + self.amplitude cos(2π t / self.period + self.phase) +} # The Gulf of Thailand has predominantly DIURNAL tides, meaning # K1 and O1 dominate over the semidiurnal M2 and S2 components. @@ -48,7 +48,7 @@ let O1 = Constituent { let mean_sea_level = 0.5 m -fn tide_height(t: Time) -> Length = mean_sea_level + height(M2, t) + height(S2, t) + height(K1, t) + height(O1, t) +fn tide_height(t: Time) -> Length = mean_sea_level + M2.height(t) + S2.height(t) + K1.height(t) + O1.height(t) let t_start = 0 days let duration = 30 days diff --git a/examples/numbat_syntax.nbt b/examples/numbat_syntax.nbt index fb020107..965e0be8 100644 --- a/examples/numbat_syntax.nbt +++ b/examples/numbat_syntax.nbt @@ -127,6 +127,9 @@ struct Element { # Define a struct name: String, atomic_number: Scalar, density: MassDensity, + + fn neptunium() -> Self = + Self { name: "Neptunium", atomic_number: 93, density: 20.45 g/cm³ } } let hydrogen = Element { # Instantiate it @@ -137,7 +140,22 @@ let hydrogen = Element { # Instantiate it hydrogen.density # Access the field of a struct +let neptunium = # Instantiate using a constructor method + Element::neptunium() +print(neptunium) + struct Vec2 { # A generic struct with type parameter x: D, y: D, + + # Methods can have their own type generics and inherit the struct's generics + fn new(x: D, y: D) -> Self = Self { x: x, y: y } + fn scale(self, factor: Scalar) -> Self = Self { x: self.x * factor, y: self.y * factor } + fn dot(self, other: Vec2) -> D * E = self.x * other.x + self.y * other.y } + +let v = Vec2::new(3 m, 4 m) # Constructor call +let w = Vec2::new(300 cm, 400 cm) +let scaled_v = v.scale(2) # Returning `Self` creates a new value (no in-place mutation) + +assert_eq(v.dot(w) -> m², 25 m²) # Method generic + unit conversion diff --git a/examples/paper_size.nbt b/examples/paper_size.nbt index 21d84411..0da0807f 100644 --- a/examples/paper_size.nbt +++ b/examples/paper_size.nbt @@ -5,30 +5,32 @@ struct PaperSize { width: Length, height: Length, -} -fn paper_size_A(n: Scalar) -> PaperSize = - if n == 0 - then - PaperSize { - width: 841 mm, - height: 1189 mm - } - else - PaperSize { - width: floor_in(mm, paper_size_A(n - 1).height / 2), - height: paper_size_A(n - 1).width, - } + fn a(n: Scalar) -> Self = + if n == 0 + then + PaperSize { + width: 841 mm, + height: 1189 mm + } + else + PaperSize { + width: floor_in(mm, PaperSize::a(n - 1).height / 2), + height: PaperSize::a(n - 1).width, + } -assert_eq(paper_size_A(4).width, 210 mm) -assert_eq(paper_size_A(4).height, 297 mm) + fn area(self) -> Area = + self.width * self.height -fn paper_area(size: PaperSize) -> Area = - size.width * size.height + fn as_string(self) -> String = + "{self.width:>4} × {self.height:>5} {self.area() -> cm²:>6.1f}" +} +assert_eq(PaperSize::a(4).width, 210 mm) +assert_eq(PaperSize::a(4).height, 297 mm) -fn size_as_string(size: PaperSize) = "{size.width:>4} × {size.height:>5} {paper_area(size) -> cm²:>6.1f}" -fn row(n) = "A{n:<3} {size_as_string(paper_size_A(n))}" +fn row(n) = "A{n:<3} {size.as_string()}" + where size = PaperSize::a(n) print("Name Width Height Area ") print("---- ------- -------- ----------") diff --git a/numbat/src/ast.rs b/numbat/src/ast.rs index 159b5a13..2f17d069 100644 --- a/numbat/src/ast.rs +++ b/numbat/src/ast.rs @@ -122,6 +122,13 @@ pub enum Expression<'a> { expr: Box>, field_name: &'a str, }, + MethodCall { + receiver: Box>, + method_name_span: Span, + method_name: &'a str, + args: Vec>, + full_span: Span, + }, List(Span, Vec>), } @@ -156,6 +163,7 @@ impl Expression<'_> { Expression::String(span, _) => *span, Expression::InstantiateStruct { full_span, .. } => *full_span, Expression::AccessField { full_span, .. } => *full_span, + Expression::MethodCall { full_span, .. } => *full_span, Expression::List(span, _) => *span, Expression::TypedHole(span) => *span, } @@ -496,6 +504,7 @@ pub enum Statement<'a> { struct_name: &'a str, type_parameters: Vec<(Span, &'a str, Option)>, fields: Vec<(Span, &'a str, TypeAnnotation)>, + methods: Vec>, }, } @@ -554,12 +563,16 @@ impl Statement<'_> { Statement::DefineStruct { struct_name_span, fields, + methods, .. } => { let mut span = *struct_name_span; if let Some((last_span, _, annotation)) = fields.last() { span = span.extend(last_span).extend(&annotation.full_span()); } + if let Some(method) = methods.last() { + span = span.extend(&method.full_span()); + } span } } @@ -716,6 +729,18 @@ impl ReplaceSpans for Expression<'_> { expr: Box::new(expr.replace_spans()), field_name, }, + Expression::MethodCall { + receiver, + method_name, + args, + .. + } => Expression::MethodCall { + receiver: Box::new(receiver.replace_spans()), + method_name_span: Span::dummy(), + method_name, + args: args.iter().map(|a| a.replace_spans()).collect(), + full_span: Span::dummy(), + }, Expression::List(_, elements) => Expression::List( Span::dummy(), elements.iter().map(|e| e.replace_spans()).collect(), @@ -819,6 +844,7 @@ impl ReplaceSpans for Statement<'_> { struct_name, type_parameters, fields, + methods, .. } => Statement::DefineStruct { struct_name_span: Span::dummy(), @@ -831,6 +857,7 @@ impl ReplaceSpans for Statement<'_> { .iter() .map(|(_span, name, type_)| (Span::dummy(), *name, type_.replace_spans())) .collect(), + methods: methods.iter().map(|m| m.replace_spans()).collect(), }, } } diff --git a/numbat/src/bytecode_interpreter.rs b/numbat/src/bytecode_interpreter.rs index 9ce3d85e..01ee9eae 100644 --- a/numbat/src/bytecode_interpreter.rs +++ b/numbat/src/bytecode_interpreter.rs @@ -20,7 +20,7 @@ use crate::typed_ast::{ use crate::unit::{CanonicalName, Unit}; use crate::unit_registry::{UnitMetadata, UnitRegistry}; use crate::value::{FunctionReference, Value}; -use crate::vm::{Constant, ExecutionContext, FfiCallArg, FfiCallArgs, Op, Vm}; +use crate::vm::{Constant, ExecutionContext, FfiCallArg, FfiCallArgs, MethodCallable, Op, Vm}; use crate::{Type, decorator}; #[derive(Debug, Clone, Default)] @@ -393,6 +393,63 @@ impl BytecodeInterpreter { self.vm.add_op1(Op::BuildList, elements.len() as u16, *span); } + Expression::MethodCall { + full_span, + receiver, + method_ref, + args, + type_scheme, + .. + } => { + let arg_count: u16 = if method_ref.kind == typed_ast::StructMethodKind::Constructor + { + for arg in args { + self.compile_expression(arg); + } + args.len() as u16 + } else { + self.compile_expression(receiver); + for arg in args { + self.compile_expression(arg); + } + (args.len() + 1) as u16 + }; + + let method_callable = self + .vm + .get_method_callable(&method_ref.owner, method_ref.name) + .expect("method must be registered before call sites are compiled"); + + match method_callable { + MethodCallable::Normal(idx) => { + self.vm.add_op2(Op::Call, idx, arg_count, *full_span); + } + MethodCallable::Foreign(idx) => { + let mut ffi_args = Vec::with_capacity(args.len() + 1); + if method_ref.kind == typed_ast::StructMethodKind::Instance { + ffi_args.push(FfiCallArg { + span: receiver.full_span(), + type_: receiver.get_type_scheme(), + }); + } + ffi_args.extend(args.iter().map(|a| FfiCallArg { + span: a.full_span(), + type_: a.get_type_scheme(), + })); + let call_args_idx = self.vm.add_ffi_call_args(FfiCallArgs { + args: ffi_args, + return_type: Some(type_scheme.clone()), + }); + self.vm.add_op3( + Op::FFICallFunction, + idx, + arg_count, + call_args_idx, + *full_span, + ); + } + } + } Expression::TypedHole(_, _) => { unreachable!("Typed holes cause type inference errors") } @@ -453,8 +510,9 @@ impl BytecodeInterpreter { metadata: LocalMetadata::default(), }); } - for local_variables in local_variables { - self.compile_define_variable(local_variables); + + for local_variable in local_variables { + self.compile_define_variable(local_variable); } self.compile_expression(expr); @@ -475,12 +533,72 @@ impl BytecodeInterpreter { } => { // Declaring a foreign function does not generate any bytecode. But we register // its name and arity here to be able to distinguish it from normal functions. - self.vm .add_foreign_function(name, parameters.len()..=parameters.len()); self.functions.insert(name.to_compact_string(), true); } + Statement::DefineMethod { + struct_name, + method_name, + parameters, + body: Some(expr), + local_variables, + .. + } => { + if let Some(idx) = self + .vm + .get_empty_method_function_idx(struct_name, method_name) + { + self.vm.begin_reserved_function(idx); + } else { + let runtime_name = format!(""); + let idx = self.vm.begin_function(&runtime_name); + self.vm + .register_method_function(struct_name, method_name, idx); + } + + self.locals.push(vec![]); + + let current_depth = self.current_depth(); + for parameter in parameters { + self.locals[current_depth].push(Local { + identifiers: [parameter.1.to_compact_string()].into(), + metadata: LocalMetadata::default(), + }); + } + + for local_variable in local_variables { + self.compile_define_variable(local_variable); + } + + self.compile_expression(expr); + + self.vm.add_op(Op::Return, expr.full_span()); + + self.locals.pop(); + + self.vm.end_function(); + } + Statement::DefineMethod { + struct_name, + method_name, + parameters, + body: None, + .. + } => { + // Declaring a foreign function does not generate any bytecode. But we register + // its name and arity here to be able to distinguish it from normal functions. + self.vm + .add_foreign_function(method_name, parameters.len()..=parameters.len()); + + let ffi_idx = self + .vm + .get_ffi_callable_idx(method_name) + .expect("just registered foreign method"); + self.vm + .register_foreign_method(struct_name, method_name, ffi_idx); + } Statement::DefineDimension(_name, _dexprs) => { // Declaring a dimension is like introducing a new type. The information // is only relevant for the type checker. Nothing happens at run time. @@ -646,6 +764,26 @@ impl BytecodeInterpreter { Ok(()) } + fn preregister_struct_methods<'a>(&mut self, methods: &[Statement<'a>]) { + for statement in methods { + let Statement::DefineMethod { + struct_name, + method_name, + body: Some(_), + .. + } = statement + else { + continue; + }; + + let runtime_name = format!(""); + let idx = self.vm.begin_function(&runtime_name); + self.vm + .register_method_function(struct_name, method_name, idx); + self.vm.end_function(); + } + } + fn run( &mut self, settings: &mut InterpreterSettings, @@ -720,8 +858,34 @@ impl Interpreter for BytecodeInterpreter { prefix_transformer: &crate::prefix_transformer::Transformer, typechecker: &TypeChecker, ) -> Result { - for statement in statements { - self.compile_statement(statement, typechecker)?; + let mut idx = 0; + while idx < statements.len() { + if let Statement::DefineStruct(struct_info) = &statements[idx] { + self.compile_statement(&statements[idx], typechecker)?; + + let mut end = idx + 1; + while end < statements.len() { + match &statements[end] { + Statement::DefineMethod { + struct_name, + method_name: _, + .. + } if struct_name == &struct_info.name => { + end += 1; + } + _ => break, + } + } + + self.preregister_struct_methods(&statements[idx + 1..end]); + for statement in &statements[idx + 1..end] { + self.compile_statement(statement, typechecker)?; + } + idx = end; + } else { + self.compile_statement(&statements[idx], typechecker)?; + idx += 1; + } } self.run(settings, prefix_transformer, typechecker) diff --git a/numbat/src/diagnostic.rs b/numbat/src/diagnostic.rs index fa04a4c7..2b58d5ed 100644 --- a/numbat/src/diagnostic.rs +++ b/numbat/src/diagnostic.rs @@ -432,6 +432,16 @@ impl ErrorDiagnostic for TypeCheckError { .diagnostic_label(LabelStyle::Secondary) .with_message(type_.to_string()), ]), + TypeCheckError::MethodCallOnNonStructType(span, _, _) + | TypeCheckError::MethodNotFound(span, _, _) + | TypeCheckError::ConstructorCalledAsMethod(span, _, _) + | TypeCheckError::InstanceMethodCalledAsConstructor(span, _, _) + | TypeCheckError::InvalidSelfParameterType(span, _, _) + | TypeCheckError::SelfTypeOutsideStructMethod(span) + | TypeCheckError::InvalidStructMember(span) => d.with_labels(vec![ + span.diagnostic_label(LabelStyle::Primary) + .with_message(inner_error), + ]), TypeCheckError::IncompatibleTypesForStructField( expected_field_span, _expected_type, @@ -466,6 +476,18 @@ impl ErrorDiagnostic for TypeCheckError { .diagnostic_label(LabelStyle::Secondary) .with_message("Already defined here"), ]), + TypeCheckError::DuplicateMemberInStructDefinition( + this_field_span, + that_field_span, + _attr_name, + ) => d.with_labels(vec![ + this_field_span + .diagnostic_label(LabelStyle::Primary) + .with_message(inner_error), + that_field_span + .diagnostic_label(LabelStyle::Secondary) + .with_message("Already defined here"), + ]), TypeCheckError::MissingFieldsInStructInstantiation( construction_span, defn_span, diff --git a/numbat/src/ffi/lookup.rs b/numbat/src/ffi/lookup.rs index 5fa0c2fe..ebda8dd7 100644 --- a/numbat/src/ffi/lookup.rs +++ b/numbat/src/ffi/lookup.rs @@ -95,10 +95,11 @@ pub fn _get_chemical_element_data_raw( kind: StructKind::Instance(vec![]), definition_span: unknown_span, fields, + methods: IndexMap::new(), }; Ok(Value::StructInstance( Arc::new(info), - vec![ + Arc::from(vec![ Value::String(element.symbol().into()), Value::String(element.name().into()), Value::Quantity(Quantity::from_scalar(element.atomic_number() as f64)), @@ -151,7 +152,7 @@ pub fn _get_chemical_element_data_raw( .map(|KiloJoulePerMole(e)| e) .unwrap_or(f64::NAN), )), - ], + ]), )) } else { Err(Box::new(RuntimeErrorKind::ChemicalElementNotFound( diff --git a/numbat/src/ffi/plot.rs b/numbat/src/ffi/plot.rs index 58a7a033..0616acc9 100644 --- a/numbat/src/ffi/plot.rs +++ b/numbat/src/ffi/plot.rs @@ -19,11 +19,12 @@ use crate::value::Value; #[cfg(feature = "plotting")] fn line_plot(mut args: Args) -> Result> { - let mut fields = arg!(args).unsafe_as_struct_fields(); - let ys = fields.pop().unwrap(); - let xs = fields.pop().unwrap(); - let y_label = fields.pop().unwrap().unsafe_as_string(); - let x_label = fields.pop().unwrap().unsafe_as_string(); + let plot = arg!(args); + let fields = plot.unsafe_as_struct_fields(); + let x_label = fields[0].clone().unsafe_as_string(); + let y_label = fields[1].clone().unsafe_as_string(); + let xs = fields[2].clone(); + let ys = fields[3].clone(); let xs = xs.unsafe_as_list(); let ys = ys.unsafe_as_list(); @@ -70,10 +71,11 @@ fn line_plot(mut args: Args) -> Result> { #[cfg(feature = "plotting")] fn bar_chart(mut args: Args) -> Result> { - let mut fields = arg!(args).unsafe_as_struct_fields(); - let x_labels = fields.pop().unwrap(); - let values = fields.pop().unwrap(); - let value_label = fields.pop().unwrap().unsafe_as_string(); + let plot = arg!(args); + let fields = plot.unsafe_as_struct_fields(); + let value_label = fields[0].clone().unsafe_as_string(); + let values = fields[1].clone(); + let x_labels = fields[2].clone(); let x_labels = x_labels .unsafe_as_list() diff --git a/numbat/src/parser.rs b/numbat/src/parser.rs index 7c1c7e47..32b6db09 100644 --- a/numbat/src/parser.rs +++ b/numbat/src/parser.rs @@ -41,7 +41,7 @@ //! power ::= factorial ( "^" "-" ? power ) ? //! factorial ::= unicode_power "!" * //! unicode_power ::= call ( "⁻" ? ( "¹" | "²" | "³" | "⁴" | "⁵" | "⁶" | "⁷" | "⁸" | "⁹" ) ) ? -//! call ::= primary ( ( "(" arguments? ")" ) | "." identifier ) * +//! call ::= primary ( ( "(" arguments? ")" ) | "." identifier | "::" identifier "(" arguments? ")" ) * //! arguments ::= expression ( "," expression ) * //! primary ::= boolean | string | hex_number | oct_number | bin_number | number | identifier ( struct_expr ? ) | typed_hole | list_expr | "(" expression ")" //! struct_expr ::= "{" ( identifier ":" type_annotation "," )* ( identifier ":" expression "," ? ) ? "}" @@ -235,6 +235,15 @@ pub enum ParseErrorKind { #[error("Expected '{{' after struct name")] ExpectedLeftCurlyAfterStructName, + #[error("Expected field or method definition in struct body")] + ExpectedFieldOrMethodInStruct, + + #[error("Fields must be declared before methods in struct body")] + FieldAfterMethodInStruct, + + #[error("Expected '(' after method name in '::' call")] + ExpectedLeftParenAfterStaticMethodName, + #[error("Expected ',' or ']' in list expression")] ExpectedCommaOrRightBracketInList, @@ -966,16 +975,32 @@ impl<'a> Parser<'a> { self.skip_empty_lines(tokens); let mut fields = vec![]; - while self.match_exact(tokens, TokenKind::RightCurly).is_none() { + let mut methods = vec![]; + let mut parsing_methods = false; + while self.peek(tokens).kind != TokenKind::RightCurly { self.skip_empty_lines(tokens); + if let Some(fn_token) = self.match_exact(tokens, TokenKind::Fn) { + parsing_methods = true; + methods.push(self.parse_function_declaration(tokens, fn_token.span)?); + self.skip_empty_lines(tokens); + continue; + } + let Some(field_name) = self.match_exact(tokens, TokenKind::Identifier) else { return Err(ParseError { - kind: ParseErrorKind::ExpectedFieldNameInStruct, + kind: ParseErrorKind::ExpectedFieldOrMethodInStruct, span: self.peek(tokens).span, }); }; + if parsing_methods { + return Err(ParseError { + kind: ParseErrorKind::FieldAfterMethodInStruct, + span: field_name.span, + }); + } + self.skip_empty_lines(tokens); if self.match_exact(tokens, TokenKind::Colon).is_none() { @@ -994,7 +1019,10 @@ impl<'a> Parser<'a> { self.skip_empty_lines(tokens); - if !has_comma && self.peek(tokens).kind != TokenKind::RightCurly { + if !has_comma + && self.peek(tokens).kind != TokenKind::RightCurly + && self.peek(tokens).kind != TokenKind::Fn + { return Err(ParseError { kind: ParseErrorKind::ExpectedCommaOrRightCurlyInStructFieldList, span: self.peek(tokens).span, @@ -1003,12 +1031,14 @@ impl<'a> Parser<'a> { fields.push((field_name.span, field_name.lexeme, attr_type)); } + self.match_exact(tokens, TokenKind::RightCurly); Ok(Statement::DefineStruct { struct_name_span: name_span, struct_name: name, type_parameters, fields, + methods, }) } @@ -1419,16 +1449,52 @@ impl<'a> Parser<'a> { callable: Box::new(expr), args, }; - } else if self.match_exact(tokens, TokenKind::Period).is_some() { + } else if self.match_exact(tokens, TokenKind::DoubleColon).is_some() { let ident = self.identifier(tokens)?; let ident_span = self.last(tokens).unwrap().span; - let full_span = expr.full_span().extend(&ident_span); - expr = Expression::AccessField { + if self.match_exact(tokens, TokenKind::LeftParen).is_none() { + return Err(ParseError { + kind: ParseErrorKind::ExpectedLeftParenAfterStaticMethodName, + span: self.peek(tokens).span, + }); + } + + let args = self.arguments(tokens)?; + let full_span = expr.full_span().extend(&self.last(tokens).unwrap().span); + + expr = Expression::MethodCall { + receiver: Box::new(expr), + method_name_span: ident_span, + method_name: ident, + args, full_span, - ident_span, - expr: Box::new(expr), - field_name: ident, + }; + } else if self.match_exact(tokens, TokenKind::Period).is_some() { + let ident = self.identifier(tokens)?; + let ident_span = self.last(tokens).unwrap().span; + + // Check if this is a method call (has parentheses) or field access (no parentheses) + if self.match_exact(tokens, TokenKind::LeftParen).is_some() { + let args = self.arguments(tokens)?; + let full_span = expr.full_span().extend(&self.last(tokens).unwrap().span); + + expr = Expression::MethodCall { + receiver: Box::new(expr), + method_name_span: ident_span, + method_name: ident, + args, + full_span, + }; + } else { + let full_span = expr.full_span().extend(&ident_span); + + expr = Expression::AccessField { + full_span, + ident_span, + expr: Box::new(expr), + field_name: ident, + }; } } else { return Ok(expr); @@ -3250,6 +3316,147 @@ mod tests { ); } + #[test] + fn struct_methods_and_method_calls() { + parse_as( + &["struct Point { x: Length, fn magnitude(self) = self.x }"], + Statement::DefineStruct { + struct_name_span: Span::dummy(), + struct_name: "Point", + type_parameters: vec![], + fields: vec![( + Span::dummy(), + "x", + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + Span::dummy(), + "Length".into(), + vec![], + )), + )], + methods: vec![Statement::DefineFunction { + fn_keyword_span: Span::dummy(), + function_name_span: Span::dummy(), + function_name: "magnitude", + type_parameters: vec![], + parameters: vec![(Span::dummy(), "self", None)], + body: Some(Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("self")), + field_name: "x", + }), + local_variables: vec![], + return_type_annotation: None, + decorators: vec![], + }], + }, + ); + + parse_as( + &[ + "struct Box { inner: T, fn replace(self, value: U) -> Box = Box { inner: value } }", + ], + Statement::DefineStruct { + struct_name_span: Span::dummy(), + struct_name: "Box", + type_parameters: vec![(Span::dummy(), "T", None)], + fields: vec![( + Span::dummy(), + "inner", + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + Span::dummy(), + "T".into(), + vec![], + )), + )], + methods: vec![Statement::DefineFunction { + fn_keyword_span: Span::dummy(), + function_name_span: Span::dummy(), + function_name: "replace", + type_parameters: vec![(Span::dummy(), "U", None)], + parameters: vec![ + (Span::dummy(), "self", None), + ( + Span::dummy(), + "value", + Some(TypeAnnotation::TypeExpression( + TypeExpression::TypeIdentifier(Span::dummy(), "U".into(), vec![]), + )), + ), + ], + body: Some(struct_! { + Box, + inner: identifier!("value") + }), + local_variables: vec![], + return_type_annotation: Some(TypeAnnotation::TypeExpression( + TypeExpression::TypeIdentifier( + Span::dummy(), + "Box".into(), + vec![TypeAnnotation::TypeExpression( + TypeExpression::TypeIdentifier(Span::dummy(), "U".into(), vec![]), + )], + ), + )), + decorators: vec![], + }], + }, + ); + + parse_as_expression( + &["p.translate(1, 2).magnitude()"], + Expression::MethodCall { + receiver: Box::new(Expression::MethodCall { + receiver: Box::new(identifier!("p")), + method_name_span: Span::dummy(), + method_name: "translate", + args: vec![scalar!(1.0), scalar!(2.0)], + full_span: Span::dummy(), + }), + method_name_span: Span::dummy(), + method_name: "magnitude", + args: vec![], + full_span: Span::dummy(), + }, + ); + + parse_as_expression( + &["Point::new(1, 2)"], + Expression::MethodCall { + receiver: Box::new(identifier!("Point")), + method_name_span: Span::dummy(), + method_name: "new", + args: vec![scalar!(1.0), scalar!(2.0)], + full_span: Span::dummy(), + }, + ); + + parse_as_expression( + &["Self { x: 1 }"], + Expression::InstantiateStruct { + full_span: Span::dummy(), + ident_span: Span::dummy(), + name: "Self", + fields: vec![(Span::dummy(), "x", scalar!(1.0))], + }, + ); + + should_fail_with( + &["Point::new"], + ParseErrorKind::ExpectedLeftParenAfterStaticMethodName, + ); + + should_fail_with( + &["struct Point { fn magnitude(self) = self.x, x: Length }"], + ParseErrorKind::ExpectedFieldOrMethodInStruct, + ); + + should_fail_with( + &["struct Point { fn magnitude(self) = self.x x: Length }"], + ParseErrorKind::ExpectedFieldOrMethodInStruct, + ); + } + #[test] fn postfix_apply() { parse_as_expression( @@ -3651,6 +3858,7 @@ mod tests { )), ), ], + methods: vec![], }, ); @@ -3684,6 +3892,7 @@ mod tests { ), (Span::dummy(), "name", TypeAnnotation::String(Span::dummy())), ], + methods: vec![], }, ); diff --git a/numbat/src/prefix_transformer.rs b/numbat/src/prefix_transformer.rs index 6a08d799..43089b8d 100644 --- a/numbat/src/prefix_transformer.rs +++ b/numbat/src/prefix_transformer.rs @@ -156,6 +156,12 @@ impl Transformer { Expression::AccessField { expr, .. } => { self.transform_expression(expr); } + Expression::MethodCall { receiver, args, .. } => { + self.transform_expression(receiver); + for arg in args { + self.transform_expression(arg); + } + } Expression::List(_, elements) => { for e in elements { self.transform_expression(e); @@ -231,7 +237,46 @@ impl Transformer { fn transform_statement(&mut self, statement: &mut Statement) -> Result<()> { match statement { - Statement::DefineStruct { .. } | Statement::ModuleImport(_, _) => {} + Statement::DefineStruct { methods, .. } => { + for method in methods { + let Statement::DefineFunction { + parameters, + body, + local_variables, + .. + } = method + else { + continue; + }; + + // Struct methods live in the struct namespace, not the global namespace. + // Do not register their names as global identifiers. + let mut method_body_transformer = self.clone(); + for (param_span, param, _) in &*parameters { + method_body_transformer + .prefix_parser + .add_shadowing_identifier(param, *param_span)?; + } + + for def in &mut *local_variables { + method_body_transformer + .variable_names + .push(def.identifier.to_compact_string()); + method_body_transformer + .prefix_parser + .add_shadowing_identifier(def.identifier, def.identifier_span)?; + } + + if let Some(expr) = body { + method_body_transformer.transform_expression(expr); + } + + for def in local_variables { + method_body_transformer.transform_expression(&mut def.expr); + } + } + } + Statement::ModuleImport(_, _) => {} Statement::Expression(expr) => { self.transform_expression(expr); diff --git a/numbat/src/traversal.rs b/numbat/src/traversal.rs index d874aa6c..cc0e943c 100644 --- a/numbat/src/traversal.rs +++ b/numbat/src/traversal.rs @@ -105,6 +105,18 @@ impl ForAllTypeSchemes for Expression<'_> { } f(type_scheme); } + Expression::MethodCall { + receiver, + args, + type_scheme, + .. + } => { + receiver.for_all_type_schemes(f); + for arg in args { + arg.for_all_type_schemes(f); + } + f(type_scheme); + } Expression::TypedHole(_, type_) => { f(type_); } @@ -127,6 +139,12 @@ impl ForAllTypeSchemes for Statement<'_> { local_variables, fn_type, .. + } + | Statement::DefineMethod { + body, + local_variables, + fn_type, + .. } => { for local_variable in local_variables { local_variable.expr.for_all_type_schemes(f); @@ -170,6 +188,11 @@ impl ForAllExpressions for Statement<'_> { body, local_variables, .. + } + | Statement::DefineMethod { + body, + local_variables, + .. } => { for local_variable in local_variables { local_variable.expr.for_all_expressions(f); @@ -243,6 +266,12 @@ impl ForAllExpressions for Expression<'_> { element.for_all_expressions(f); } } + Expression::MethodCall { receiver, args, .. } => { + receiver.for_all_expressions(f); + for arg in args { + arg.for_all_expressions(f); + } + } Expression::TypedHole(_, _) => {} } } diff --git a/numbat/src/typechecker/const_evaluation.rs b/numbat/src/typechecker/const_evaluation.rs index c19546f1..cd5a93e2 100644 --- a/numbat/src/typechecker/const_evaluation.rs +++ b/numbat/src/typechecker/const_evaluation.rs @@ -108,6 +108,7 @@ pub fn evaluate_const_expr(expr: &typed_ast::Expression) -> Result { typed_ast::Expression::InstantiateStruct { .. } => "instantiate struct", typed_ast::Expression::AccessField { .. } => "access field of struct", typed_ast::Expression::List { .. } => "lists", + typed_ast::Expression::MethodCall { .. } => "method call", typed_ast::Expression::TypedHole(_, _) => "typed hole", }; diff --git a/numbat/src/typechecker/error.rs b/numbat/src/typechecker/error.rs index 2efd101a..2984bb6b 100644 --- a/numbat/src/typechecker/error.rs +++ b/numbat/src/typechecker/error.rs @@ -144,6 +144,9 @@ pub enum TypeCheckError { #[error("Duplicate field '{2}' in struct definition")] DuplicateFieldInStructDefinition(Span, Span, String), + #[error("Duplicate member '{2}' in struct definition")] + DuplicateMemberInStructDefinition(Span, Span, String), + #[error("Duplicate field '{2}' in struct instantiation")] DuplicateFieldInStructInstantiation(Span, Span, String), @@ -153,6 +156,27 @@ pub enum TypeCheckError { #[error("Field '{2}' does not exist in struct '{3}'")] UnknownFieldAccess(Span, Span, String, Type), + #[error("Can not call method '{1}' on non struct type '{2}'")] + MethodCallOnNonStructType(Span, String, Type), + + #[error("Method '{1}' does not exist on struct '{2}'")] + MethodNotFound(Span, String, String), + + #[error("Constructor '{1}' of struct '{2}' can not be called as an instance method")] + ConstructorCalledAsMethod(Span, String, String), + + #[error("Instance method '{1}' of struct '{2}' can not be called as a constructor")] + InstanceMethodCalledAsConstructor(Span, String, String), + + #[error("Type of 'self' parameter in method '{1}' must be '{2}'")] + InvalidSelfParameterType(Span, String, String), + + #[error("`Self` can only be used inside struct method definitions")] + SelfTypeOutsideStructMethod(Span), + + #[error("Only function definitions are allowed in struct method section")] + InvalidStructMember(Span), + #[error("Missing fields in struct instantiation")] MissingFieldsInStructInstantiation(Span, Span, Vec<(CompactString, Type)>), diff --git a/numbat/src/typechecker/mod.rs b/numbat/src/typechecker/mod.rs index 143bd1ba..031a7057 100644 --- a/numbat/src/typechecker/mod.rs +++ b/numbat/src/typechecker/mod.rs @@ -28,6 +28,7 @@ use crate::pretty_print::PrettyPrint; use crate::span::Span; use crate::type_variable::TypeVariable; use crate::typed_ast::{self, DType, DTypeFactor, Expression, StructInfo, StructKind, Type}; +use crate::typed_ast::{StructMethodInfo, StructMethodKind}; use crate::{decorator, ffi, suggestion}; use compact_str::{CompactString, ToCompactString, format_compact}; @@ -190,6 +191,17 @@ pub struct TypeChecker { name_generator: NameGenerator, constraints: ConstraintSet, + + methods: HashMap>, + checking_struct_method: bool, + current_struct_name: Option, + non_generic_struct_instances: HashMap, +} + +/// Stores information about a method defined in a struct method section +#[derive(Clone)] +struct MethodInfo { + signature: FunctionSignature, } struct ElaborationDefinitionArgs<'a, 'b> { @@ -205,6 +217,141 @@ struct ElaborationDefinitionArgs<'a, 'b> { } impl TypeChecker { + fn rewrite_self_in_type_expression( + type_expression: &mut TypeExpression, + struct_name: &str, + inside_struct_method: bool, + self_type_args: &[TypeAnnotation], + ) -> Result<()> { + match type_expression { + TypeExpression::Unity(_) => Ok(()), + TypeExpression::TypeIdentifier(span, name, type_args) => { + if name == "Self" { + if !inside_struct_method { + return Err(Box::new(TypeCheckError::SelfTypeOutsideStructMethod(*span))); + } + *name = struct_name.to_compact_string(); + *type_args = self_type_args.to_vec(); + } + + for type_arg in type_args { + Self::rewrite_self_in_type_annotation( + type_arg, + struct_name, + inside_struct_method, + self_type_args, + )?; + } + Ok(()) + } + TypeExpression::Multiply(_, lhs, rhs) | TypeExpression::Divide(_, lhs, rhs) => { + Self::rewrite_self_in_type_expression( + lhs, + struct_name, + inside_struct_method, + self_type_args, + )?; + Self::rewrite_self_in_type_expression( + rhs, + struct_name, + inside_struct_method, + self_type_args, + ) + } + TypeExpression::Power(_, lhs, _, _) => Self::rewrite_self_in_type_expression( + lhs, + struct_name, + inside_struct_method, + self_type_args, + ), + } + } + + fn rewrite_self_in_type_annotation( + annotation: &mut TypeAnnotation, + struct_name: &str, + inside_struct_method: bool, + self_type_args: &[TypeAnnotation], + ) -> Result<()> { + match annotation { + TypeAnnotation::TypeExpression(type_expression) => { + Self::rewrite_self_in_type_expression( + type_expression, + struct_name, + inside_struct_method, + self_type_args, + ) + } + TypeAnnotation::Bool(_) | TypeAnnotation::String(_) | TypeAnnotation::DateTime(_) => { + Ok(()) + } + TypeAnnotation::Fn(_, parameters, return_type) => { + for parameter in parameters { + Self::rewrite_self_in_type_annotation( + parameter, + struct_name, + inside_struct_method, + self_type_args, + )?; + } + Self::rewrite_self_in_type_annotation( + return_type, + struct_name, + inside_struct_method, + self_type_args, + ) + } + TypeAnnotation::List(_, element_type) => Self::rewrite_self_in_type_annotation( + element_type, + struct_name, + inside_struct_method, + self_type_args, + ), + } + } + + fn rewrite_self_in_method_signature( + method: &mut ast::Statement<'_>, + struct_name: &str, + self_type_args: &[TypeAnnotation], + ) -> Result<()> { + let ast::Statement::DefineFunction { + parameters, + return_type_annotation, + local_variables, + .. + } = method + else { + return Ok(()); + }; + + for (_, _, parameter_annotation) in parameters { + if let Some(annotation) = parameter_annotation { + Self::rewrite_self_in_type_annotation( + annotation, + struct_name, + true, + self_type_args, + )?; + } + } + if let Some(annotation) = return_type_annotation { + Self::rewrite_self_in_type_annotation(annotation, struct_name, true, self_type_args)?; + } + for local_variable in local_variables { + if let Some(annotation) = &mut local_variable.type_annotation { + Self::rewrite_self_in_type_annotation( + annotation, + struct_name, + true, + self_type_args, + )?; + } + } + + Ok(()) + } + fn fresh_type_variable(&mut self) -> Type { Type::TVar(self.name_generator.fresh_type_variable()) } @@ -217,6 +364,90 @@ impl TypeChecker { self.constraints.add_dtype_constraint(type_) } + /// Register a method for a struct + fn register_method( + &mut self, + struct_name: CompactString, + method_name: CompactString, + signature: FunctionSignature, + ) { + self.methods + .entry(struct_name) + .or_default() + .insert(method_name, MethodInfo { signature }); + } + + /// Look up a method for a struct + fn lookup_method(&self, struct_name: &str, method_name: &str) -> Option<&MethodInfo> { + self.methods.get(struct_name)?.get(method_name) + } + + fn provisional_method_signature( + &mut self, + method: &ast::Statement<'_>, + ) -> Result { + let ast::Statement::DefineFunction { + function_name_span, + function_name, + type_parameters, + parameters, + return_type_annotation, + .. + } = method + else { + unreachable!("method members are checked to be functions"); + }; + + let initial_type_parameter_count = self.registry.introduced_type_parameters.len(); + for (span, type_parameter, bound) in type_parameters { + self.registry.introduced_type_parameters.push(( + *span, + (*type_parameter).to_compact_string(), + bound.clone(), + )); + } + + let parameter_types = parameters + .iter() + .map(|(_, _, annotation)| { + annotation + .as_ref() + .map(|a| self.type_from_annotation(a)) + .transpose() + .map(|a| a.unwrap_or_else(|| self.fresh_type_variable())) + }) + .collect::>>(); + + let return_type = return_type_annotation + .as_ref() + .map(|annotation| self.type_from_annotation(annotation)) + .transpose(); + + self.registry + .introduced_type_parameters + .truncate(initial_type_parameter_count); + + let parameter_types = parameter_types?; + let return_type = return_type?.unwrap_or_else(|| self.fresh_type_variable()); + + Ok(FunctionSignature { + name: (*function_name).to_compact_string(), + definition_span: *function_name_span, + type_parameters: type_parameters + .iter() + .map(|(span, name, bound)| (*span, (*name).to_compact_string(), bound.clone())) + .collect(), + parameters: parameters + .iter() + .map(|(span, name, annotation)| { + (*span, (*name).to_compact_string(), annotation.clone()) + }) + .collect(), + return_type_annotation: return_type_annotation.clone(), + fn_type: TypeScheme::Concrete(Type::Fn(parameter_types, Box::new(return_type))), + }) + } + fn enforce_dtype(&mut self, type_: &Type, span: Span) -> Result<()> { if self .constraints @@ -233,6 +464,9 @@ impl TypeChecker { } fn type_from_annotation(&self, annotation: &TypeAnnotation) -> Result { + let mut checked_annotation = annotation.clone(); + Self::rewrite_self_in_type_annotation(&mut checked_annotation, "", false, &[])?; + match annotation { TypeAnnotation::TypeExpression(dexpr) => { if let TypeExpression::TypeIdentifier(span, name, type_args) = dexpr @@ -263,6 +497,7 @@ impl TypeChecker { name: struct_info.name.clone(), kind: StructKind::Instance(vec![]), fields: struct_info.fields.clone(), + methods: struct_info.methods.clone(), }))); } @@ -286,6 +521,7 @@ impl TypeChecker { name: struct_info.name.clone(), kind: StructKind::Instance(concrete_type_args), fields: instantiated_fields, + methods: struct_info.methods.clone(), }))); } @@ -1026,27 +1262,44 @@ impl TypeChecker { fields, } => { let name = *name; + let resolved_name: CompactString = if name == "Self" { + self.current_struct_name.clone().ok_or_else(|| { + Box::new(TypeCheckError::SelfTypeOutsideStructMethod(*ident_span)) + })? + } else { + name.to_compact_string() + }; let fields_checked = fields .iter() .map(|(_, n, v)| Ok((*n, self.elaborate_expression(v)?))) .collect::>>()?; - let Some(struct_info) = self.structs.get(name).cloned() else { + let Some(struct_info) = self.structs.get(&resolved_name).cloned() else { return Err(Box::new(TypeCheckError::UnknownStruct( *ident_span, - name.to_owned(), + resolved_name.to_string(), ))); }; // For generic structs, instantiate type parameters with fresh type variables let instantiated_struct_info = match &struct_info.kind { StructKind::Definition(type_parameters) if type_parameters.is_empty() => { - // Non-generic struct: create instance with empty type arguments - StructInfo { - definition_span: struct_info.definition_span, - name: struct_info.name.clone(), - kind: StructKind::Instance(vec![]), - fields: struct_info.fields.clone(), + // Non-generic struct: cache a single instantiated view and clone cheaply. + if let Some(cached) = + self.non_generic_struct_instances.get(&struct_info.name) + { + cached.clone() + } else { + let instantiated = StructInfo { + definition_span: struct_info.definition_span, + name: struct_info.name.clone(), + kind: StructKind::Instance(vec![]), + fields: struct_info.fields.clone(), + methods: struct_info.methods.clone(), + }; + self.non_generic_struct_instances + .insert(struct_info.name.clone(), instantiated.clone()); + instantiated } } StructKind::Definition(type_parameters) => { @@ -1088,6 +1341,7 @@ impl TypeChecker { variables.iter().map(|v| Type::TVar(v.clone())).collect(), ), fields: instantiated_fields, + methods: struct_info.methods.clone(), } } StructKind::Instance(_) => { @@ -1260,6 +1514,174 @@ impl TypeChecker { let type_ = self.fresh_type_variable(); typed_ast::Expression::TypedHole(*span, TypeScheme::concrete(type_)) } + ast::Expression::MethodCall { + receiver, + method_name_span, + method_name, + args, + full_span, + } => { + if let ast::Expression::Identifier(receiver_span, receiver_name) = receiver.as_ref() + && self.structs.contains_key(*receiver_name) + { + let struct_info = self + .structs + .get(*receiver_name) + .cloned() + .expect("receiver name was already checked as known struct"); + let method_meta = struct_info.methods.get(*method_name).ok_or_else(|| { + Box::new(TypeCheckError::MethodNotFound( + *method_name_span, + method_name.to_string(), + (*receiver_name).to_string(), + )) + })?; + if method_meta.kind == StructMethodKind::Instance { + return Err(Box::new(TypeCheckError::InstanceMethodCalledAsConstructor( + *method_name_span, + method_name.to_string(), + (*receiver_name).to_string(), + ))); + } + + let method = self + .lookup_method(receiver_name, method_name) + .ok_or_else(|| { + Box::new(TypeCheckError::MethodNotFound( + *method_name_span, + method_name.to_string(), + (*receiver_name).to_string(), + )) + })? + .clone(); + + let method_arguments = args + .iter() + .map(|e| self.elaborate_expression(e)) + .collect::>>()?; + + let method_argument_types = method_arguments + .iter() + .map(typed_ast::Expression::get_type) + .collect::>(); + + let checked_call = proper_function_call(ProperFunctionCallArgs { + registry: &self.registry, + constraints: &mut self.constraints, + name_generator: &mut self.name_generator, + span: method_name_span, + full_span, + function_name: method_name, + signature: &method.signature, + arguments: method_arguments.clone(), + argument_types: method_argument_types, + })?; + + let typed_ast::Expression::FunctionCall { type_scheme, .. } = checked_call + else { + unreachable!("proper function call returns typed function call"); + }; + + return Ok(typed_ast::Expression::MethodCall { + full_span: *full_span, + receiver: Box::new(typed_ast::Expression::Identifier { + span: *receiver_span, + name: receiver_name, + type_scheme: TypeScheme::concrete(Type::Struct(Box::new( + struct_info.clone(), + ))), + }), + method_ref: typed_ast::MethodRef { + owner: struct_info.name.clone(), + name: method_name, + kind: StructMethodKind::Constructor, + }, + method_name_span: *method_name_span, + args: method_arguments, + type_scheme, + }); + } + + let receiver_checked = self.elaborate_expression(receiver)?; + let receiver_type = receiver_checked.get_type(); + + let Type::Struct(struct_info) = &receiver_type else { + return Err(Box::new(TypeCheckError::MethodCallOnNonStructType( + *method_name_span, + method_name.to_string(), + receiver_type, + ))); + }; + + let method_meta = struct_info.methods.get(*method_name).ok_or_else(|| { + Box::new(TypeCheckError::MethodNotFound( + *method_name_span, + method_name.to_string(), + struct_info.name.to_string(), + )) + })?; + if method_meta.kind == StructMethodKind::Constructor { + return Err(Box::new(TypeCheckError::ConstructorCalledAsMethod( + *method_name_span, + method_name.to_string(), + struct_info.name.to_string(), + ))); + } + + let method = self + .lookup_method(&struct_info.name, method_name) + .ok_or_else(|| { + Box::new(TypeCheckError::MethodNotFound( + *method_name_span, + method_name.to_string(), + struct_info.name.to_string(), + )) + })? + .clone(); + + let arguments_checked = args + .iter() + .map(|e| self.elaborate_expression(e)) + .collect::>>()?; + + let mut method_arguments = Vec::with_capacity(arguments_checked.len() + 1); + method_arguments.push(receiver_checked.clone()); + method_arguments.extend(arguments_checked.clone()); + + let method_argument_types = method_arguments + .iter() + .map(typed_ast::Expression::get_type) + .collect::>(); + + let checked_call = proper_function_call(ProperFunctionCallArgs { + registry: &self.registry, + constraints: &mut self.constraints, + name_generator: &mut self.name_generator, + span: method_name_span, + full_span, + function_name: method_name, + signature: &method.signature, + arguments: method_arguments, + argument_types: method_argument_types, + })?; + + let typed_ast::Expression::FunctionCall { type_scheme, .. } = checked_call else { + unreachable!("proper function call returns typed function call"); + }; + + typed_ast::Expression::MethodCall { + full_span: *full_span, + receiver: Box::new(receiver_checked), + method_ref: typed_ast::MethodRef { + owner: struct_info.name.clone(), + name: method_name, + kind: StructMethodKind::Instance, + }, + method_name_span: *method_name_span, + args: arguments_checked, + type_scheme, + } + } }) } @@ -1502,22 +1924,24 @@ impl TypeChecker { decorators, .. } => { - if body.is_none() { - self.value_namespace - .add_identifier( - function_name.to_compact_string(), - *function_name_span, - CompactString::const_new("foreign function"), - ) - .map_err(|err| Box::new(err.into()))?; - } else { - self.value_namespace - .add_identifier_allow_override( - function_name.to_compact_string(), - *function_name_span, - CompactString::const_new("function"), - ) - .map_err(|err| Box::new(err.into()))?; + if !self.checking_struct_method { + if body.is_none() { + self.value_namespace + .add_identifier( + function_name.to_compact_string(), + *function_name_span, + CompactString::const_new("foreign function"), + ) + .map_err(|err| Box::new(err.into()))?; + } else { + self.value_namespace + .add_identifier_allow_override( + function_name.to_compact_string(), + *function_name_span, + CompactString::const_new("function"), + ) + .map_err(|err| Box::new(err.into()))?; + } } // Save the environment and namespaces to avoid polluting @@ -1737,11 +2161,13 @@ impl TypeChecker { self.value_namespace.restore(); self.type_namespace.restore(); self.env.restore(); - self.env.add_function( - function_name.to_compact_string(), - signature.clone(), - metadata.clone(), - ); + if !self.checking_struct_method { + self.env.add_function( + function_name.to_compact_string(), + signature.clone(), + metadata.clone(), + ); + } typed_ast::Statement::DefineFunction { function_name, @@ -1918,6 +2344,7 @@ impl TypeChecker { struct_name, type_parameters, fields, + methods, } => { self.type_namespace .add_identifier( @@ -1969,6 +2396,33 @@ impl TypeChecker { seen_fields.insert(field, *span); } + let mut seen_member_spans: HashMap<&str, Span> = fields + .iter() + .map(|(span, name, _)| (*name, *span)) + .collect(); + for method in methods { + if let ast::Statement::DefineFunction { + function_name_span, + function_name, + .. + } = method + && let Some(previous_span) = seen_member_spans.get(function_name) + { + return Err(Box::new(TypeCheckError::DuplicateMemberInStructDefinition( + *function_name_span, + *previous_span, + (*function_name).to_string(), + ))); + } else if let ast::Statement::DefineFunction { + function_name_span, + function_name, + .. + } = method + { + seen_member_spans.insert(function_name, *function_name_span); + } + } + let struct_info = StructInfo { definition_span: *struct_name_span, name: struct_name.to_compact_string(), @@ -1989,9 +2443,39 @@ impl TypeChecker { )) }) .collect::>()?, + methods: methods + .iter() + .map(|method| { + let ast::Statement::DefineFunction { + function_name_span, + function_name, + parameters, + .. + } = method + else { + unreachable!("non-function members are rejected earlier"); + }; + + let kind = if !parameters.is_empty() && parameters[0].1 == "self" { + StructMethodKind::Instance + } else { + StructMethodKind::Constructor + }; + + ( + (*function_name).to_compact_string(), + StructMethodInfo { + definition_span: *function_name_span, + kind, + }, + ) + }) + .collect(), }; self.structs .insert(struct_name.to_compact_string(), struct_info.clone()); + self.non_generic_struct_instances + .remove(&struct_name.to_compact_string()); typed_ast::Statement::DefineStruct(struct_info) } @@ -2031,6 +2515,20 @@ impl TypeChecker { self.env.apply(&substitution).map_err(|e| { TypeCheckError::SubstitutionError(elaborated_statement.pretty_print().to_string(), e) })?; + for method_infos in self.methods.values_mut() { + for method_info in method_infos.values_mut() { + method_info + .signature + .fn_type + .apply(&substitution) + .map_err(|e| { + TypeCheckError::SubstitutionError( + elaborated_statement.pretty_print().to_string(), + e, + ) + })?; + } + } if let typed_ast::Statement::DefineDerivedUnit { expr, type_scheme, .. @@ -2088,6 +2586,11 @@ impl TypeChecker { elaborated_statement.update_readable_types(&self.registry); self.env.generalize_types(&dtype_variables); + for method_infos in self.methods.values_mut() { + for method_info in method_infos.values_mut() { + method_info.signature.fn_type.generalize(&dtype_variables); + } + } // Check if there is a typed hole in the statement if let Some((span, type_of_hole)) = elaborated_statement.find_typed_hole()? { @@ -2126,6 +2629,241 @@ impl TypeChecker { let mut checked_statements = vec![]; for statement in statements { + if let ast::Statement::DefineStruct { + struct_name_span, + struct_name, + type_parameters, + fields, + methods, + .. + } = statement + { + let mut seen_member_spans: HashMap<&str, Span> = fields + .iter() + .map(|(span, name, _)| (*name, *span)) + .collect(); + for method in methods { + match method { + ast::Statement::DefineFunction { + function_name_span, + function_name, + .. + } => { + if let Some(previous_span) = seen_member_spans.get(function_name) { + return Err(Box::new( + TypeCheckError::DuplicateMemberInStructDefinition( + *function_name_span, + *previous_span, + (*function_name).to_string(), + ), + )); + } + seen_member_spans.insert(function_name, *function_name_span); + } + _ => { + return Err(Box::new(TypeCheckError::InvalidStructMember( + method.full_span(), + ))); + } + } + } + + checked_statements.push(self.check_statement(statement)?); + + let mut prepared_methods = vec![]; + + for method in methods { + let ast::Statement::DefineFunction { parameters, .. } = method else { + unreachable!("non-function members are rejected above"); + }; + + let method_kind = if !parameters.is_empty() && parameters[0].1 == "self" { + StructMethodKind::Instance + } else { + StructMethodKind::Constructor + }; + + let mut method_to_check = method.clone(); + if let ast::Statement::DefineFunction { + type_parameters: method_type_parameters, + .. + } = &mut method_to_check + { + let mut merged_type_parameters = type_parameters.clone(); + merged_type_parameters.extend(method_type_parameters.clone()); + *method_type_parameters = merged_type_parameters; + } + + if method_kind == StructMethodKind::Instance + && let ast::Statement::DefineFunction { parameters, .. } = + &mut method_to_check + && parameters[0].2.is_none() + { + let self_type_args = type_parameters + .iter() + .map(|(span, name, _)| { + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + *span, + (*name).to_compact_string(), + vec![], + )) + }) + .collect(); + parameters[0].2 = Some(TypeAnnotation::TypeExpression( + TypeExpression::TypeIdentifier( + *struct_name_span, + (*struct_name).to_compact_string(), + self_type_args, + ), + )); + } + let self_type_args: Vec<_> = type_parameters + .iter() + .map(|(span, name, _)| { + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + *span, + (*name).to_compact_string(), + vec![], + )) + }) + .collect(); + Self::rewrite_self_in_method_signature( + &mut method_to_check, + struct_name, + &self_type_args, + )?; + + prepared_methods.push((method_to_check, method_kind)); + } + + for (method_to_check, _method_kind) in &prepared_methods { + let ast::Statement::DefineFunction { function_name, .. } = method_to_check + else { + unreachable!("method_to_check is DefineFunction"); + }; + + let signature = self.provisional_method_signature(method_to_check)?; + self.register_method( + struct_name.to_compact_string(), + function_name.to_compact_string(), + signature, + ); + } + + for (method_to_check, method_kind) in prepared_methods { + self.checking_struct_method = true; + self.current_struct_name = Some(struct_name.to_compact_string()); + let checked_method = self.check_statement(&method_to_check); + self.current_struct_name = None; + self.checking_struct_method = false; + let mut checked_method = checked_method?; + + let ast::Statement::DefineFunction { + function_name_span, + function_name, + type_parameters, + parameters, + return_type_annotation, + .. + } = &method_to_check + else { + unreachable!("method_to_check is DefineFunction"); + }; + + let checked_fn_type = match &checked_method { + typed_ast::Statement::DefineFunction { fn_type, .. } => fn_type.clone(), + _ => unreachable!("checked method is DefineFunction"), + }; + + checked_method = match checked_method { + typed_ast::Statement::DefineFunction { + function_name: method_name, + decorators, + type_parameters, + parameters, + body, + local_variables, + fn_type, + return_type_annotation, + readable_return_type, + } => typed_ast::Statement::DefineMethod { + struct_name: struct_name.to_compact_string(), + method_name, + decorators, + type_parameters, + parameters, + body, + local_variables, + fn_type, + return_type_annotation, + readable_return_type, + }, + _ => unreachable!("checked method is DefineFunction"), + }; + + checked_statements.push(checked_method); + + let signature = FunctionSignature { + name: (*function_name).to_compact_string(), + definition_span: *function_name_span, + type_parameters: type_parameters + .iter() + .map(|(span, name, bound)| { + (*span, (*name).to_compact_string(), bound.clone()) + }) + .collect(), + parameters: parameters + .iter() + .map(|(span, name, annotation)| { + (*span, (*name).to_compact_string(), annotation.clone()) + }) + .collect(), + return_type_annotation: return_type_annotation.clone(), + fn_type: checked_fn_type, + }; + + if method_kind == StructMethodKind::Instance { + let fn_type = match &signature.fn_type { + TypeScheme::Concrete(t) => t.clone(), + TypeScheme::Quantified(_, _) => { + signature + .fn_type + .instantiate(&mut self.name_generator) + .inner + } + }; + + let Type::Fn(parameter_types, _) = &fn_type else { + unreachable!("method type is expected to be function type"); + }; + + let Some(Type::Struct(self_type)) = parameter_types.first() else { + return Err(Box::new(TypeCheckError::InvalidSelfParameterType( + *struct_name_span, + function_name.to_string(), + struct_name.to_string(), + ))); + }; + + if self_type.name.as_str() != *struct_name { + return Err(Box::new(TypeCheckError::InvalidSelfParameterType( + *struct_name_span, + function_name.to_string(), + struct_name.to_string(), + ))); + } + } + + self.register_method( + struct_name.to_compact_string(), + function_name.to_compact_string(), + signature, + ); + } + + continue; + } + checked_statements.push(self.check_statement(statement)?); } diff --git a/numbat/src/typechecker/substitutions.rs b/numbat/src/typechecker/substitutions.rs index 217d384c..767685b5 100644 --- a/numbat/src/typechecker/substitutions.rs +++ b/numbat/src/typechecker/substitutions.rs @@ -263,6 +263,18 @@ impl ApplySubstitution for Expression<'_> { } type_scheme.apply(s) } + Expression::MethodCall { + receiver, + args, + type_scheme, + .. + } => { + receiver.apply(s)?; + for arg in args { + arg.apply(s)?; + } + type_scheme.apply(s) + } Expression::TypedHole(_, type_) => type_.apply(s), } } @@ -283,6 +295,12 @@ impl ApplySubstitution for Statement<'_> { local_variables, fn_type, .. + } + | Statement::DefineMethod { + body, + local_variables, + fn_type, + .. } => { for local_variable in local_variables { local_variable.expr.apply(s)?; diff --git a/numbat/src/typechecker/tests/type_checking.rs b/numbat/src/typechecker/tests/type_checking.rs index ca35b195..172e1bfb 100644 --- a/numbat/src/typechecker/tests/type_checking.rs +++ b/numbat/src/typechecker/tests/type_checking.rs @@ -771,6 +771,254 @@ fn generic_structs() { )); } +#[test] +fn struct_methods() { + assert_successful_typecheck( + " + struct Point { + x: A, + y: A, + fn new(x: A, y: A) -> Self = Point { x: x, y: y } + fn get_x(self) -> A = self.x + fn translate(self, dx: A, dy: A) -> Self = Point { x: self.x + dx, y: self.y + dy } + } + + let p = Point::new(1 a, 2 a) + let px: A = p.get_x() + let px2: A = p.translate(1 a, 2 a).get_x() + ", + ); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn new(x: A, y: A) = Point { x: x, y: y } + } + Point { x: 1 a, y: 1 a }.new() + " + ), + TypeCheckError::ConstructorCalledAsMethod(_, name, struct_name) if name == "new" && struct_name == "Point" + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn get_x(self) -> A = self.x + } + Point::get_x() + " + ), + TypeCheckError::InstanceMethodCalledAsConstructor(_, name, struct_name) if name == "get_x" && struct_name == "Point" + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn new(x: A, y: A) = Point { x: x, y: y } + } + Point::new(1 a) + " + ), + TypeCheckError::WrongArity { .. } + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn get_x(self) -> A = self.x + } + (1 a).get_x() + " + ), + TypeCheckError::MethodCallOnNonStructType(_, name, Type::Dimension(_)) if name == "get_x" + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn get_x(self) -> A = self.x + } + Point::missing(1 a, 2 a) + " + ), + TypeCheckError::MethodNotFound(_, name, struct_name) if name == "missing" && struct_name == "Point" + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn not_self(self: A) -> A = self + } + " + ), + TypeCheckError::InvalidSelfParameterType(_, name, struct_name) if name == "not_self" && struct_name == "Point" + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn x(self) = self.x + fn x() = 1 a + } + ", + ), + TypeCheckError::DuplicateMemberInStructDefinition(_, _, name) if name == "x" + )); + + assert_successful_typecheck( + " + struct Box { + inner: T, + fn new(value: T) -> Self = Box { inner: value } + fn get(self) -> T = self.inner + fn replace(self, value: U) -> Box = Box { inner: value } + } + + let box_a: A = Box::new(1 a).get() + let box_scalar: Scalar = Box::new(1 a).replace(3).get() + ", + ); + + assert!(matches!( + get_typecheck_error( + " + struct Box { + inner: T, + fn bad(self) -> T = self.inner + } + " + ), + TypeCheckError::TypeParameterNameClash(_, name) if name == "T" + )); + + assert_successful_typecheck( + " + struct Box { + inner: T, + fn into_list(self) -> List = [self] + fn map_self(self, f: Fn[(Self) -> Self]) -> Self = f(self) + } + + fn id_box(b) = b + + let xs: List> = Box { inner: 1 a }.into_list() + let y: Box = Box { inner: 1 a }.map_self(id_box) + ", + ); + + assert!(matches!( + get_typecheck_error("fn id(x: Self) -> Self = x"), + TypeCheckError::SelfTypeOutsideStructMethod(_) + )); + + assert_successful_typecheck( + " + struct Flag { + n: Scalar, + fn even(self) -> Bool = + if self.n == 0 then true else Flag { n: self.n - 1 }.odd() + fn odd(self) -> Bool = + if self.n == 0 then false else Flag { n: self.n - 1 }.even() + } + + let x: Bool = Flag { n: 4 }.even() + let y: Bool = Flag { n: 5 }.odd() + ", + ); + + assert_successful_typecheck( + " + struct Inner { + x: A, + } + + struct Outer { + inner: T, + fn wrap(value: T) -> Self = Self { inner: value } + fn replace(self, value: U) -> Outer = Outer { inner: value } + } + + struct OuterPoint { + inner: Inner, + fn shift(self, dx: A) -> Self = Self { inner: Inner { x: self.inner.x + dx } } + } + + struct Counter { + value: Scalar, + fn new(n: Scalar) -> Self = result + where result = + if n == 0 then Self { value: 0 } else Counter::new(n - 1).inc() + fn inc(self) -> Self = Self { value: self.value + 1 } + } + + let p: A = OuterPoint { inner: Inner { x: 1 a } }.shift(2 a).inner.x + let q: Scalar = Outer::wrap(1 a).replace(5).inner + let count: Scalar = Counter::new(4).value + ", + ); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn translate(self, dx: A, dy: A) -> Self = Point { x: self.x + dx, y: self.y + dy } + } + Point { x: 1 a, y: 2 a }.translate(1 a) + " + ), + TypeCheckError::WrongArity { .. } + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + fn translate(self, dx: A, dy: A) -> Self = Point { x: self.x + dx, y: self.y + dy } + } + Point { x: 1 a, y: 2 a }.translate(1 a, 1 b) + " + ), + TypeCheckError::IncompatibleDimensions(..) + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + fn missing(self: Self) -> Self + } + " + ), + TypeCheckError::UnknownForeignFunction(_, name) if name == "missing" + )); +} + #[test] fn lists() { assert_successful_typecheck("[]"); diff --git a/numbat/src/typed_ast.rs b/numbat/src/typed_ast.rs index 4e803427..3380e418 100644 --- a/numbat/src/typed_ast.rs +++ b/numbat/src/typed_ast.rs @@ -337,6 +337,19 @@ pub struct StructInfo { pub name: CompactString, pub kind: StructKind, pub fields: IndexMap, + pub methods: IndexMap, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StructMethodKind { + Constructor, + Instance, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StructMethodInfo { + pub definition_span: Span, + pub kind: StructMethodKind, } /// A monomorphic type (no quantifiers). @@ -607,6 +620,7 @@ impl Type { name: info.name.clone(), kind: instantiated_kind, fields: instantiated_fields, + methods: info.methods.clone(), })) } Type::List(element_type) => { @@ -740,6 +754,14 @@ pub enum Expression<'a> { struct_type: TypeScheme, field_type: TypeScheme, }, + MethodCall { + full_span: Span, + receiver: Box>, + method_ref: MethodRef<'a>, + method_name_span: Span, + args: Vec>, + type_scheme: TypeScheme, + }, List { span: Span, elements: Vec>, @@ -748,6 +770,13 @@ pub enum Expression<'a> { TypedHole(Span, TypeScheme), } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct MethodRef<'a> { + pub owner: CompactString, + pub name: &'a str, + pub kind: StructMethodKind, +} + impl Expression<'_> { pub fn full_span(&self) -> Span { match self { @@ -782,6 +811,7 @@ impl Expression<'_> { Expression::String(span, _) => *span, Expression::InstantiateStruct { span, .. } => *span, Expression::AccessField { full_span, .. } => *full_span, + Expression::MethodCall { full_span, .. } => *full_span, Expression::List { span, .. } => *span, Expression::TypedHole(span, _) => *span, } @@ -818,6 +848,23 @@ pub enum Statement<'a> { return_type_annotation: Option, readable_return_type: Markup, }, + DefineMethod { + struct_name: CompactString, + method_name: &'a str, + decorators: Vec>, + type_parameters: Vec<(&'a str, Option)>, + parameters: Vec<( + Span, // span of the parameter + &'a str, // parameter name + Option, // parameter type annotation + Markup, // readable parameter type + )>, + body: Option>, + local_variables: Vec>, + fn_type: TypeScheme, + return_type_annotation: Option, + readable_return_type: Markup, + }, DefineDimension(&'a str, Vec), DefineBaseUnit { name: &'a str, @@ -889,6 +936,15 @@ impl Statement<'_> { return_type_annotation, readable_return_type, .. + } + | Statement::DefineMethod { + type_parameters, + parameters, + local_variables, + fn_type, + return_type_annotation, + readable_return_type, + .. } => { let (fn_type, _) = fn_type.instantiate_for_printing(Some(type_parameters.iter().map(|(n, _)| *n))); @@ -989,6 +1045,12 @@ impl Statement<'_> { local_variables, fn_type, .. + } + | Statement::DefineMethod { + parameters, + local_variables, + fn_type, + .. } => { let mut bindings = Vec::new(); @@ -1046,6 +1108,7 @@ impl Expression<'_> { Expression::List { type_scheme, .. } => { Type::List(Box::new(type_scheme.unsafe_as_concrete())) } + Expression::MethodCall { type_scheme, .. } => type_scheme.unsafe_as_concrete(), Expression::TypedHole(_, type_) => type_.unsafe_as_concrete(), } } @@ -1077,6 +1140,7 @@ impl Expression<'_> { }, ), }, + Expression::MethodCall { type_scheme, .. } => type_scheme.clone(), Expression::TypedHole(_, type_) => type_.clone(), } } @@ -1246,6 +1310,16 @@ impl PrettyPrint for Statement<'_> { fn_type, readable_return_type, .. + } + | Statement::DefineMethod { + method_name: function_name, + type_parameters, + parameters, + body, + local_variables, + fn_type, + readable_return_type, + .. } => { let (fn_type, type_parameters) = fn_type.instantiate_for_printing(Some(type_parameters.iter().map(|(n, _)| *n))); @@ -1412,6 +1486,7 @@ fn with_parens(expr: &Expression) -> Markup { | Expression::String(..) | Expression::InstantiateStruct { .. } | Expression::AccessField { .. } + | Expression::MethodCall { .. } | Expression::List { .. } | Expression::TypedHole(_, _) => expr.pretty_print(), Expression::UnaryOperator { .. } @@ -1726,6 +1801,26 @@ impl PrettyPrint for Expression<'_> { .sum() + m::operator("]") } + MethodCall { + receiver, + method_ref, + args, + .. + } => { + receiver.pretty_print() + + m::operator(match method_ref.kind { + StructMethodKind::Constructor => "::", + StructMethodKind::Instance => ".", + }) + + m::identifier(method_ref.name.to_compact_string()) + + m::operator("(") + + itertools::Itertools::intersperse( + args.iter().map(|e: &Expression| e.pretty_print()), + m::operator(",") + m::space(), + ) + .sum() + + m::operator(")") + } TypedHole(_, _) => m::operator("?"), } } diff --git a/numbat/src/value.rs b/numbat/src/value.rs index 77da60cc..b863c878 100644 --- a/numbat/src/value.rs +++ b/numbat/src/value.rs @@ -40,7 +40,7 @@ pub enum Value { DateTime(Zoned), FunctionReference(FunctionReference), FormatSpecifiers(Option), - StructInstance(Arc, Vec), + StructInstance(Arc, Arc<[Value]>), List(NumbatList), } @@ -91,7 +91,7 @@ impl Value { } #[track_caller] - pub fn unsafe_as_struct_fields(self) -> Vec { + pub fn unsafe_as_struct_fields(&self) -> &[Value] { if let Value::StructInstance(_, values) = self { values } else { @@ -134,7 +134,7 @@ impl std::fmt::Display for Value { struct_info .fields .keys() - .zip(values) + .zip(values.iter()) .map(|(name, value)| name.to_owned() + ": " + &value.to_string()) .join(", ") ) @@ -175,12 +175,16 @@ impl Value { } else { crate::markup::space() + itertools::Itertools::intersperse( - struct_info.fields.keys().zip(values).map(|(name, val)| { - crate::markup::identifier(name.clone()) - + crate::markup::operator(":") - + crate::markup::space() - + val.pretty_print_with(options) - }), + struct_info + .fields + .keys() + .zip(values.iter()) + .map(|(name, val)| { + crate::markup::identifier(name.clone()) + + crate::markup::operator(":") + + crate::markup::space() + + val.pretty_print_with(options) + }), crate::markup::operator(",") + crate::markup::space(), ) .sum() diff --git a/numbat/src/vm.rs b/numbat/src/vm.rs index d35c6133..707dd540 100644 --- a/numbat/src/vm.rs +++ b/numbat/src/vm.rs @@ -278,6 +278,12 @@ pub struct ExecutionContext<'a> { pub typechecker: &'a TypeChecker, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MethodCallable { + Normal(u16), + Foreign(u16), +} + /// Metadata for a single FFI call argument #[derive(Clone)] pub struct FfiCallArg { @@ -308,6 +314,8 @@ pub struct Vm { /// struct metadata, used so we can display struct fields at runtime struct_infos: IndexMap>, + /// Mapping from (struct name, method name) to callable target + method_callables: HashMap<(CompactString, CompactString), MethodCallable>, /// Unit prefixes in use prefixes: Vec, @@ -348,6 +356,7 @@ impl Vm { current_chunk_index: 0, constants: vec![], struct_infos: IndexMap::new(), + method_callables: HashMap::new(), prefixes: vec![], strings: vec![], unit_information: vec![], @@ -486,9 +495,14 @@ impl Vm { (self.unit_information.len() - 1) as u16 // TODO: this can overflow, see above } - pub(crate) fn begin_function(&mut self, name: &str) { + pub(crate) fn begin_function(&mut self, name: &str) -> u16 { self.bytecode.push((name.into(), vec![], vec![])); - self.current_chunk_index = self.bytecode.len() - 1 + self.current_chunk_index = self.bytecode.len() - 1; + self.current_chunk_index as u16 + } + + pub(crate) fn begin_reserved_function(&mut self, idx: u16) { + self.current_chunk_index = idx as usize; } pub(crate) fn end_function(&mut self) { @@ -510,6 +524,42 @@ impl Vm { position as u16 } + pub(crate) fn register_method_function(&mut self, owner: &str, method: &str, idx: u16) { + self.method_callables.insert( + (owner.to_compact_string(), method.to_compact_string()), + MethodCallable::Normal(idx), + ); + } + + pub(crate) fn register_foreign_method(&mut self, owner: &str, method: &str, ffi_idx: u16) { + self.method_callables.insert( + (owner.to_compact_string(), method.to_compact_string()), + MethodCallable::Foreign(ffi_idx), + ); + } + + pub(crate) fn get_method_callable(&self, owner: &str, method: &str) -> Option { + self.method_callables + .get(&(owner.to_compact_string(), method.to_compact_string())) + .copied() + } + + pub(crate) fn get_empty_method_function_idx(&self, owner: &str, method: &str) -> Option { + let MethodCallable::Normal(idx) = self.get_method_callable(owner, method)? else { + return None; + }; + let idx_usize = idx as usize; + if self + .bytecode + .get(idx_usize) + .is_some_and(|(_, code, spans)| code.is_empty() && spans.is_empty()) + { + Some(idx) + } else { + None + } + } + pub(crate) fn add_foreign_function(&mut self, name: &str, arity: ArityRange) { // `key: &'static str`, whereas `name: &'non_static str` let (key, ff) = ffi::functions().get_key_value(name).unwrap(); @@ -1163,14 +1213,19 @@ impl Vm { content.push(self.pop()); } - self.stack.push(Value::StructInstance(struct_info, content)); + self.stack.push(Value::StructInstance( + struct_info, + Arc::<[Value]>::from(content), + )); } Op::AccessStructField => { let field_idx = self.read_u16(); - - let mut fields = self.pop().unsafe_as_struct_fields(); - - let value = fields.swap_remove(field_idx as usize); + let struct_value = self.pop(); + let fields = struct_value.unsafe_as_struct_fields(); + let value = fields + .get(field_idx as usize) + .expect("field index should be valid") + .clone(); self.stack.push(value); } Op::BuildList => { diff --git a/numbat/tests/interpreter.rs b/numbat/tests/interpreter.rs index e92a28cb..9a11a363 100644 --- a/numbat/tests/interpreter.rs +++ b/numbat/tests/interpreter.rs @@ -1243,6 +1243,300 @@ fn test_temperature_syntactic_sugar() { expect_pretty_print("let t = 68 fahrenheit", "let t: Temperature = 68 °F"); } +#[test] +fn test_struct_methods() { + expect_output( + " + struct Point { + x: Scalar, + y: Scalar, + fn origin() = Point { x: 0, y: 0 } + fn shift_x(self, dx: Scalar) = Point { x: self.x + dx, y: self.y } + fn get_x(self) = self.x + } + + Point::origin().shift_x(3).get_x() + ", + "3", + ); + + expect_output( + " + struct Point { + x: Length, + y: Length, + fn new(x: Length, y: Length) -> Self = Point { x: x, y: y } + fn translate(self, dx: Length, dy: Length) -> Self = Self { x: self.x + dx, y: self.y + dy } + fn get_x(self) -> Length = self.x + } + Point::new(3 m, 4 m).translate(2 m, 0 m).get_x() + ", + "5 m", + ); + + expect_output( + " + struct Box { + inner: T, + fn wrap(value: T) -> Box = Box { inner: value } + fn replace(self, value: U) -> Box = Box { inner: value } + } + Box::wrap(2 m).replace(5).inner + ", + "5", + ); + + expect_output( + " + struct Flag { + n: Scalar, + fn even(self) -> Bool = + if self.n == 0 then true else Flag { n: self.n - 1 }.odd() + fn odd(self) -> Bool = + if self.n == 0 then false else Flag { n: self.n - 1 }.even() + } + + if Flag { n: 7 }.odd() then 1 else 0 + ", + "1", + ); + + expect_output( + " + fn combine(x: Scalar, y: Scalar) -> Scalar = x + y + + struct Vec2 { + x: Length, + y: Length, + fn dot(self, other: Vec2) -> Area = self.x * other.x + self.y * other.y + } + + let v1 = Vec2 { x: 1 m, y: 2 m } + let v2 = Vec2 { x: 3 m, y: 4 m } + v1.dot(v2) + combine(1, 2) m^2 + ", + "14 m²", + ); + + expect_failure( + " + struct Point { + fn magnitude(self) = self.x + x: Scalar + } + ", + "Fields must be declared before methods in struct body", + ); + + expect_failure( + " + fn id(x: Self) -> Self = x + ", + "`Self` can only be used inside struct method definitions", + ); +} + +#[test] +fn test_struct_method_namespace_stress() { + // Method names should be struct-scoped even when they match prelude function names. + expect_output( + " + struct Sample { + x: Scalar, + fn sqrt(self) -> Scalar = self.x + 1 + fn sin(self) -> Scalar = self.x + 2 + } + + let sample_value = Sample { x: 9 } + sample_value.sqrt() + sample_value.sin() + sqrt(9) + sin(0) + ", + "24", + ); + + // Constructor-style methods should also be isolated from prelude/global names. + expect_output( + " + struct AngleBox { + value: Scalar, + fn atan2(x: Scalar, y: Scalar) -> Self = AngleBox { value: x + y } + } + + AngleBox::atan2(2, 3).value + atan2(10, 0) / (pi / 2) + ", + "6", + ); + + // Imported module functions should not conflict with methods of the same name. + let mut ctx = get_test_context(); + let _ = ctx + .interpret("use extra::algebra", CodeSource::Internal) + .unwrap(); + expect_output_with_context( + &mut ctx, + " + struct Eqn { + x: Scalar, + fn quadratic_equation(self) -> Scalar = self.x + } + + Eqn { x: 5 }.quadratic_equation() + len(quadratic_equation(1, 0, -1)) + ", + "7", + ); + + // Cross-struct same-name methods must dispatch to the correct implementation. + expect_output( + " + struct A { + x: Scalar, + fn value(self) -> Scalar = self.x + 100 + } + + struct B { + x: Scalar, + fn value(self) -> Scalar = self.x + 1000 + } + + A { x: 1 }.value() + B { x: 1 }.value() + ", + "1102", + ); +} + +#[test] +fn test_struct_field_projection_stability() { + expect_output( + " + struct Pair { + x: Scalar, + y: Scalar, + fn score(self) -> Scalar = self.x + self.y + self.x + self.y + } + + Pair { x: 2, y: 3 }.score() + ", + "10", + ); + + expect_output( + " + struct Pair { + x: Scalar, + y: Scalar, + } + + struct Wrap { + pair: Pair, + fn twice_x(self) -> Scalar = self.pair.x + self.pair.x + } + + Wrap { pair: Pair { x: 4, y: 1 } }.twice_x() + ", + "8", + ); +} + +#[test] +fn test_self_struct_instantiation_contexts() { + expect_output( + " + struct Outer { + inner: T, + fn wrap(value: T) -> Self = Self { inner: value } + fn replace(self, value: U) -> Outer = Outer { inner: value } + } + + Outer::wrap(2 m).replace(5).inner + ", + "5", + ); + + expect_output( + " + struct Inner { + x: Length, + } + + struct Outer { + inner: Inner, + fn shift(self, dx: Length) -> Self = Self { + inner: Inner { x: self.inner.x + dx } + } + } + + Outer { inner: Inner { x: 1 m } }.shift(2 m).inner.x + ", + "3 m", + ); + + expect_output( + " + struct Counter { + value: Scalar, + + fn new(n: Scalar) -> Self = result + where result = + if n == 0 then Self { value: 0 } else Counter::new(n - 1).inc() + + fn inc(self) -> Self = Self { value: self.value + 1 } + } + + Counter::new(4).value + ", + "4", + ); +} + +#[test] +fn test_ffi_struct_shape_for_chemical_elements() { + expect_output( + " + use chemistry::elements + element(\"H\").symbol + ", + "\"H\"", + ); + + expect_output( + " + use chemistry::elements + element(\"hydrogen\").atomic_number + ", + "1", + ); + + expect_failure( + " + use chemistry::elements + element(\"definitely-not-an-element\") + ", + "Chemical element not found", + ); +} + +#[test] +fn test_where_locals_are_evaluated_once() { + expect_output( + " + fn once() -> Scalar = x - x where x = random() + once() + ", + "0", + ); + + expect_output( + " + struct Sample { + value: Scalar, + fn once(self) -> Scalar = x - x where x = random() + self.value + } + + Sample { value: 5 }.once() + ", + "0", + ); +} + #[cfg(test)] mod tests { use super::*; diff --git a/numbat/tests/snapshots/example_snapshots__numbat_syntax@numbat_syntax.nbt.snap b/numbat/tests/snapshots/example_snapshots__numbat_syntax@numbat_syntax.nbt.snap index 43008720..9f323135 100644 --- a/numbat/tests/snapshots/example_snapshots__numbat_syntax@numbat_syntax.nbt.snap +++ b/numbat/tests/snapshots/example_snapshots__numbat_syntax@numbat_syntax.nbt.snap @@ -9,4 +9,5 @@ value of pi = 3.14159 sqrt(10) = 3.16228 value of π ≈ 3.142 = Length / Time +Element { name: "Neptunium", atomic_number: 93, density: 20.45 g/cm³ } = 0.08988 g/L