diff --git a/book/build.py b/book/build.py index 64f03c6e..05060235 100755 --- a/book/build.py +++ b/book/build.py @@ -15,7 +15,6 @@ import os import sys - SCRIPT_DIR = Path(__file__).parent.resolve() WORKSPACE_DIR = SCRIPT_DIR.parent NUMBAT_CRATE_DIR = WORKSPACE_DIR / "numbat" @@ -137,6 +136,7 @@ def generate_all_examples(): generate_example("body_mass_index", "Body mass index") generate_example("factorial", "Factorial", strip_asserts=False) generate_example("medication_dosage", "Medication dosage") + generate_example("matrix", "Matrices") generate_example("molarity", "Molarity") generate_example("musical_note_frequency", "Musical note frequency") generate_example("paper_size", "Paper sizes") diff --git a/book/src/basics/lists.md b/book/src/basics/lists.md index aa110e88..0288abf9 100644 --- a/book/src/basics/lists.md +++ b/book/src/basics/lists.md @@ -16,6 +16,17 @@ Lists can be created using the `[…]` syntax. For example: The type of a list is written as `List`, where `T` is the type of the elements. The types of the lists above are `List`, `List`, and `List>`, respectively. +Lists can be indexed with square brackets: + +```nbt +let xs = [30 cm, 110 cm, 2 m] + +xs[0] # 30 cm +xs[2] # 2 m +``` + +List indices must be non-negative integers. An out-of-bounds index yields a runtime error. + The standard library provides a [number of functions](../prelude/functions/lists.md) to work with lists. Some useful things to do with lists are: ```nbt diff --git a/book/src/basics/structs.md b/book/src/basics/structs.md index 4e2e6446..5bce6801 100644 --- a/book/src/basics/structs.md +++ b/book/src/basics/structs.md @@ -24,6 +24,23 @@ let tungsten = Element { } ``` +Struct fields can also define default values: + +```nbt +struct PaperSize { + width: Length = 210 mm, + height: Length = 297 mm, +} +``` + +When instantiating, fields with defaults may be omitted: + +```nbt +let a4 = PaperSize {} +let custom = PaperSize { width: 100 mm } # height defaults to 297 mm +let swapped = PaperSize { height: 400 mm } # order does not matter +``` + Fields can be accessed using dot notation: ```nbt @@ -33,6 +50,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 +99,95 @@ 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, + + @add + fn add(self, rhs: Self) -> Self = + Vec { x: self.x + rhs.x, y: self.y + rhs.y } + + @mul + fn scale(self, factor: Scalar) -> Self = + Vec { x: self.x * factor, y: self.y * factor } + + @rmul + fn scale_from_left(self, lhs: Scalar) -> Self = + Vec { x: lhs * self.x, y: lhs * self.y } + + fn dot(self, other: Vec) -> X * Y = + self.x * other.x + self.y * other.y + + fn new(x: X, y: X) -> Self = + Vec { x: x, y: y } +} + +let v1 = Vec::new(1 m, 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 combined = v1 + v2 +let v4 = 2 * v1 +let dp_m = v1.dot(v2) # 11 m² +let dp_cm = v1.dot(v2_cm) # 110_000 cm² +``` + +Operator overloads are declared on struct instance methods with decorators such as `@add`, `@sub`, `@mul`, and `@div`. The decorator itself is bare; the compiler infers the left-hand side from `self`, the right-hand side from the second parameter, and the result type from the return type: + +```nbt +struct Vec { + x: X, + y: X, + + @add + fn add(self, rhs: Self) -> Self = + Vec { x: self.x + rhs.x, y: self.y + rhs.y } + + @mul + fn scale(self, factor: Scalar) -> Self = + Vec { x: self.x * factor, y: self.y * factor } + + @rmul + fn scale_from_left(self, lhs: Scalar) -> Self = + Vec { x: lhs * self.x, y: lhs * self.y } +} + +let v = Vec { x: 3 m, y: 4 m } +let w = Vec { x: 300 cm, y: 400 cm } + +let combined = v + w +let scaled = v * 2 +let scaled_from_left = 2 * v +``` + +Reverse decorators `@radd`, `@rsub`, `@rmul`, and `@rdiv` allow the struct to handle expressions where it appears on the right-hand side. In the example above, `@rmul fn scale_from_left(self, lhs: Scalar)` enables `2 * v`. + +Structs can also define indexing behavior with the `@index` decorator on an instance method. The parameters after `self` determine how many index arguments the struct accepts: + +```nbt +struct Grid2 { + a: Scalar, + b: Scalar, + c: Scalar, + d: Scalar, + + @index + fn get(self, row: Scalar, col: Scalar) -> Scalar = + if row == 0 then + if col == 0 then self.a else self.b + else + if col == 0 then self.c else self.d +} + +let g = Grid2 { a: 1, b: 2, c: 3, d: 4 } + +assert_eq(g[0, 1], 2) +assert_eq(g[1, 0], 3) +``` + +This is directional and struct-owned, just like operator methods. Builtin lists also support indexing with `xs[i]`. diff --git a/book/src/examples/example-matrix.md b/book/src/examples/example-matrix.md new file mode 100644 index 00000000..51a20623 --- /dev/null +++ b/book/src/examples/example-matrix.md @@ -0,0 +1,86 @@ + + +# Matrices + +[:material-play-circle: Run this example](https://numbat.dev/?q=%23+A+dynamically+sized+matrix+backed+by+nested+lists.%0A%23%0A%23+This+example+uses%3A%0A%23+-+a+struct+with+List%3CList%3CScalar%3E%3E+storage%0A%23+-+builtin+list+indexing+%28%60xs%5Bi%5D%60%29%0A%23+-+custom+struct+indexing+via+%60matrix%5Brow%2C+col%5D%60%0A%23+-+struct-owned+matrix+multiplication+via+%60%40mul%60%0A%0Astruct+Matrix+%7B%0A++++rows%3A+Scalar%2C%0A++++cols%3A+Scalar%2C%0A++++data%3A+List%3CList%3CScalar%3E%3E%2C%0A%0A++++fn+from_rows%28data%3A+List%3CList%3CScalar%3E%3E%29+-%3E+Self+%3D%0A++++++++Self+%7B%0A++++++++++++rows%3A+len%28data%29%2C%0A++++++++++++cols%3A+if+is_empty%28data%29+then+0+else+len%28data%5B0%5D%29%2C%0A++++++++++++data%3A+data%2C%0A++++++++%7D%0A%0A++++%40index%0A++++fn+get%28self%2C+row%3A+Scalar%2C+col%3A+Scalar%29+-%3E+Scalar+%3D%0A++++++++self.data%5Brow%5D%5Bcol%5D%0A%0A++++fn+dot_product%28lhs%3A+Self%2C+rhs%3A+Self%2C+row%3A+Scalar%2C+col%3A+Scalar%2C+i%3A+Scalar%29+-%3E+Scalar+%3D%0A++++++++if+i+%3D%3D+lhs.cols%0A++++++++++++then+0%0A++++++++++++else+lhs%5Brow%2C+i%5D+%2A+rhs%5Bi%2C+col%5D+%2B+Matrix%3A%3Adot_product%28lhs%2C+rhs%2C+row%2C+col%2C+i+%2B+1%29%0A%0A++++fn+product_row%28lhs%3A+Self%2C+rhs%3A+Self%2C+row%3A+Scalar%2C+col%3A+Scalar%29+-%3E+List%3CScalar%3E+%3D%0A++++++++if+col+%3D%3D+rhs.cols%0A++++++++++++then+%5B%5D%0A++++++++++++else+cons%28%0A++++++++++++++++Matrix%3A%3Adot_product%28lhs%2C+rhs%2C+row%2C+col%2C+0%29%2C%0A++++++++++++++++Matrix%3A%3Aproduct_row%28lhs%2C+rhs%2C+row%2C+col+%2B+1%29%2C%0A++++++++++++%29%0A%0A++++fn+product_rows%28lhs%3A+Self%2C+rhs%3A+Self%2C+row%3A+Scalar%29+-%3E+List%3CList%3CScalar%3E%3E+%3D%0A++++++++if+row+%3D%3D+lhs.rows%0A++++++++++++then+%5B%5D%0A++++++++++++else+cons%28%0A++++++++++++++++Matrix%3A%3Aproduct_row%28lhs%2C+rhs%2C+row%2C+0%29%2C%0A++++++++++++++++Matrix%3A%3Aproduct_rows%28lhs%2C+rhs%2C+row+%2B+1%29%2C%0A++++++++++++%29%0A%0A++++%40mul%0A++++fn+mul%28self%2C+rhs%3A+Self%29+-%3E+Self+%3D%0A++++++++Matrix%3A%3Afrom_rows%28Matrix%3A%3Aproduct_rows%28self%2C+rhs%2C+0%29%29%0A%7D%0A%0Alet+identity_2+%3D+Matrix%3A%3Afrom_rows%28%5B%0A++++%5B1%2C+0%5D%2C%0A++++%5B0%2C+1%5D%2C%0A%5D%29%0A%0Alet+wide+%3D+Matrix%3A%3Afrom_rows%28%5B%0A++++%5B1%2C+2%2C+3%5D%2C%0A++++%5B4%2C+5%2C+6%5D%2C%0A%5D%29%0A%0Alet+tall+%3D+Matrix%3A%3Afrom_rows%28%5B%0A++++%5B1%2C+2%5D%2C%0A++++%5B3%2C+4%5D%2C%0A++++%5B5%2C+6%5D%2C%0A%5D%29%0A%0Alet+product+%3D+wide+%2A+tall%0Alet+unchanged+%3D+identity_2+%2A+identity_2%0A%0A%0A%0A%0A%0A%0Aprint%28%222x2+identity%3A+%7Bidentity_2%7D%22%29%0Aprint%28%222x3+matrix%3A+++%7Bwide%7D%22%29%0Aprint%28%223x2+matrix%3A+++%7Btall%7D%22%29%0Aprint%28%22product%3A++++++%7Bproduct%7D%22%29%0A){ .md-button .md-button--primary } + +```numbat +# A dynamically sized matrix backed by nested lists. +# +# This example uses: +# - a struct with List> storage +# - builtin list indexing (`xs[i]`) +# - custom struct indexing via `matrix[row, col]` +# - struct-owned matrix multiplication via `@mul` + +struct Matrix { + rows: Scalar, + cols: Scalar, + data: List>, + + fn from_rows(data: List>) -> Self = + Self { + rows: len(data), + cols: if is_empty(data) then 0 else len(data[0]), + data: data, + } + + @index + fn get(self, row: Scalar, col: Scalar) -> Scalar = + self.data[row][col] + + fn dot_product(lhs: Self, rhs: Self, row: Scalar, col: Scalar, i: Scalar) -> Scalar = + if i == lhs.cols + then 0 + else lhs[row, i] * rhs[i, col] + Matrix::dot_product(lhs, rhs, row, col, i + 1) + + fn product_row(lhs: Self, rhs: Self, row: Scalar, col: Scalar) -> List = + if col == rhs.cols + then [] + else cons( + Matrix::dot_product(lhs, rhs, row, col, 0), + Matrix::product_row(lhs, rhs, row, col + 1), + ) + + fn product_rows(lhs: Self, rhs: Self, row: Scalar) -> List> = + if row == lhs.rows + then [] + else cons( + Matrix::product_row(lhs, rhs, row, 0), + Matrix::product_rows(lhs, rhs, row + 1), + ) + + @mul + fn mul(self, rhs: Self) -> Self = + Matrix::from_rows(Matrix::product_rows(self, rhs, 0)) +} + +let identity_2 = Matrix::from_rows([ + [1, 0], + [0, 1], +]) + +let wide = Matrix::from_rows([ + [1, 2, 3], + [4, 5, 6], +]) + +let tall = Matrix::from_rows([ + [1, 2], + [3, 4], + [5, 6], +]) + +let product = wide * tall +let unchanged = identity_2 * identity_2 + + + + + + +print("2x2 identity: {identity_2}") +print("2x3 matrix: {wide}") +print("3x2 matrix: {tall}") +print("product: {product}") +``` diff --git a/book/src/examples/example-numbat_syntax.md b/book/src/examples/example-numbat_syntax.md index 35611e6d..afc18586 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,57 @@ 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, + y: D = 0 m, # Struct defaults can be provided per field + + # 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 } + @add + fn add(self, rhs: Self) -> Self = Self { x: self.x + rhs.x, y: self.y + rhs.y } + @mul + fn scale(self, factor: Scalar) -> Self = Self { x: self.x * factor, y: self.y * factor } + @rmul + fn scale_from_left(self, lhs: Scalar) -> Self = Self { x: lhs * self.x, y: lhs * self.y } + 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) +let combined = v + w # Operator methods can provide struct-owned arithmetic +let left_scaled = 2 * v # Reverse operator methods handle lhs-owned scalar syntax + +assert_eq(v.dot(w) -> m², 25 m²) # Method generic + unit conversion +assert_eq(combined.x -> m, 6 m) +assert_eq(left_scaled.y -> m, 8 m) + +let only_x = Vec2 { x: 5 m } # Omitted fields with defaults are auto-filled + +assert_eq(only_x.x, 5 m) +assert_eq(only_x.y, 0 m) + +let distances = [3 m, 4 m, 5 m] # Builtin lists support indexing +assert_eq(distances[1], 4 m) + +struct Grid2 { # Struct methods can define custom indexing syntax + a: Scalar, + b: Scalar, + c: Scalar, + d: Scalar, + + @index + fn get(self, row: Scalar, col: Scalar) -> Scalar = + if row == 0 then + if col == 0 then self.a else self.b + else + if col == 0 then self.c else self.d +} + +let grid = Grid2 { a: 1, b: 2, c: 3, d: 4 } +assert_eq(grid[1, 0], 3) ``` 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/book/zensical.toml b/book/zensical.toml index be2f8ddc..a61f4c44 100644 --- a/book/zensical.toml +++ b/book/zensical.toml @@ -26,6 +26,7 @@ nav = [ "examples/example-factorial.md", "examples/example-pipe_flow_rate.md", "examples/example-medication_dosage.md", + "examples/example-matrix.md", "examples/example-molarity.md", "examples/example-musical_note_frequency.md", "examples/example-paper_size.md", 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/matrix.nbt b/examples/matrix.nbt new file mode 100644 index 00000000..54d9dd10 --- /dev/null +++ b/examples/matrix.nbt @@ -0,0 +1,96 @@ +# A dynamically sized matrix backed by nested lists. +# +# This example uses: +# - a struct with List> storage +# - builtin list indexing (`xs[i]`) +# - custom struct indexing via `matrix[row, col]` +# - struct-owned matrix multiplication via `@mul` + +struct Matrix { + rows: Scalar, + cols: Scalar, + data: List>, + + fn from_rows(data: List>) -> Self = + Self { + rows: len(data), + cols: if is_empty(data) then 0 else len(data[0]), + data: data, + } + + @index + fn get(self, row: Scalar, col: Scalar) -> Scalar = + self.data[row][col] + + fn dot_product(lhs: Self, rhs: Self, row: Scalar, col: Scalar, i: Scalar) -> Scalar = + if i == lhs.cols + then 0 + else lhs[row, i] * rhs[i, col] + Matrix::dot_product(lhs, rhs, row, col, i + 1) + + fn product_row(lhs: Self, rhs: Self, row: Scalar, col: Scalar) -> List = + if col == rhs.cols + then [] + else cons( + Matrix::dot_product(lhs, rhs, row, col, 0), + Matrix::product_row(lhs, rhs, row, col + 1), + ) + + fn product_rows(lhs: Self, rhs: Self, row: Scalar) -> List> = + if row == lhs.rows + then [] + else cons( + Matrix::product_row(lhs, rhs, row, 0), + Matrix::product_rows(lhs, rhs, row + 1), + ) + + @mul + fn mul(self, rhs: Self) -> Self = + Matrix::from_rows(Matrix::product_rows(self, rhs, 0)) +} + +let identity_2 = Matrix::from_rows([ + [1, 0], + [0, 1], +]) + +let wide = Matrix::from_rows([ + [1, 2, 3], + [4, 5, 6], +]) + +let tall = Matrix::from_rows([ + [1, 2], + [3, 4], + [5, 6], +]) + +let product = wide * tall +let unchanged = identity_2 * identity_2 + +assert_eq(identity_2.rows, 2) +assert_eq(identity_2.cols, 2) +assert_eq(identity_2[1, 1], 1) + +assert_eq(wide.rows, 2) +assert_eq(wide.cols, 3) +assert_eq(wide[0, 2], 3) +assert_eq(wide[1, 0], 4) + +assert_eq(tall.rows, 3) +assert_eq(tall.cols, 2) +assert_eq(tall[2, 1], 6) + +assert_eq(product.rows, 2) +assert_eq(product.cols, 2) +assert_eq(product[0, 0], 22) +assert_eq(product[0, 1], 28) +assert_eq(product[1, 0], 49) +assert_eq(product[1, 1], 64) + +assert_eq(unchanged[0, 0], 1) +assert_eq(unchanged[1, 1], 1) + +print("2x2 identity: {identity_2}") +print("2x3 matrix: {wide}") +print("3x2 matrix: {tall}") +print("product: {product}") diff --git a/examples/numbat_syntax.nbt b/examples/numbat_syntax.nbt index fb020107..bc3b33e6 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,56 @@ 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, + y: D = 0 m, # Struct defaults can be provided per field + + # 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 } + @add + fn add(self, rhs: Self) -> Self = Self { x: self.x + rhs.x, y: self.y + rhs.y } + @mul + fn scale(self, factor: Scalar) -> Self = Self { x: self.x * factor, y: self.y * factor } + @rmul + fn scale_from_left(self, lhs: Scalar) -> Self = Self { x: lhs * self.x, y: lhs * self.y } + 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) +let combined = v + w # Operator methods can provide struct-owned arithmetic +let left_scaled = 2 * v # Reverse operator methods handle lhs-owned scalar syntax + +assert_eq(v.dot(w) -> m², 25 m²) # Method generic + unit conversion +assert_eq(combined.x -> m, 6 m) +assert_eq(left_scaled.y -> m, 8 m) + +let only_x = Vec2 { x: 5 m } # Omitted fields with defaults are auto-filled + +assert_eq(only_x.x, 5 m) +assert_eq(only_x.y, 0 m) + +let distances = [3 m, 4 m, 5 m] # Builtin lists support indexing +assert_eq(distances[1], 4 m) + +struct Grid2 { # Struct methods can define custom indexing syntax + a: Scalar, + b: Scalar, + c: Scalar, + d: Scalar, + + @index + fn get(self, row: Scalar, col: Scalar) -> Scalar = + if row == 0 then + if col == 0 then self.a else self.b + else + if col == 0 then self.c else self.d +} + +let grid = Grid2 { a: 1, b: 2, c: 3, d: 4 } +assert_eq(grid[1, 0], 3) 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/examples/tests/numerics.nbt b/examples/tests/numerics.nbt index e9c68f60..f320f10a 100644 --- a/examples/tests/numerics.nbt +++ b/examples/tests/numerics.nbt @@ -36,8 +36,8 @@ assert_eq(velocity(2.0 s), 2.0 s × g0, 1e-3 m/s) # Differential equations -let t_min = 0 s -let t_max = 1 s +let t_min: Time = 0 s +let t_max: Time = 1 s let n_points = 1_000 let μ = 0.7 / s @@ -51,7 +51,7 @@ fn ode(t, x) = μ x let result = dsolve_runge_kutta(ode, t_min, t_max, x0, n_points) -fn numerical_solution(t) = element_at(idx, result.ys) +fn numerical_solution(t: Time) = element_at(idx, result.ys) where t_range = t_max - t_min and idx = floor((t - t_min) / t_range * (n_points - 1)) diff --git a/numbat/examples/inspect.rs b/numbat/examples/inspect.rs index c02e4416..be2b79f0 100644 --- a/numbat/examples/inspect.rs +++ b/numbat/examples/inspect.rs @@ -233,7 +233,7 @@ fn replace_equation_delimiters(text_in: &str) -> CompactString { fn prepare_context(module_path: &Path) -> Context { let mut importer = FileSystemImporter::default(); - importer.add_path(module_path.to_path_buf()); + importer.add_path(module_path); Context::new(importer) } diff --git a/numbat/src/ast.rs b/numbat/src/ast.rs index 159b5a13..b077499d 100644 --- a/numbat/src/ast.rs +++ b/numbat/src/ast.rs @@ -122,6 +122,18 @@ 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, + }, + IndexCall { + receiver: Box>, + args: Vec>, + full_span: Span, + }, List(Span, Vec>), } @@ -156,6 +168,8 @@ 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::IndexCall { full_span, .. } => *full_span, Expression::List(span, _) => *span, Expression::TypedHole(span) => *span, } @@ -495,7 +509,8 @@ pub enum Statement<'a> { struct_name_span: Span, struct_name: &'a str, type_parameters: Vec<(Span, &'a str, Option)>, - fields: Vec<(Span, &'a str, TypeAnnotation)>, + fields: Vec<(Span, &'a str, TypeAnnotation, Option>)>, + methods: Vec>, }, } @@ -554,11 +569,18 @@ impl Statement<'_> { Statement::DefineStruct { struct_name_span, fields, + methods, .. } => { let mut span = *struct_name_span; - if let Some((last_span, _, annotation)) = fields.last() { + if let Some((last_span, _, annotation, default_expr)) = fields.last() { span = span.extend(last_span).extend(&annotation.full_span()); + if let Some(default_expr) = default_expr { + span = span.extend(&default_expr.full_span()); + } + } + if let Some(method) = methods.last() { + span = span.extend(&method.full_span()); } span } @@ -571,6 +593,34 @@ pub trait ReplaceSpans { fn replace_spans(&self) -> Self; } +#[cfg(test)] +impl ReplaceSpans for Decorator<'_> { + fn replace_spans(&self) -> Self { + match self { + Decorator::MetricPrefixes => Decorator::MetricPrefixes, + Decorator::BinaryPrefixes => Decorator::BinaryPrefixes, + Decorator::Abbreviation => Decorator::Abbreviation, + Decorator::Aliases(aliases) => Decorator::Aliases( + aliases + .iter() + .map(|(name, accepts_prefix, _)| (*name, *accepts_prefix, Span::dummy())) + .collect(), + ), + Decorator::Url(url) => Decorator::Url(url.clone()), + Decorator::Name(name) => Decorator::Name(name.clone()), + Decorator::Description(description) => Decorator::Description(description.clone()), + Decorator::Example(code, description) => { + Decorator::Example(code.clone(), description.clone()) + } + Decorator::BinaryOperator { operator, reverse } => Decorator::BinaryOperator { + operator: *operator, + reverse: *reverse, + }, + Decorator::Index => Decorator::Index, + } + } +} + #[cfg(test)] impl ReplaceSpans for TypeAnnotation { fn replace_spans(&self) -> Self { @@ -716,6 +766,23 @@ 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::IndexCall { receiver, args, .. } => Expression::IndexCall { + receiver: Box::new(receiver.replace_spans()), + 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(), @@ -733,7 +800,11 @@ impl ReplaceSpans for DefineVariable<'_> { identifier: self.identifier, expr: self.expr.replace_spans(), type_annotation: self.type_annotation.as_ref().map(|t| t.replace_spans()), - decorators: self.decorators.clone(), + decorators: self + .decorators + .iter() + .map(Decorator::replace_spans) + .collect(), } } } @@ -779,7 +850,7 @@ impl ReplaceSpans for Statement<'_> { .map(DefineVariable::replace_spans) .collect(), return_type_annotation: return_type_annotation.as_ref().map(|t| t.replace_spans()), - decorators: decorators.clone(), + decorators: decorators.iter().map(Decorator::replace_spans).collect(), }, Statement::DefineDimension(_, name, dexprs) => Statement::DefineDimension( Span::dummy(), @@ -805,7 +876,7 @@ impl ReplaceSpans for Statement<'_> { expr: expr.replace_spans(), type_annotation_span: type_annotation_span.map(|_| Span::dummy()), type_annotation: type_annotation.as_ref().map(|t| t.replace_spans()), - decorators: decorators.clone(), + decorators: decorators.iter().map(Decorator::replace_spans).collect(), }, Statement::ProcedureCall(_, proc, args) => Statement::ProcedureCall( Span::dummy(), @@ -819,6 +890,7 @@ impl ReplaceSpans for Statement<'_> { struct_name, type_parameters, fields, + methods, .. } => Statement::DefineStruct { struct_name_span: Span::dummy(), @@ -829,8 +901,16 @@ impl ReplaceSpans for Statement<'_> { .collect(), fields: fields .iter() - .map(|(_span, name, type_)| (Span::dummy(), *name, type_.replace_spans())) + .map(|(_span, name, type_, default_expr)| { + ( + Span::dummy(), + *name, + type_.replace_spans(), + default_expr.as_ref().map(Expression::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..517226b6 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)] @@ -215,6 +215,206 @@ impl BytecodeInterpreter { }), ); } + Expression::DeferredBinaryOperator { + op_span, + op: operator, + lhs, + rhs, + type_scheme, + } => { + let lhs_type = lhs.get_type_scheme().to_concrete_type(); + let rhs_type = rhs.get_type_scheme().to_concrete_type(); + let output_type = type_scheme.to_concrete_type(); + + match (&lhs_type, &rhs_type) { + (Type::DateTime, Type::DateTime) => { + self.compile_expression(lhs); + self.compile_expression(rhs); + + let second_idx = self.unit_name_to_constant_index.get("second"); + self.vm.add_op1( + Op::LoadConstant, + *second_idx.unwrap(), + op_span.unwrap_or_else(|| { + crate::span::Span::in_between(lhs.full_span(), rhs.full_span()) + }), + ); + self.vm.add_op( + Op::DiffDateTime, + op_span.unwrap_or_else(|| { + crate::span::Span::in_between(lhs.full_span(), rhs.full_span()) + }), + ); + } + (Type::DateTime, Type::Dimension(dtype)) if dtype.is_time_dimension() => { + self.compile_expression(lhs); + self.compile_expression(rhs); + + let op = match operator { + BinaryOperator::Add => Op::AddToDateTime, + BinaryOperator::Sub => Op::SubFromDateTime, + _ => unreachable!(), + }; + + self.vm.add_op( + op, + op_span.unwrap_or_else(|| { + crate::span::Span::in_between(lhs.full_span(), rhs.full_span()) + }), + ); + } + _ if matches!(lhs_type, Type::Struct(_)) + || matches!(rhs_type, Type::Struct(_)) => + { + let canonical_method_name = match (operator, false) { + (BinaryOperator::Add, false) => Some("add"), + (BinaryOperator::Sub, false) => Some("sub"), + (BinaryOperator::Mul, false) => Some("mul"), + (BinaryOperator::Div, false) => Some("div"), + _ => None, + }; + let reverse_canonical_method_name = match (operator, true) { + (BinaryOperator::Add, true) => Some("radd"), + (BinaryOperator::Sub, true) => Some("rsub"), + (BinaryOperator::Mul, true) => Some("rmul"), + (BinaryOperator::Div, true) => Some("rdiv"), + _ => None, + }; + + let lhs_match = + match &lhs_type { + Type::Struct(struct_info) => struct_info.methods.iter().find_map( + |(method_name, method_info)| { + if let Some(operator_impl) = + method_info.operator_impl.as_ref() + && operator_impl.operator == *operator + && !operator_impl.reverse + && operator_impl + .rhs_type + .equals_ignoring_struct_methods(&rhs_type) + && operator_impl + .output_type + .equals_ignoring_struct_methods(&output_type) + { + return Some((false, &struct_info.name, method_name)); + } + + canonical_method_name + .filter(|canonical| method_name.as_str() == *canonical) + .map(|_| (false, &struct_info.name, method_name)) + }, + ), + _ => None, + }; + + let rhs_match = match &rhs_type { + Type::Struct(rhs_struct_info) => rhs_struct_info + .methods + .iter() + .find_map(|(method_name, method_info)| { + if let Some(operator_impl) = method_info.operator_impl.as_ref() + && operator_impl.operator == *operator + && operator_impl.reverse + && operator_impl + .rhs_type + .equals_ignoring_struct_methods(&lhs_type) + && operator_impl + .output_type + .equals_ignoring_struct_methods(&output_type) + { + return Some(( + true, + rhs_struct_info.name.clone(), + method_name.clone(), + )); + } + + reverse_canonical_method_name + .filter(|canonical| method_name.as_str() == *canonical) + .map(|_| { + ( + true, + rhs_struct_info.name.clone(), + method_name.clone(), + ) + }) + }), + _ => None, + }; + + let (reverse, owner, method_name) = lhs_match + .map(|(reverse, owner, method_name)| { + (reverse, owner.clone(), method_name.clone()) + }) + .or(rhs_match) + .expect("deferred operator must resolve to a struct method"); + + if reverse { + self.compile_expression(rhs); + self.compile_expression(lhs); + } else { + self.compile_expression(lhs); + self.compile_expression(rhs); + } + + let arg_count = 2; + let method_callable = self + .vm + .get_method_callable(&owner, &method_name) + .expect("operator method must be registered before call sites"); + + match method_callable { + MethodCallable::Normal(idx) => { + self.vm.add_op2( + Op::Call, + idx, + arg_count, + lhs.full_span().extend(&rhs.full_span()), + ); + } + MethodCallable::Foreign(idx) => { + let call_args_idx = self.vm.add_ffi_call_args(FfiCallArgs { + args: vec![ + FfiCallArg { + span: lhs.full_span(), + type_: lhs.get_type_scheme(), + }, + FfiCallArg { + span: rhs.full_span(), + type_: rhs.get_type_scheme(), + }, + ], + return_type: Some(type_scheme.clone()), + }); + self.vm.add_op3( + Op::FFICallFunction, + idx, + arg_count, + call_args_idx, + lhs.full_span().extend(&rhs.full_span()), + ); + } + } + } + _ => { + self.compile_expression(lhs); + self.compile_expression(rhs); + let op = match operator { + BinaryOperator::Add => Op::Add, + BinaryOperator::Sub => Op::Subtract, + BinaryOperator::Mul => Op::Multiply, + BinaryOperator::Div => Op::Divide, + _ => unreachable!("unsupported deferred binary operator"), + }; + self.vm.add_op( + op, + op_span.unwrap_or_else(|| { + crate::span::Span::in_between(lhs.full_span(), rhs.full_span()) + }), + ); + } + } + } Expression::FunctionCall { full_span, name, @@ -227,6 +427,12 @@ impl BytecodeInterpreter { self.compile_expression(arg); } + if self.vm.get_ffi_callable_idx(name).is_none() + && crate::ffi::functions().contains_key(name) + { + self.vm.add_foreign_function(name, args.len()..=args.len()); + } + if let Some(idx) = self.vm.get_ffi_callable_idx(name) { // TODO: check overflow: let call_args = FfiCallArgs { @@ -393,6 +599,66 @@ 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::IndexCall { .. } => { + unreachable!("index expressions must be resolved before bytecode compilation") + } Expression::TypedHole(_, _) => { unreachable!("Typed holes cause type inference errors") } @@ -453,8 +719,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 +742,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 +973,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 +1067,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/decorator.rs b/numbat/src/decorator.rs index a9aff3a4..bf4b6dbe 100644 --- a/numbat/src/decorator.rs +++ b/numbat/src/decorator.rs @@ -1,8 +1,8 @@ use compact_str::CompactString; -use crate::{prefix_parser::AcceptsPrefix, span::Span, unit::CanonicalName}; +use crate::{ast::BinaryOperator, prefix_parser::AcceptsPrefix, span::Span, unit::CanonicalName}; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq)] pub enum Decorator<'a> { MetricPrefixes, BinaryPrefixes, @@ -12,6 +12,11 @@ pub enum Decorator<'a> { Name(CompactString), Description(CompactString), Example(CompactString, Option), + BinaryOperator { + operator: BinaryOperator, + reverse: bool, + }, + Index, } /// Get an iterator of data computed from a name and/or its alias's `AcceptsPrefix` and @@ -172,3 +177,30 @@ pub fn contains_abbreviation(decorators: &[Decorator]) -> bool { false } + +pub fn contains_binary_operator_decorators(decorators: &[Decorator]) -> bool { + decorators + .iter() + .any(|decorator| matches!(decorator, Decorator::BinaryOperator { .. })) +} + +pub fn contains_index_decorator(decorators: &[Decorator]) -> bool { + decorators + .iter() + .any(|decorator| matches!(decorator, Decorator::Index)) +} + +pub fn contains_method_only_decorators(decorators: &[Decorator]) -> bool { + contains_binary_operator_decorators(decorators) || contains_index_decorator(decorators) +} + +pub fn binary_operator<'a>(decorators: &'a [Decorator<'a>]) -> Option<(BinaryOperator, bool)> { + decorators.iter().find_map(|decorator| match decorator { + Decorator::BinaryOperator { operator, reverse } => Some((*operator, *reverse)), + _ => None, + }) +} + +pub fn index_decorator(decorators: &[Decorator<'_>]) -> bool { + contains_index_decorator(decorators) +} diff --git a/numbat/src/diagnostic.rs b/numbat/src/diagnostic.rs index fa04a4c7..0e22c456 100644 --- a/numbat/src/diagnostic.rs +++ b/numbat/src/diagnostic.rs @@ -432,6 +432,24 @@ 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::InvalidOperatorMethodSignature(span, _) + | TypeCheckError::InvalidIndexMethodSignature(span, _) + | TypeCheckError::IndexCallOnNonStructType(span, _) + | TypeCheckError::InvalidListIndexArity(span, _) + | TypeCheckError::InvalidListIndexType(span, _) + | TypeCheckError::IndexMethodNotFound(span, _, _) + | TypeCheckError::AmbiguousIndexOverload(span, _, _) + | TypeCheckError::AmbiguousOperatorOverload(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 +484,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/functions.rs b/numbat/src/ffi/functions.rs index f295b752..a4f0f60c 100644 --- a/numbat/src/ffi/functions.rs +++ b/numbat/src/ffi/functions.rs @@ -89,6 +89,7 @@ pub(crate) fn functions() -> &'static HashMap<&'static str, ForeignFunction> { insert_function!(tail, 1..=1); insert_function!(cons, 2..=2); insert_function!(cons_end, 2..=2); + insert_function!("_list_at", list_at, 2..=2); // Strings insert_function!(str_length, 1..=1); diff --git a/numbat/src/ffi/lists.rs b/numbat/src/ffi/lists.rs index 159b9daa..0359b83e 100644 --- a/numbat/src/ffi/lists.rs +++ b/numbat/src/ffi/lists.rs @@ -63,3 +63,21 @@ pub fn cons_end( return_list!(list) } + +pub fn list_at( + _ctx: &mut FfiContext, + mut args: Args, + _return_type: &TypeScheme, +) -> Result> { + let list = list_arg!(args); + let index = scalar_arg!(args).to_f64(); + + if !index.is_finite() || index < 0.0 || index.fract() != 0.0 { + return Err(Box::new(RuntimeErrorKind::InvalidListIndex)); + } + + let index = index as usize; + list.get(index) + .cloned() + .ok_or_else(|| Box::new(RuntimeErrorKind::ListIndexOutOfBounds(index, list.len()))) +} 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/interpreter/mod.rs b/numbat/src/interpreter/mod.rs index 52c65055..7af37037 100644 --- a/numbat/src/interpreter/mod.rs +++ b/numbat/src/interpreter/mod.rs @@ -81,6 +81,12 @@ pub enum RuntimeErrorKind { #[error("Empty list")] EmptyList, + #[error("List index must be a non-negative integer")] + InvalidListIndex, + + #[error("List index {0} is out of bounds for a list of length {1}")] + ListIndexOutOfBounds(usize, usize), + #[error("Could not write to file: {0:?}")] FileWrite(std::path::PathBuf), diff --git a/numbat/src/list.rs b/numbat/src/list.rs index c0eec205..1cd6622f 100644 --- a/numbat/src/list.rs +++ b/numbat/src/list.rs @@ -94,6 +94,15 @@ impl NumbatList { let (start, end) = self.view.map_or((0, self.alloc.len()), |view| view); self.alloc.iter().skip(start).take(end - start) } + + pub fn get(&self, index: usize) -> Option<&T> { + let (start, end) = self.view.map_or((0, self.alloc.len()), |view| view); + if index >= end - start { + None + } else { + self.alloc.get(start + index) + } + } } impl NumbatList { diff --git a/numbat/src/parser.rs b/numbat/src/parser.rs index 7c1c7e47..7af6cfc6 100644 --- a/numbat/src/parser.rs +++ b/numbat/src/parser.rs @@ -16,7 +16,7 @@ //! module_import ::= "use" ident ( "::" ident) * //! procedure_call ::= ( "print" | "assert" | "assert_eq" | "type" ) "(" arguments? ")" //! -//! decorator ::= "@" ( "metric_prefixes" | "binary_prefixes" | ( "aliases(" list_of_aliases ")" ) ) +//! decorator ::= "@" ( "metric_prefixes" | "binary_prefixes" | "add" | "sub" | "mul" | "div" | "radd" | "rsub" | "rmul" | "rdiv" | "index" | ( "aliases(" list_of_aliases ")" ) ) //! //! type_annotation ::= "Bool" | "String" | "List<" type ">" | dimension_expr //! dimension_expr ::= dim_factor @@ -41,7 +41,7 @@ //! power ::= factorial ( "^" "-" ? power ) ? //! factorial ::= unicode_power "!" * //! unicode_power ::= call ( "⁻" ? ( "¹" | "²" | "³" | "⁴" | "⁵" | "⁶" | "⁷" | "⁸" | "⁹" ) ) ? -//! call ::= primary ( ( "(" arguments? ")" ) | "." identifier ) * +//! call ::= primary ( ( "(" arguments? ")" ) | "[" 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 "," ? ) ? "}" @@ -211,6 +211,12 @@ pub enum ParseErrorKind { #[error("Example decorators can only be used on functions.")] ExampleUsedOnUnsuitableKind, + #[error("Operator and index decorators can only be used on struct instance methods.")] + MethodDecoratorUsedOutsideStructMethod, + + #[error("Expected comma in decorator arguments")] + ExpectedCommaInDecorator, + #[error("Numerical overflow in dimension exponent")] OverflowInDimensionExponent, @@ -235,6 +241,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, @@ -280,6 +295,8 @@ static PROCEDURES: &[TokenKind] = &[ struct Parser<'a> { current: usize, decorator_stack: Vec>, + parsing_struct_method: bool, + expression_group_depth: usize, /// When we split `>=` into `>` and `=`, we store the span of the `=` here. /// This is used for parsing something like `let v: Vec= ...`, where /// the `>=` is being tokenized as a TokenKind::GreaterOrEqual. @@ -291,6 +308,8 @@ impl<'a> Parser<'a> { Parser { current: 0, decorator_stack: vec![], + parsing_struct_method: false, + expression_group_depth: 0, pending_equals: None, } } @@ -495,7 +514,7 @@ impl<'a> Parser<'a> { self.parse_variable(tokens, true) .map(Statement::DefineVariable) } else if let Some(fn_token) = self.match_exact(tokens, TokenKind::Fn) { - self.parse_function_declaration(tokens, fn_token.span) + self.parse_function_declaration(tokens, fn_token.span, false) } else if self.match_exact(tokens, TokenKind::Dimension).is_some() { self.parse_dimension_declaration(tokens) } else if self.match_exact(tokens, TokenKind::At).is_some() { @@ -575,6 +594,7 @@ impl<'a> Parser<'a> { &mut self, tokens: &[Token<'a>], fn_keyword_span: Span, + allow_operator_decorators: bool, ) -> Result> { if let Some(fn_name) = self.match_exact(tokens, TokenKind::Identifier) { let function_name_span = self.last(tokens).unwrap().span; @@ -681,6 +701,16 @@ impl<'a> Parser<'a> { }); } + if !allow_operator_decorators + && !self.parsing_struct_method + && decorator::contains_method_only_decorators(&self.decorator_stack) + { + return Err(ParseError { + kind: ParseErrorKind::MethodDecoratorUsedOutsideStructMethod, + span: self.peek(tokens).span, + }); + } + let mut decorators = vec![]; std::mem::swap(&mut decorators, &mut self.decorator_stack); @@ -838,6 +868,22 @@ impl<'a> Parser<'a> { }); } } + "add" | "sub" | "mul" | "div" | "radd" | "rsub" | "rmul" | "rdiv" => { + let (operator, reverse) = match decorator.lexeme { + "add" => (BinaryOperator::Add, false), + "sub" => (BinaryOperator::Sub, false), + "mul" => (BinaryOperator::Mul, false), + "div" => (BinaryOperator::Div, false), + "radd" => (BinaryOperator::Add, true), + "rsub" => (BinaryOperator::Sub, true), + "rmul" => (BinaryOperator::Mul, true), + "rdiv" => (BinaryOperator::Div, true), + _ => unreachable!(), + }; + + Decorator::BinaryOperator { operator, reverse } + } + "index" => Decorator::Index, _ => { return Err(ParseError { kind: ParseErrorKind::UnknownDecorator, @@ -966,16 +1012,48 @@ 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, true)?); + self.skip_empty_lines(tokens); + continue; + } + + if self.match_exact(tokens, TokenKind::At).is_some() { + parsing_methods = true; + self.parsing_struct_method = true; + let method = self.parse_decorators(tokens)?; + self.parsing_struct_method = false; + let Statement::DefineFunction { .. } = method else { + return Err(ParseError { + kind: ParseErrorKind::ExpectedFieldOrMethodInStruct, + span: self.peek(tokens).span, + }); + }; + methods.push(method); + 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() { @@ -990,25 +1068,39 @@ impl<'a> Parser<'a> { self.skip_empty_lines(tokens); + let default_expr = if self.match_exact(tokens, TokenKind::Equal).is_some() { + self.skip_empty_lines(tokens); + Some(self.expression(tokens)?) + } else { + None + }; + + self.skip_empty_lines(tokens); + let has_comma = self.match_exact(tokens, TokenKind::Comma).is_some(); 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, }); } - fields.push((field_name.span, field_name.lexeme, attr_type)); + fields.push((field_name.span, field_name.lexeme, attr_type, default_expr)); } + self.match_exact(tokens, TokenKind::RightCurly); Ok(Statement::DefineStruct { struct_name_span: name_span, struct_name: name, type_parameters, fields, + methods, }) } @@ -1046,8 +1138,18 @@ impl<'a> Parser<'a> { next_parser: impl Fn(&mut Self) -> Result>, ) -> Result> { let mut expr = next_parser(self)?; - while let Some(matched) = self.match_any(tokens, op_symbol) { + loop { + if self.expression_group_depth > 0 { + self.skip_empty_lines(tokens); + } + + let Some(matched) = self.match_any(tokens, op_symbol) else { + break; + }; let span_op = Some(self.last(tokens).unwrap().span); + if self.expression_group_depth > 0 { + self.skip_empty_lines(tokens); + } let rhs = next_parser(self)?; expr = Expression::BinaryOperator { @@ -1419,17 +1521,62 @@ 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 if self.match_exact(tokens, TokenKind::LeftBracket).is_some() { + let args = self.index_arguments(tokens)?; + let full_span = expr.full_span().extend(&self.last(tokens).unwrap().span); + + expr = Expression::IndexCall { + receiver: Box::new(expr), + args, + full_span, + }; } else { return Ok(expr); } @@ -1473,6 +1620,35 @@ impl<'a> Parser<'a> { Ok(args) } + fn index_arguments(&mut self, tokens: &[Token<'a>]) -> Result>> { + self.skip_empty_lines(tokens); + if self.match_exact(tokens, TokenKind::RightBracket).is_some() { + return Ok(vec![]); + } + + let mut args = vec![self.expression(tokens)?]; + loop { + self.skip_empty_lines(tokens); + + if self.match_exact(tokens, TokenKind::Comma).is_some() { + self.skip_empty_lines(tokens); + if self.match_exact(tokens, TokenKind::RightBracket).is_some() { + break; + } + args.push(self.expression(tokens)?); + } else if self.match_exact(tokens, TokenKind::RightBracket).is_some() { + break; + } else { + return Err(ParseError { + kind: ParseErrorKind::ExpectedCommaOrRightBracketInList, + span: self.peek(tokens).span, + }); + } + } + + Ok(args) + } + fn primary(&mut self, tokens: &[Token<'a>]) -> Result> { // This function needs to be kept in sync with `next_token_could_start_power_expression` below. @@ -1658,7 +1834,11 @@ impl<'a> Parser<'a> { Ok(Expression::String(span_full_string, parts)) } else if self.match_exact(tokens, TokenKind::LeftParen).is_some() { + self.expression_group_depth += 1; + self.skip_empty_lines(tokens); let inner = self.expression(tokens)?; + self.skip_empty_lines(tokens); + self.expression_group_depth -= 1; if self.match_exact(tokens, TokenKind::RightParen).is_none() { return Err(ParseError::new( @@ -2183,12 +2363,9 @@ mod tests { use std::fmt::Write; use super::*; - use crate::{ - ast::{ - ReplaceSpans, binop, boolean, conditional, factorial, identifier, list, logical_neg, - negate, scalar, struct_, - }, - span::ByteIndex, + use crate::ast::{ + ReplaceSpans, binop, boolean, conditional, factorial, identifier, list, logical_neg, + negate, scalar, struct_, }; #[track_caller] @@ -2661,24 +2838,8 @@ mod tests { decorators: vec![ decorator::Decorator::Name("myvar".into()), decorator::Decorator::Aliases(vec![ - ( - "foo", - None, - Span { - start: ByteIndex(24), - end: ByteIndex(27), - code_source_id: 0, - }, - ), - ( - "bar", - None, - Span { - start: ByteIndex(29), - end: ByteIndex(32), - code_source_id: 0, - }, - ), + ("foo", None, Span::dummy()), + ("bar", None, Span::dummy()), ]), ], }), @@ -3250,6 +3411,359 @@ 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![], + )), + None, + )], + 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![], + )), + None, + )], + 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( + &[ + "struct Vec2 { x: Length, y: Length, @add fn add(self, rhs: Self) -> Self = Vec2 { x: self.x + rhs.x, y: self.y + rhs.y } }", + ], + Statement::DefineStruct { + struct_name_span: Span::dummy(), + struct_name: "Vec2", + type_parameters: vec![], + fields: vec![ + ( + Span::dummy(), + "x", + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + Span::dummy(), + "Length".into(), + vec![], + )), + None, + ), + ( + Span::dummy(), + "y", + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + Span::dummy(), + "Length".into(), + vec![], + )), + None, + ), + ], + methods: vec![Statement::DefineFunction { + fn_keyword_span: Span::dummy(), + function_name_span: Span::dummy(), + function_name: "add", + type_parameters: vec![], + parameters: vec![ + (Span::dummy(), "self", None), + ( + Span::dummy(), + "rhs", + Some(TypeAnnotation::TypeExpression( + TypeExpression::TypeIdentifier( + Span::dummy(), + "Self".into(), + vec![], + ), + )), + ), + ], + body: Some(struct_! { + Vec2, + x: binop!( + Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("self")), + field_name: "x", + }, + Add, + Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("rhs")), + field_name: "x", + } + ), + y: binop!( + Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("self")), + field_name: "y", + }, + Add, + Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("rhs")), + field_name: "y", + } + ) + }), + local_variables: vec![], + return_type_annotation: Some(TypeAnnotation::TypeExpression( + TypeExpression::TypeIdentifier(Span::dummy(), "Self".into(), vec![]), + )), + decorators: vec![decorator::Decorator::BinaryOperator { + operator: BinaryOperator::Add, + reverse: false, + }], + }], + }, + ); + + parse_as( + &[ + "struct Shift { amount: Length, @radd fn add_to_vec(self, lhs: Vec2) -> Vec2 = Vec2 { x: lhs.x + self.amount, y: lhs.y + self.amount } }", + ], + Statement::DefineStruct { + struct_name_span: Span::dummy(), + struct_name: "Shift", + type_parameters: vec![], + fields: vec![( + Span::dummy(), + "amount", + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + Span::dummy(), + "Length".into(), + vec![], + )), + None, + )], + methods: vec![Statement::DefineFunction { + fn_keyword_span: Span::dummy(), + function_name_span: Span::dummy(), + function_name: "add_to_vec", + type_parameters: vec![], + parameters: vec![ + (Span::dummy(), "self", None), + ( + Span::dummy(), + "lhs", + Some(TypeAnnotation::TypeExpression( + TypeExpression::TypeIdentifier( + Span::dummy(), + "Vec2".into(), + vec![], + ), + )), + ), + ], + body: Some(struct_! { + Vec2, + x: binop!( + Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("lhs")), + field_name: "x", + }, + Add, + Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("self")), + field_name: "amount", + } + ), + y: binop!( + Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("lhs")), + field_name: "y", + }, + Add, + Expression::AccessField { + full_span: Span::dummy(), + ident_span: Span::dummy(), + expr: Box::new(identifier!("self")), + field_name: "amount", + } + ) + }), + local_variables: vec![], + return_type_annotation: Some(TypeAnnotation::TypeExpression( + TypeExpression::TypeIdentifier(Span::dummy(), "Vec2".into(), vec![]), + )), + decorators: vec![decorator::Decorator::BinaryOperator { + operator: BinaryOperator::Add, + reverse: true, + }], + }], + }, + ); + + 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( + &["grid[1, 2]"], + Expression::IndexCall { + receiver: Box::new(identifier!("grid")), + args: vec![scalar!(1.0), scalar!(2.0)], + full_span: Span::dummy(), + }, + ); + + parse_as_expression( + &["(\n 1\n + 2\n)"], + Expression::BinaryOperator { + op: BinaryOperator::Add, + lhs: Box::new(scalar!(1.0)), + rhs: Box::new(scalar!(2.0)), + span_op: Some(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( + &["@add fn add(x: Scalar, y: Scalar) = x + y"], + ParseErrorKind::MethodDecoratorUsedOutsideStructMethod, + ); + + should_fail_with( + &["@radd fn add(x: Scalar, y: Scalar) = x + y"], + ParseErrorKind::MethodDecoratorUsedOutsideStructMethod, + ); + + should_fail_with( + &["@index fn get(x: Scalar, y: Scalar) = x + y"], + ParseErrorKind::MethodDecoratorUsedOutsideStructMethod, + ); + + 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( @@ -3640,6 +4154,7 @@ mod tests { CompactString::const_new("Scalar"), vec![], )), + None, ), ( Span::dummy(), @@ -3649,8 +4164,10 @@ mod tests { CompactString::const_new("Scalar"), vec![], )), + None, ), ], + methods: vec![], }, ); @@ -3672,6 +4189,7 @@ mod tests { CompactString::const_new("T"), vec![], )), + None, ), ( Span::dummy(), @@ -3681,9 +4199,48 @@ mod tests { CompactString::const_new("D"), vec![], )), + None, + ), + ( + Span::dummy(), + "name", + TypeAnnotation::String(Span::dummy()), + None, + ), + ], + methods: vec![], + }, + ); + + parse_as( + &["struct Foo { x: Scalar = 2, y: Scalar }"], + Statement::DefineStruct { + struct_name_span: Span::dummy(), + struct_name: "Foo", + type_parameters: vec![], + fields: vec![ + ( + Span::dummy(), + "x", + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + Span::dummy(), + CompactString::const_new("Scalar"), + vec![], + )), + Some(Expression::Scalar(Span::dummy(), Number::from_f64(2.0))), + ), + ( + Span::dummy(), + "y", + TypeAnnotation::TypeExpression(TypeExpression::TypeIdentifier( + Span::dummy(), + CompactString::const_new("Scalar"), + vec![], + )), + None, ), - (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..8f387866 100644 --- a/numbat/src/prefix_transformer.rs +++ b/numbat/src/prefix_transformer.rs @@ -156,6 +156,18 @@ 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::IndexCall { 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 +243,54 @@ impl Transformer { fn transform_statement(&mut self, statement: &mut Statement) -> Result<()> { match statement { - Statement::DefineStruct { .. } | Statement::ModuleImport(_, _) => {} + Statement::DefineStruct { + fields, methods, .. + } => { + for (_, _, _, default_expr) in fields { + if let Some(default_expr) = default_expr { + self.transform_expression(default_expr); + } + } + + 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..6f298c44 100644 --- a/numbat/src/traversal.rs +++ b/numbat/src/traversal.rs @@ -43,6 +43,16 @@ impl ForAllTypeSchemes for Expression<'_> { rhs.for_all_type_schemes(f); f(type_scheme); } + Expression::DeferredBinaryOperator { + lhs, + rhs, + type_scheme, + .. + } => { + lhs.for_all_type_schemes(f); + rhs.for_all_type_schemes(f); + f(type_scheme); + } Expression::FunctionCall { args, type_scheme, .. } => { @@ -95,6 +105,18 @@ impl ForAllTypeSchemes for Expression<'_> { f(struct_type); f(field_type); } + Expression::IndexCall { + receiver, + args, + type_scheme, + .. + } => { + receiver.for_all_type_schemes(f); + for arg in args { + arg.for_all_type_schemes(f); + } + f(type_scheme); + } Expression::List { elements, type_scheme, @@ -105,6 +127,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 +161,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 +210,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); @@ -207,6 +252,10 @@ impl ForAllExpressions for Expression<'_> { lhs.for_all_expressions(f); rhs.for_all_expressions(f); } + Expression::DeferredBinaryOperator { lhs, rhs, .. } => { + lhs.for_all_expressions(f); + rhs.for_all_expressions(f); + } Expression::FunctionCall { args, .. } => { for arg in args { arg.for_all_expressions(f); @@ -238,11 +287,23 @@ impl ForAllExpressions for Expression<'_> { Expression::AccessField { expr, .. } => { expr.for_all_expressions(f); } + Expression::IndexCall { receiver, args, .. } => { + receiver.for_all_expressions(f); + for arg in args { + arg.for_all_expressions(f); + } + } Expression::List { elements, .. } => { for element in elements { 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..f5d999b1 100644 --- a/numbat/src/typechecker/const_evaluation.rs +++ b/numbat/src/typechecker/const_evaluation.rs @@ -105,9 +105,12 @@ pub fn evaluate_const_expr(expr: &typed_ast::Expression) -> Result { typed_ast::Expression::String(_, _) => "String", typed_ast::Expression::Condition { .. } => "Conditional", typed_ast::Expression::BinaryOperatorForDate { .. } => "binary operator for datetimes", + typed_ast::Expression::DeferredBinaryOperator { .. } => "deferred binary operator", typed_ast::Expression::InstantiateStruct { .. } => "instantiate struct", typed_ast::Expression::AccessField { .. } => "access field of struct", + typed_ast::Expression::IndexCall { .. } => "index expression", typed_ast::Expression::List { .. } => "lists", + typed_ast::Expression::MethodCall { .. } => "method call", typed_ast::Expression::TypedHole(_, _) => "typed hole", }; diff --git a/numbat/src/typechecker/constraints.rs b/numbat/src/typechecker/constraints.rs index 7c87fad0..65531a2a 100644 --- a/numbat/src/typechecker/constraints.rs +++ b/numbat/src/typechecker/constraints.rs @@ -4,7 +4,7 @@ use compact_str::{CompactString, format_compact}; use super::substitutions::{ApplySubstitution, Substitution, SubstitutionError}; use crate::type_variable::TypeVariable; -use crate::typed_ast::{DType, DTypeFactor, StructKind, Type}; +use crate::typed_ast::{BinaryOperator, DType, DTypeFactor, StructKind, Type}; #[derive(Debug, Clone, PartialEq, Eq)] pub enum ConstraintSolverError { @@ -82,7 +82,13 @@ impl ConstraintSet { self.constraints.clear(); } - pub fn solve(&mut self) -> Result<(Substitution, Vec), ConstraintSolverError> { + pub fn solve( + &mut self, + mut resolve_binary_operator: F, + ) -> Result<(Substitution, Vec), ConstraintSolverError> + where + F: FnMut(BinaryOperator, &Type, &Type, &Type) -> Option, + { let mut substitution = Substitution::empty(); let mut made_progress = true; @@ -94,7 +100,7 @@ impl ConstraintSet { if let Some(Satisfied { new_constraints, new_substitution, - }) = c.try_satisfy() + }) = c.try_satisfy(&mut resolve_binary_operator) { new_constraint_set.remove(i); new_constraint_set.constraints.extend(new_constraints); @@ -191,13 +197,14 @@ pub enum Constraint { IsDType(Type), EqualScalar(DType), HasField(Type, CompactString, Type), + HasBinaryOperator(BinaryOperator, Type, Type, Type), } impl Constraint { fn try_trivial_resolution(&self) -> TrivialResolution { match self { Constraint::Equal(t1, t2) if t1.is_closed() && t2.is_closed() => { - if t1 == t2 { + if t1.equals_ignoring_struct_methods(t2) { TrivialResolution::Satisfied } else { TrivialResolution::Violated @@ -228,11 +235,15 @@ impl Constraint { // Trivial resolution handling for structs is done directly in the type checker TrivialResolution::Unknown } + Constraint::HasBinaryOperator(_, _, _, _) => TrivialResolution::Unknown, } } /// Try to solve a constraint. Returns `None` if the constaint can not (yet) be solved. - fn try_satisfy(&self) -> Option { + fn try_satisfy(&self, resolve_binary_operator: &mut F) -> Option + where + F: FnMut(BinaryOperator, &Type, &Type, &Type) -> Option, + { match self { Constraint::Equal(t1, t2) if t1 == t2 => Some(Satisfied::trivially()), Constraint::Equal(Type::TVar(x), t) | Constraint::Equal(t, Type::TVar(x)) @@ -355,6 +366,9 @@ impl Constraint { } } Constraint::HasField(_, _, _) => None, + Constraint::HasBinaryOperator(op, lhs, rhs, output) => { + resolve_binary_operator(*op, lhs, rhs, output) + } } } @@ -368,6 +382,25 @@ impl Constraint { Constraint::HasField(struct_type, field_name, field_type) => { format_compact!("HasField({struct_type}, \"{field_name}\", {field_type})") } + Constraint::HasBinaryOperator(op, lhs, rhs, output) => { + let op = match op { + BinaryOperator::Add => "+", + BinaryOperator::Sub => "-", + BinaryOperator::Mul => "*", + BinaryOperator::Div => "/", + BinaryOperator::Power => "^", + BinaryOperator::ConvertTo => "->", + BinaryOperator::LessThan => "<", + BinaryOperator::GreaterThan => ">", + BinaryOperator::LessOrEqual => "<=", + BinaryOperator::GreaterOrEqual => ">=", + BinaryOperator::Equal => "==", + BinaryOperator::NotEqual => "!=", + BinaryOperator::LogicalAnd => "&&", + BinaryOperator::LogicalOr => "||", + }; + format_compact!("HasBinaryOperator({lhs}, {op}, {rhs}, {output})") + } } } @@ -396,6 +429,11 @@ impl ApplySubstitution for Constraint { struct_type.apply(substitution)?; field_type.apply(substitution)?; } + Constraint::HasBinaryOperator(_, lhs, rhs, output) => { + lhs.apply(substitution)?; + rhs.apply(substitution)?; + output.apply(substitution)?; + } } Ok(()) } diff --git a/numbat/src/typechecker/error.rs b/numbat/src/typechecker/error.rs index 2efd101a..f30e1bb2 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,12 +156,61 @@ 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( + "Operator decorator on method '{1}' requires an instance method with exactly one non-self parameter" + )] + InvalidOperatorMethodSignature(Span, String), + + #[error( + "Index decorator on method '{1}' requires an instance method with at least one non-self parameter" + )] + InvalidIndexMethodSignature(Span, String), + + #[error("Can not index value of type '{1}'")] + IndexCallOnNonStructType(Span, Type), + + #[error("List indexing expects exactly one argument, got {1}")] + InvalidListIndexArity(Span, usize), + + #[error("List indices must be scalars, got '{1}'")] + InvalidListIndexType(Span, Type), + + #[error("No matching index overload with {2} argument(s) exists on struct '{1}'")] + IndexMethodNotFound(Span, String, usize), + + #[error("Multiple index overloads matched on struct '{1}' with {2} argument(s)")] + AmbiguousIndexOverload(Span, String, usize), + + #[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)>), #[error("Incompatible types in list: expected '{1}', got '{3}' instead")] IncompatibleTypesInList(Span, Type, Span, Type), + #[error("Multiple operator overloads matched for '{1:?}' with operands '{2}' and '{3}'")] + AmbiguousOperatorOverload(Span, BinaryOperator, Type, Type), + #[error(transparent)] NameResolutionError(#[from] NameResolutionError), diff --git a/numbat/src/typechecker/mod.rs b/numbat/src/typechecker/mod.rs index 143bd1ba..f05b49ef 100644 --- a/numbat/src/typechecker/mod.rs +++ b/numbat/src/typechecker/mod.rs @@ -12,7 +12,7 @@ pub mod qualified_type; mod substitutions; pub mod type_scheme; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::ops::Deref; use std::sync::Arc; @@ -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,25 @@ 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, + operator_impl: Option, + index_impl: bool, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +struct MethodOperatorInfo { + operator: BinaryOperator, + reverse: bool, } struct ElaborationDefinitionArgs<'a, 'b> { @@ -205,6 +225,464 @@ struct ElaborationDefinitionArgs<'a, 'b> { } impl TypeChecker { + fn collect_struct_field_defaults<'a>( + statements: &[ast::Statement<'a>], + ) -> HashMap<&'a str, Vec<(&'a str, ast::Expression<'a>)>> { + let mut defaults = HashMap::new(); + + for statement in statements { + if let ast::Statement::DefineStruct { + struct_name, + fields, + .. + } = statement + { + let field_defaults = fields + .iter() + .filter_map(|(_, field_name, _, default_expr)| { + default_expr.clone().map(|expr| (*field_name, expr)) + }) + .collect::>(); + + if !field_defaults.is_empty() { + defaults.insert(*struct_name, field_defaults); + } + } + } + + defaults + } + + fn apply_struct_defaults_expression<'a>( + expr: &ast::Expression<'a>, + defaults: &HashMap<&'a str, Vec<(&'a str, ast::Expression<'a>)>>, + ) -> ast::Expression<'a> { + match expr { + ast::Expression::Scalar(..) + | ast::Expression::Identifier(..) + | ast::Expression::UnitIdentifier { .. } + | ast::Expression::TypedHole(..) + | ast::Expression::Boolean(..) => expr.clone(), + ast::Expression::UnaryOperator { op, expr, span_op } => { + ast::Expression::UnaryOperator { + op: *op, + expr: Box::new(Self::apply_struct_defaults_expression(expr, defaults)), + span_op: *span_op, + } + } + ast::Expression::BinaryOperator { + op, + lhs, + rhs, + span_op, + } => ast::Expression::BinaryOperator { + op: *op, + lhs: Box::new(Self::apply_struct_defaults_expression(lhs, defaults)), + rhs: Box::new(Self::apply_struct_defaults_expression(rhs, defaults)), + span_op: *span_op, + }, + ast::Expression::FunctionCall { + ident_span, + full_span, + callable, + args, + } => ast::Expression::FunctionCall { + ident_span: *ident_span, + full_span: *full_span, + callable: Box::new(Self::apply_struct_defaults_expression(callable, defaults)), + args: args + .iter() + .map(|arg| Self::apply_struct_defaults_expression(arg, defaults)) + .collect(), + }, + ast::Expression::String(span, parts) => ast::Expression::String( + *span, + parts + .iter() + .map(|part| match part { + ast::StringPart::Fixed(s) => ast::StringPart::Fixed(s.clone()), + ast::StringPart::Interpolation { + span, + expr, + format_specifiers, + } => ast::StringPart::Interpolation { + span: *span, + expr: Box::new(Self::apply_struct_defaults_expression(expr, defaults)), + format_specifiers: *format_specifiers, + }, + }) + .collect(), + ), + ast::Expression::Condition { + span, + condition, + then_expr, + else_expr, + } => ast::Expression::Condition { + span: *span, + condition: Box::new(Self::apply_struct_defaults_expression(condition, defaults)), + then_expr: Box::new(Self::apply_struct_defaults_expression(then_expr, defaults)), + else_expr: Box::new(Self::apply_struct_defaults_expression(else_expr, defaults)), + }, + ast::Expression::InstantiateStruct { + full_span, + ident_span, + name, + fields, + } => { + let mut transformed_fields = fields + .iter() + .map(|(span, field_name, value)| { + ( + *span, + *field_name, + Self::apply_struct_defaults_expression(value, defaults), + ) + }) + .collect::>(); + + if let Some(defaults_for_struct) = defaults.get(name) { + let provided_fields = transformed_fields + .iter() + .map(|(_, field_name, _)| *field_name) + .collect::>(); + + for (field_name, default_expr) in defaults_for_struct { + if !provided_fields.contains(field_name) { + transformed_fields.push(( + default_expr.full_span(), + *field_name, + Self::apply_struct_defaults_expression(default_expr, defaults), + )); + } + } + } + + ast::Expression::InstantiateStruct { + full_span: *full_span, + ident_span: *ident_span, + name, + fields: transformed_fields, + } + } + ast::Expression::AccessField { + full_span, + ident_span, + expr, + field_name, + } => ast::Expression::AccessField { + full_span: *full_span, + ident_span: *ident_span, + expr: Box::new(Self::apply_struct_defaults_expression(expr, defaults)), + field_name, + }, + ast::Expression::MethodCall { + receiver, + method_name_span, + method_name, + args, + full_span, + } => ast::Expression::MethodCall { + receiver: Box::new(Self::apply_struct_defaults_expression(receiver, defaults)), + method_name_span: *method_name_span, + method_name, + args: args + .iter() + .map(|arg| Self::apply_struct_defaults_expression(arg, defaults)) + .collect(), + full_span: *full_span, + }, + ast::Expression::IndexCall { + receiver, + args, + full_span, + } => ast::Expression::IndexCall { + receiver: Box::new(Self::apply_struct_defaults_expression(receiver, defaults)), + args: args + .iter() + .map(|arg| Self::apply_struct_defaults_expression(arg, defaults)) + .collect(), + full_span: *full_span, + }, + ast::Expression::List(span, elements) => ast::Expression::List( + *span, + elements + .iter() + .map(|e| Self::apply_struct_defaults_expression(e, defaults)) + .collect(), + ), + } + } + + fn apply_struct_defaults_statement<'a>( + statement: &ast::Statement<'a>, + defaults: &HashMap<&'a str, Vec<(&'a str, ast::Expression<'a>)>>, + ) -> ast::Statement<'a> { + match statement { + ast::Statement::Expression(expr) => { + ast::Statement::Expression(Self::apply_struct_defaults_expression(expr, defaults)) + } + ast::Statement::DefineVariable(variable) => { + ast::Statement::DefineVariable(ast::DefineVariable { + identifier_span: variable.identifier_span, + identifier: variable.identifier, + expr: Self::apply_struct_defaults_expression(&variable.expr, defaults), + type_annotation: variable.type_annotation.clone(), + decorators: variable.decorators.clone(), + }) + } + ast::Statement::DefineFunction { + fn_keyword_span, + function_name_span, + function_name, + type_parameters, + parameters, + body, + local_variables, + return_type_annotation, + decorators, + } => ast::Statement::DefineFunction { + fn_keyword_span: *fn_keyword_span, + function_name_span: *function_name_span, + function_name, + type_parameters: type_parameters.clone(), + parameters: parameters.clone(), + body: body + .as_ref() + .map(|expr| Self::apply_struct_defaults_expression(expr, defaults)), + local_variables: local_variables + .iter() + .map(|v| ast::DefineVariable { + identifier_span: v.identifier_span, + identifier: v.identifier, + expr: Self::apply_struct_defaults_expression(&v.expr, defaults), + type_annotation: v.type_annotation.clone(), + decorators: v.decorators.clone(), + }) + .collect(), + return_type_annotation: return_type_annotation.clone(), + decorators: decorators.clone(), + }, + ast::Statement::DefineDimension(span, name, dexprs) => { + ast::Statement::DefineDimension(*span, name, dexprs.clone()) + } + ast::Statement::DefineBaseUnit(span, name, type_expr, decorators) => { + ast::Statement::DefineBaseUnit(*span, name, type_expr.clone(), decorators.clone()) + } + ast::Statement::DefineDerivedUnit { + identifier_span, + identifier, + expr, + type_annotation_span, + type_annotation, + decorators, + } => ast::Statement::DefineDerivedUnit { + identifier_span: *identifier_span, + identifier, + expr: Self::apply_struct_defaults_expression(expr, defaults), + type_annotation_span: *type_annotation_span, + type_annotation: type_annotation.clone(), + decorators: decorators.clone(), + }, + ast::Statement::ProcedureCall(span, kind, args) => ast::Statement::ProcedureCall( + *span, + kind.clone(), + args.iter() + .map(|arg| Self::apply_struct_defaults_expression(arg, defaults)) + .collect(), + ), + ast::Statement::ModuleImport(span, module_path) => { + ast::Statement::ModuleImport(*span, module_path.clone()) + } + ast::Statement::DefineStruct { + struct_name_span, + struct_name, + type_parameters, + fields, + methods, + } => ast::Statement::DefineStruct { + struct_name_span: *struct_name_span, + struct_name, + type_parameters: type_parameters.clone(), + fields: fields + .iter() + .map(|(span, field_name, field_type, default_expr)| { + ( + *span, + *field_name, + field_type.clone(), + default_expr + .as_ref() + .map(|expr| Self::apply_struct_defaults_expression(expr, defaults)), + ) + }) + .collect(), + methods: methods + .iter() + .map(|method| Self::apply_struct_defaults_statement(method, defaults)) + .collect(), + }, + } + } + + 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 rewrite_self_in_method_decorators( + method: &mut ast::Statement<'_>, + struct_name: &str, + self_type_args: &[TypeAnnotation], + ) -> Result<()> { + let ast::Statement::DefineFunction { decorators, .. } = method else { + return Ok(()); + }; + + let _ = (decorators, struct_name, self_type_args); + + Ok(()) + } + + fn method_operator_impl(decorators: &[decorator::Decorator<'_>]) -> Option { + decorator::binary_operator(decorators) + .map(|(operator, reverse)| MethodOperatorInfo { operator, reverse }) + } + + fn method_index_impl(decorators: &[decorator::Decorator<'_>]) -> bool { + decorator::index_decorator(decorators) + } + fn fresh_type_variable(&mut self) -> Type { Type::TVar(self.name_generator.fresh_type_variable()) } @@ -217,6 +695,821 @@ 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, + operator_impl: Option, + index_impl: bool, + ) { + self.methods.entry(struct_name).or_default().insert( + method_name, + MethodInfo { + signature, + operator_impl, + index_impl, + }, + ); + } + + /// 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 validate_operator_method_decorator( + &self, + method_kind: StructMethodKind, + definition_span: Span, + method_name: &str, + decorators: &[decorator::Decorator<'_>], + fn_type: &TypeScheme, + type_parameters: &[(Span, &str, Option)], + ) -> Result> { + let Some((operator, reverse)) = decorator::binary_operator(decorators) else { + return Ok(None); + }; + + if method_kind != StructMethodKind::Instance { + return Err(Box::new(TypeCheckError::InvalidOperatorMethodSignature( + definition_span, + method_name.to_string(), + ))); + } + + let (instantiated_fn_type, _) = fn_type + .instantiate_for_printing(Some(type_parameters.iter().map(|(_, name, _)| *name))); + let Type::Fn(parameter_types, return_type) = instantiated_fn_type.inner else { + unreachable!("method type is expected to be a function type"); + }; + + if parameter_types.len() != 2 { + return Err(Box::new(TypeCheckError::InvalidOperatorMethodSignature( + definition_span, + method_name.to_string(), + ))); + } + + let rhs_type = parameter_types[1].clone(); + let output_type = return_type.as_ref().clone(); + + Ok(Some(typed_ast::StructMethodOperatorInfo { + operator, + reverse, + rhs_type, + output_type, + })) + } + + fn validate_index_method_decorator( + &self, + method_kind: StructMethodKind, + definition_span: Span, + method_name: &str, + decorators: &[decorator::Decorator<'_>], + fn_type: &TypeScheme, + type_parameters: &[(Span, &str, Option)], + ) -> Result { + if !decorator::index_decorator(decorators) { + return Ok(false); + } + + if method_kind != StructMethodKind::Instance { + return Err(Box::new(TypeCheckError::InvalidIndexMethodSignature( + definition_span, + method_name.to_string(), + ))); + } + + let (instantiated_fn_type, _) = fn_type + .instantiate_for_printing(Some(type_parameters.iter().map(|(_, name, _)| *name))); + let Type::Fn(parameter_types, _) = instantiated_fn_type.inner else { + unreachable!("method type is expected to be a function type"); + }; + + if parameter_types.len() < 2 { + return Err(Box::new(TypeCheckError::InvalidIndexMethodSignature( + definition_span, + method_name.to_string(), + ))); + } + + Ok(true) + } + + fn resolve_operator_candidates( + &mut self, + operator: BinaryOperator, + lhs: &Type, + rhs: &Type, + output: &Type, + reverse: bool, + ) -> Vec { + let struct_type = if reverse { rhs } else { lhs }; + let Some(struct_info) = (match struct_type { + Type::Struct(struct_info) => Some(struct_info), + _ => None, + }) else { + return vec![]; + }; + let Some(methods) = self.methods.get(&struct_info.name) else { + return vec![]; + }; + + let mut candidates = Vec::new(); + for method_info in methods.values() { + let Some(operator_impl) = method_info.operator_impl.as_ref() else { + continue; + }; + if operator_impl.operator != operator || operator_impl.reverse != reverse { + continue; + } + + let qualified_type = match &method_info.signature.fn_type { + TypeScheme::Concrete(type_) => { + crate::typechecker::qualified_type::QualifiedType::new( + type_.clone(), + qualified_type::Bounds::none(), + ) + } + TypeScheme::Quantified(_, _) => method_info + .signature + .fn_type + .instantiate(&mut self.name_generator), + }; + + let Type::Fn(parameter_types, return_type) = qualified_type.inner else { + continue; + }; + if parameter_types.len() != 2 { + continue; + } + + let receiver_type = if reverse { rhs } else { lhs }; + let other_type = if reverse { lhs } else { rhs }; + + let mut constraint_set = ConstraintSet::default(); + constraint_set + .add(Constraint::Equal( + parameter_types[0].clone(), + receiver_type.clone(), + )) + .ok(); + constraint_set + .add(Constraint::Equal( + parameter_types[1].clone(), + other_type.clone(), + )) + .ok(); + constraint_set + .add(Constraint::Equal( + return_type.as_ref().clone(), + output.clone(), + )) + .ok(); + for bound in qualified_type.bounds.iter() { + match bound { + Bound::IsDim(type_) => { + constraint_set.add(Constraint::IsDType(type_.clone())).ok(); + } + } + } + + if let Ok((substitution, dtype_vars)) = + constraint_set.solve(|_, _, _, _| None::) + && dtype_vars.is_empty() + { + candidates.push(constraints::Satisfied::with_substitution(substitution)); + } + } + + candidates + } + + fn try_resolve_binary_operator_constraint( + &mut self, + operator: BinaryOperator, + lhs: &Type, + rhs: &Type, + output: &Type, + ) -> Option { + let mut candidates = Vec::new(); + + match (operator, lhs, rhs) { + ( + BinaryOperator::Add | BinaryOperator::Sub, + Type::Dimension(lhs_dtype), + Type::Dimension(rhs_dtype), + ) if lhs_dtype == rhs_dtype => { + let mut constraints = vec![ + Constraint::Equal(lhs.clone(), rhs.clone()), + Constraint::Equal(lhs.clone(), output.clone()), + ]; + constraints.push(Constraint::IsDType(lhs.clone())); + candidates.push(constraints::Satisfied::with_new_constraints(constraints)); + } + (BinaryOperator::Mul, Type::Dimension(lhs_dtype), Type::Dimension(rhs_dtype)) => { + let result = Type::Dimension(lhs_dtype.multiply(rhs_dtype)); + candidates.push(constraints::Satisfied::with_new_constraints(vec![ + Constraint::Equal(output.clone(), result), + ])); + } + (BinaryOperator::Div, Type::Dimension(lhs_dtype), Type::Dimension(rhs_dtype)) => { + let result = Type::Dimension(lhs_dtype.divide(rhs_dtype)); + candidates.push(constraints::Satisfied::with_new_constraints(vec![ + Constraint::Equal(output.clone(), result), + ])); + } + _ => {} + } + + match (operator, lhs, rhs) { + (BinaryOperator::Add, Type::DateTime, Type::Dimension(dtype)) + | (BinaryOperator::Sub, Type::DateTime, Type::Dimension(dtype)) + if dtype.is_time_dimension() => + { + candidates.push(constraints::Satisfied::with_new_constraints(vec![ + Constraint::Equal(output.clone(), Type::DateTime), + ])); + } + (BinaryOperator::Sub, Type::DateTime, Type::DateTime) => { + candidates.push(constraints::Satisfied::with_new_constraints(vec![ + Constraint::Equal( + output.clone(), + Type::Dimension(DType::base_dimension("Time")), + ), + ])); + } + _ => {} + } + + let lhs_candidates = self.resolve_operator_candidates(operator, lhs, rhs, output, false); + if lhs_candidates.len() == 1 { + return lhs_candidates.into_iter().next(); + } + if lhs_candidates.len() > 1 { + return None; + } + + let rhs_candidates = self.resolve_operator_candidates(operator, lhs, rhs, output, true); + if rhs_candidates.len() == 1 { + return rhs_candidates.into_iter().next(); + } + + if candidates.len() == 1 { + candidates.pop() + } else { + None + } + } + + fn find_operator_method_call( + &mut self, + operator: BinaryOperator, + lhs: &Type, + rhs: &Type, + output: &Type, + ) -> Option<(bool, CompactString, CompactString)> { + let find_match = |this: &mut Self, + struct_type: &Type, + reverse: bool| + -> Option<(CompactString, CompactString)> { + let Type::Struct(struct_info) = struct_type else { + return None; + }; + let methods = this.methods.get(&struct_info.name)?; + let receiver_type = if reverse { rhs } else { lhs }; + let other_type = if reverse { lhs } else { rhs }; + + let mut matches = methods.iter().filter_map(|(method_name, method_info)| { + let operator_impl = method_info.operator_impl.as_ref()?; + if operator_impl.operator != operator || operator_impl.reverse != reverse { + return None; + } + + let qualified_type = match &method_info.signature.fn_type { + TypeScheme::Concrete(type_) => { + crate::typechecker::qualified_type::QualifiedType::new( + type_.clone(), + qualified_type::Bounds::none(), + ) + } + TypeScheme::Quantified(_, _) => method_info + .signature + .fn_type + .instantiate(&mut this.name_generator), + }; + + let Type::Fn(parameter_types, return_type) = qualified_type.inner else { + return None; + }; + if parameter_types.len() != 2 { + return None; + } + + let mut constraint_set = ConstraintSet::default(); + if constraint_set + .add(Constraint::Equal( + parameter_types[0].clone(), + receiver_type.clone(), + )) + .is_trivially_violated() + { + return None; + } + if constraint_set + .add(Constraint::Equal( + parameter_types[1].clone(), + other_type.clone(), + )) + .is_trivially_violated() + { + return None; + } + if constraint_set + .add(Constraint::Equal( + return_type.as_ref().clone(), + output.clone(), + )) + .is_trivially_violated() + { + return None; + } + + if let Ok((_, dtype_vars)) = + constraint_set.solve(|_, _, _, _| None::) + && dtype_vars.is_empty() + { + Some((struct_info.name.clone(), method_name.clone())) + } else { + None + } + }); + + let first = matches.next()?; + if matches.next().is_some() { + None + } else { + Some(first) + } + }; + + if let Some((owner, method_name)) = find_match(self, lhs, false) { + return Some((false, owner, method_name)); + } + + find_match(self, rhs, true).map(|(owner, method_name)| (true, owner, method_name)) + } + + fn find_index_method_call( + &mut self, + span: Span, + receiver_type: &Type, + index_types: &[Type], + ) -> Result<(CompactString, CompactString)> { + let Type::Struct(struct_info) = receiver_type else { + return Err(Box::new(TypeCheckError::IndexCallOnNonStructType( + span, + receiver_type.clone(), + ))); + }; + + let Some(methods) = self.methods.get(&struct_info.name) else { + return Err(Box::new(TypeCheckError::IndexMethodNotFound( + span, + struct_info.name.to_string(), + index_types.len(), + ))); + }; + + let mut matches = vec![]; + for (method_name, method_info) in methods { + if !method_info.index_impl { + continue; + } + + let qualified_type = match &method_info.signature.fn_type { + TypeScheme::Concrete(type_) => { + crate::typechecker::qualified_type::QualifiedType::new( + type_.clone(), + qualified_type::Bounds::none(), + ) + } + TypeScheme::Quantified(_, _) => method_info + .signature + .fn_type + .instantiate(&mut self.name_generator), + }; + + let Type::Fn(parameter_types, _) = qualified_type.inner else { + continue; + }; + if parameter_types.len() != index_types.len() + 1 { + continue; + } + + let mut constraint_set = ConstraintSet::default(); + if constraint_set + .add(Constraint::Equal( + parameter_types[0].clone(), + receiver_type.clone(), + )) + .is_trivially_violated() + { + continue; + } + + let mut valid = true; + for (parameter_type, index_type) in + parameter_types.iter().skip(1).zip(index_types.iter()) + { + if constraint_set + .add(Constraint::Equal( + parameter_type.clone(), + index_type.clone(), + )) + .is_trivially_violated() + { + valid = false; + break; + } + } + if !valid { + continue; + } + + if let Ok((_, dtype_vars)) = + constraint_set.solve(|_, _, _, _| None::) + && dtype_vars.is_empty() + { + matches.push((struct_info.name.clone(), method_name.clone())); + } + } + + match matches.len() { + 0 => Err(Box::new(TypeCheckError::IndexMethodNotFound( + span, + struct_info.name.to_string(), + index_types.len(), + ))), + 1 => Ok(matches.pop().unwrap()), + _ => Err(Box::new(TypeCheckError::AmbiguousIndexOverload( + span, + struct_info.name.to_string(), + index_types.len(), + ))), + } + } + + fn resolve_deferred_binary_operator_expression<'a>( + &mut self, + expr: &mut typed_ast::Expression<'a>, + ) -> Result<()> { + match expr { + typed_ast::Expression::UnaryOperator { expr, .. } => { + self.resolve_deferred_binary_operator_expression(expr)?; + } + typed_ast::Expression::BinaryOperator { lhs, rhs, .. } + | typed_ast::Expression::BinaryOperatorForDate { lhs, rhs, .. } => { + self.resolve_deferred_binary_operator_expression(lhs)?; + self.resolve_deferred_binary_operator_expression(rhs)?; + } + typed_ast::Expression::DeferredBinaryOperator { + op_span, + op, + lhs, + rhs, + type_scheme, + } => { + self.resolve_deferred_binary_operator_expression(lhs)?; + self.resolve_deferred_binary_operator_expression(rhs)?; + + let lhs_type = lhs.get_type_scheme().to_concrete_type(); + let rhs_type = rhs.get_type_scheme().to_concrete_type(); + let output_type = type_scheme.to_concrete_type(); + let full_span = lhs.full_span().extend(&rhs.full_span()); + + let replacement = if lhs_type == Type::DateTime + && (rhs_type == Type::DateTime + || matches!( + rhs_type, + Type::Dimension(ref dtype) if dtype.is_time_dimension() + )) { + typed_ast::Expression::BinaryOperatorForDate { + op_span: *op_span, + op: *op, + lhs: lhs.clone(), + rhs: rhs.clone(), + type_scheme: type_scheme.clone(), + } + } else if let Some((reverse, owner, method_name)) = + self.find_operator_method_call(*op, &lhs_type, &rhs_type, &output_type) + { + let receiver = if reverse { rhs.clone() } else { lhs.clone() }; + let arg = if reverse { + (**lhs).clone() + } else { + (**rhs).clone() + }; + let method_name_span = op_span.unwrap_or(full_span); + + typed_ast::Expression::MethodCall { + full_span, + receiver, + method_ref: typed_ast::MethodRef { + owner, + name: Box::leak(method_name.to_string().into_boxed_str()), + kind: StructMethodKind::Instance, + }, + method_name_span, + args: vec![arg], + type_scheme: type_scheme.clone(), + } + } else { + typed_ast::Expression::DeferredBinaryOperator { + op_span: *op_span, + op: *op, + lhs: lhs.clone(), + rhs: rhs.clone(), + type_scheme: type_scheme.clone(), + } + }; + + *expr = replacement; + } + typed_ast::Expression::FunctionCall { args, .. } => { + for arg in args { + self.resolve_deferred_binary_operator_expression(arg)?; + } + } + typed_ast::Expression::CallableCall { callable, args, .. } => { + self.resolve_deferred_binary_operator_expression(callable)?; + for arg in args { + self.resolve_deferred_binary_operator_expression(arg)?; + } + } + typed_ast::Expression::String(_, parts) => { + for part in parts { + if let typed_ast::StringPart::Interpolation { expr, .. } = part { + self.resolve_deferred_binary_operator_expression(expr)?; + } + } + } + typed_ast::Expression::Condition { + condition, + then_expr, + else_expr, + .. + } => { + self.resolve_deferred_binary_operator_expression(condition)?; + self.resolve_deferred_binary_operator_expression(then_expr)?; + self.resolve_deferred_binary_operator_expression(else_expr)?; + } + typed_ast::Expression::InstantiateStruct { fields, .. } => { + for (_, value) in fields { + self.resolve_deferred_binary_operator_expression(value)?; + } + } + typed_ast::Expression::AccessField { expr, .. } => { + self.resolve_deferred_binary_operator_expression(expr)?; + } + typed_ast::Expression::MethodCall { receiver, args, .. } => { + self.resolve_deferred_binary_operator_expression(receiver)?; + for arg in args { + self.resolve_deferred_binary_operator_expression(arg)?; + } + } + typed_ast::Expression::IndexCall { + full_span, + receiver, + args, + type_scheme, + } => { + self.resolve_deferred_binary_operator_expression(receiver)?; + for arg in args.iter_mut() { + self.resolve_deferred_binary_operator_expression(arg)?; + } + + let receiver_type = receiver.get_type_scheme().to_concrete_type(); + + if let Type::List(element_type) = &receiver_type { + if args.len() != 1 { + return Err(Box::new(TypeCheckError::InvalidListIndexArity( + *full_span, + args.len(), + ))); + } + + let index_type = args[0].get_type(); + if self + .add_equal_constraint(&index_type, &Type::scalar()) + .is_trivially_violated() + { + return Err(Box::new(TypeCheckError::InvalidListIndexType( + args[0].full_span(), + index_type, + ))); + } + + let result_type = (**element_type).clone(); + *type_scheme = TypeScheme::concrete(result_type); + *expr = typed_ast::Expression::FunctionCall { + full_span: *full_span, + ident_span: *full_span, + name: "_list_at", + args: vec![(**receiver).clone(), args[0].clone()], + type_scheme: type_scheme.clone(), + }; + } else if let Type::Struct(_) = &receiver_type { + let index_types = args + .iter() + .map(typed_ast::Expression::get_type) + .collect::>(); + let (owner, method_name) = + self.find_index_method_call(*full_span, &receiver_type, &index_types)?; + + let method = self + .lookup_method(&owner, &method_name) + .ok_or_else(|| { + Box::new(TypeCheckError::IndexMethodNotFound( + *full_span, + owner.to_string(), + index_types.len(), + )) + })? + .clone(); + + let mut method_arguments = Vec::with_capacity(args.len() + 1); + method_arguments.push((**receiver).clone()); + method_arguments.extend(args.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: full_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"); + }; + + *expr = typed_ast::Expression::MethodCall { + full_span: *full_span, + receiver: receiver.clone(), + method_ref: typed_ast::MethodRef { + owner, + name: Box::leak(method_name.to_string().into_boxed_str()), + kind: StructMethodKind::Instance, + }, + method_name_span: *full_span, + args: args.clone(), + type_scheme, + }; + } else if receiver_type.is_closed() { + return Err(Box::new(TypeCheckError::IndexCallOnNonStructType( + *full_span, + receiver_type, + ))); + } + } + typed_ast::Expression::List { elements, .. } => { + for element in elements { + self.resolve_deferred_binary_operator_expression(element)?; + } + } + typed_ast::Expression::Scalar { .. } + | typed_ast::Expression::Identifier { .. } + | typed_ast::Expression::UnitIdentifier { .. } + | typed_ast::Expression::Boolean(..) + | typed_ast::Expression::TypedHole(..) => {} + } + + Ok(()) + } + fn resolve_deferred_binary_operators<'a>( + &mut self, + statement: &mut typed_ast::Statement<'a>, + ) -> Result<()> { + match statement { + typed_ast::Statement::Expression(expr) => { + self.resolve_deferred_binary_operator_expression(expr)?; + } + typed_ast::Statement::DefineVariable(typed_ast::DefineVariable { expr, .. }) => { + self.resolve_deferred_binary_operator_expression(expr)?; + } + typed_ast::Statement::DefineFunction { + body, + local_variables, + .. + } + | typed_ast::Statement::DefineMethod { + body, + local_variables, + .. + } => { + for local_variable in local_variables { + self.resolve_deferred_binary_operator_expression(&mut local_variable.expr)?; + } + if let Some(body) = body { + self.resolve_deferred_binary_operator_expression(body)?; + } + } + typed_ast::Statement::DefineDerivedUnit { expr, .. } => { + self.resolve_deferred_binary_operator_expression(expr)?; + } + typed_ast::Statement::ProcedureCall { args, .. } => { + for arg in args { + self.resolve_deferred_binary_operator_expression(arg)?; + } + } + typed_ast::Statement::DefineStruct(_) + | typed_ast::Statement::DefineDimension(_, _) + | typed_ast::Statement::DefineBaseUnit { .. } => {} + } + + Ok(()) + } + + 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 +1526,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 +1559,7 @@ impl TypeChecker { name: struct_info.name.clone(), kind: StructKind::Instance(vec![]), fields: struct_info.fields.clone(), + methods: struct_info.methods.clone(), }))); } @@ -280,12 +1577,20 @@ impl TypeChecker { for (_, field_type) in instantiated_fields.values_mut() { field_type.apply(&substitution).ok(); } + let mut instantiated_methods = struct_info.methods.clone(); + for method_info in instantiated_methods.values_mut() { + if let Some(operator_impl) = &mut method_info.operator_impl { + operator_impl.rhs_type.apply(&substitution).ok(); + operator_impl.output_type.apply(&substitution).ok(); + } + } return Ok(Type::Struct(Box::new(StructInfo { definition_span: struct_info.definition_span, name: struct_info.name.clone(), kind: StructKind::Instance(concrete_type_args), fields: instantiated_fields, + methods: instantiated_methods, }))); } @@ -490,38 +1795,44 @@ impl TypeChecker { args: vec![lhs_checked], type_scheme: TypeScheme::concrete(*return_type), } - } else if lhs_type == Type::DateTime { - // DateTime types need special handling here, since they're not scalars with dimensions, - // yet some select binary operators can be applied to them - - let rhs_is_time = dtype(&rhs_checked) - .ok() - .map(|t| t.is_time_dimension()) - .unwrap_or(false); - let rhs_is_datetime = rhs_type == Type::DateTime; - - if *op == BinaryOperator::Sub && rhs_is_datetime { - let time = DType::base_dimension("Time"); // TODO: error handling - // TODO make sure the "second" unit exists - - typed_ast::Expression::BinaryOperatorForDate { - op_span: *span_op, - op: *op, - lhs: Box::new(lhs_checked), - rhs: Box::new(rhs_checked), - type_scheme: TypeScheme::concrete(Type::Dimension(time)), - } - } else if (*op == BinaryOperator::Add || *op == BinaryOperator::Sub) - && rhs_is_time + } else if matches!( + op, + BinaryOperator::Add + | BinaryOperator::Sub + | BinaryOperator::Mul + | BinaryOperator::Div + ) && { + let lhs_overload_candidate = + matches!(lhs_type, Type::DateTime | Type::Struct(_)); + let rhs_overload_candidate = + matches!(rhs_type, Type::DateTime | Type::Struct(_)); + let lhs_is_time = matches!( + lhs_type, + Type::Dimension(ref dtype) if dtype.is_time_dimension() + ); + let rhs_is_time = matches!( + rhs_type, + Type::Dimension(ref dtype) if dtype.is_time_dimension() + ); + + lhs_overload_candidate + || rhs_overload_candidate + || matches!(op, BinaryOperator::Add | BinaryOperator::Sub) + && ((!lhs_type.is_closed() && rhs_is_time) + || (!rhs_type.is_closed() && lhs_is_time)) + } { + let result_type = self.fresh_type_variable(); + if lhs_type.is_closed() + && rhs_type.is_closed() + && self + .try_resolve_binary_operator_constraint( + *op, + &lhs_type, + &rhs_type, + &result_type, + ) + .is_none() { - typed_ast::Expression::BinaryOperatorForDate { - op_span: *span_op, - op: *op, - lhs: Box::new(lhs_checked), - rhs: Box::new(rhs_checked), - type_scheme: TypeScheme::concrete(Type::DateTime), - } - } else { return Err(Box::new(TypeCheckError::IncompatibleTypesInOperator( span_op.unwrap_or_else(|| { ast::Expression::BinaryOperator { @@ -539,6 +1850,23 @@ impl TypeChecker { rhs.full_span(), ))); } + + self.constraints + .add(Constraint::HasBinaryOperator( + *op, + lhs_type.clone(), + rhs_type.clone(), + result_type.clone(), + )) + .ok(); + + typed_ast::Expression::DeferredBinaryOperator { + op_span: *span_op, + op: *op, + lhs: Box::new(lhs_checked), + rhs: Box::new(rhs_checked), + type_scheme: TypeScheme::concrete(result_type), + } } else { let mut get_type_and_assert_equal_dtypes = || -> Result { let lhs_type = lhs_checked.get_type(); @@ -1026,27 +2354,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) => { @@ -1080,6 +2425,13 @@ impl TypeChecker { for (_, field_type) in instantiated_fields.values_mut() { field_type.apply(&substitution).ok(); } + let mut instantiated_methods = struct_info.methods.clone(); + for method_info in instantiated_methods.values_mut() { + if let Some(operator_impl) = &mut method_info.operator_impl { + operator_impl.rhs_type.apply(&substitution).ok(); + operator_impl.output_type.apply(&substitution).ok(); + } + } StructInfo { definition_span: struct_info.definition_span, @@ -1088,6 +2440,7 @@ impl TypeChecker { variables.iter().map(|v| Type::TVar(v.clone())).collect(), ), fields: instantiated_fields, + methods: instantiated_methods, } } StructKind::Instance(_) => { @@ -1260,6 +2613,282 @@ 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, + } + } + ast::Expression::IndexCall { + receiver, + args, + full_span, + } => { + let receiver_checked = self.elaborate_expression(receiver)?; + let receiver_type = receiver_checked.get_type(); + + let arguments_checked = args + .iter() + .map(|e| self.elaborate_expression(e)) + .collect::>>()?; + + if let Type::List(element_type) = &receiver_type { + if arguments_checked.len() != 1 { + return Err(Box::new(TypeCheckError::InvalidListIndexArity( + *full_span, + arguments_checked.len(), + ))); + } + + let index_type = arguments_checked[0].get_type(); + if self + .add_equal_constraint(&index_type, &Type::scalar()) + .is_trivially_violated() + { + return Err(Box::new(TypeCheckError::InvalidListIndexType( + arguments_checked[0].full_span(), + index_type, + ))); + } + + return Ok(typed_ast::Expression::FunctionCall { + full_span: *full_span, + ident_span: *full_span, + name: "_list_at", + args: vec![ + receiver_checked, + arguments_checked.into_iter().next().unwrap(), + ], + type_scheme: TypeScheme::concrete((**element_type).clone()), + }); + } + + if !receiver_type.is_closed() { + return Ok(typed_ast::Expression::IndexCall { + full_span: *full_span, + receiver: Box::new(receiver_checked), + args: arguments_checked, + type_scheme: TypeScheme::concrete(self.fresh_type_variable()), + }); + } + + 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 (owner, method_name) = self.find_index_method_call( + *full_span, + &receiver_type, + &method_argument_types[1..], + )?; + + let method = self + .lookup_method(&owner, &method_name) + .ok_or_else(|| { + Box::new(TypeCheckError::IndexMethodNotFound( + *full_span, + owner.to_string(), + method_argument_types.len() - 1, + )) + })? + .clone(); + + let checked_call = proper_function_call(ProperFunctionCallArgs { + registry: &self.registry, + constraints: &mut self.constraints, + name_generator: &mut self.name_generator, + span: full_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, + name: Box::leak(method_name.to_string().into_boxed_str()), + kind: StructMethodKind::Instance, + }, + method_name_span: *full_span, + args: arguments_checked, + type_scheme, + } + } }) } @@ -1502,22 +3131,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 +3368,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 +3551,7 @@ impl TypeChecker { struct_name, type_parameters, fields, + methods, } => { self.type_namespace .add_identifier( @@ -1957,7 +3591,7 @@ impl TypeChecker { .push((*span, type_parameter.to_compact_string(), bound.clone())); } - for (span, field, _) in fields { + for (span, field, _, _) in fields { if let Some(other_span) = seen_fields.get(field) { return Err(Box::new(TypeCheckError::DuplicateFieldInStructDefinition( *span, @@ -1969,6 +3603,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(), @@ -1982,16 +3643,47 @@ impl TypeChecker { ), fields: fields .iter() - .map(|(span, name, type_)| { + .map(|(span, name, type_, _)| { Ok(( name.to_compact_string(), (*span, typechecker_struct.type_from_annotation(type_)?), )) }) .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, + operator_impl: None, + }, + ) + }) + .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) } @@ -2011,26 +3703,44 @@ impl TypeChecker { let mut elaborated_statement = self.elaborate_statement(statement)?; // Solve constraints - let (substitution, dtype_variables) = - self.constraints.solve().map_err(|inner| match inner { - ConstraintSolverError::CouldNotSolve(constraints) => { - TypeCheckError::ConstraintSolverError(statement.full_span(), constraints) - } - ConstraintSolverError::SubstitutionError(inner) => { - TypeCheckError::SubstitutionError( - elaborated_statement.pretty_print().to_string(), - inner, - ) - } - })?; + let mut constraints = std::mem::take(&mut self.constraints); + let solve_result = constraints.solve(|op, lhs, rhs, output| { + self.try_resolve_binary_operator_constraint(op, lhs, rhs, output) + }); + self.constraints = constraints; + + let (substitution, dtype_variables) = solve_result.map_err(|inner| match inner { + ConstraintSolverError::CouldNotSolve(constraints) => { + TypeCheckError::ConstraintSolverError(statement.full_span(), constraints) + } + ConstraintSolverError::SubstitutionError(inner) => TypeCheckError::SubstitutionError( + elaborated_statement.pretty_print().to_string(), + inner, + ), + })?; elaborated_statement.apply(&substitution).map_err(|e| { TypeCheckError::SubstitutionError(elaborated_statement.pretty_print().to_string(), e) })?; + self.resolve_deferred_binary_operators(&mut elaborated_statement)?; 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 +3798,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()? { @@ -2124,9 +3839,304 @@ impl TypeChecker { statements: &[ast::Statement<'a>], ) -> Result>> { let mut checked_statements = vec![]; + let struct_field_defaults = Self::collect_struct_field_defaults(statements); for statement in statements { - checked_statements.push(self.check_statement(statement)?); + let statement = + Self::apply_struct_defaults_statement(statement, &struct_field_defaults); + + 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(), + ))); + } + } + } + + let struct_statement_idx = checked_statements.len(); + 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, + )?; + Self::rewrite_self_in_method_decorators( + &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, + decorators, + .. + } = 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, + Self::method_operator_impl(decorators), + Self::method_index_impl(decorators), + ); + } + + 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"), + }; + let checked_method_decorators = match &checked_method { + typed_ast::Statement::DefineMethod { decorators, .. } => decorators.clone(), + _ => unreachable!("checked method is DefineMethod"), + }; + + 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, + }; + + let operator_impl = self.validate_operator_method_decorator( + method_kind.clone(), + *function_name_span, + function_name, + &checked_method_decorators, + &signature.fn_type, + &type_parameters + .iter() + .map(|(span, name, bound)| (*span, *name, bound.clone())) + .collect::>(), + )?; + let index_impl = self.validate_index_method_decorator( + method_kind.clone(), + *function_name_span, + function_name, + &checked_method_decorators, + &signature.fn_type, + &type_parameters + .iter() + .map(|(span, name, bound)| (*span, *name, bound.clone())) + .collect::>(), + )?; + + 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, + Self::method_operator_impl(&checked_method_decorators), + index_impl, + ); + + if let Some(struct_info) = self.structs.get_mut(*struct_name) + && let Some(method_info) = struct_info.methods.get_mut(*function_name) + { + method_info.operator_impl = operator_impl; + } + + checked_statements.push(checked_method); + } + + if let Some(typed_ast::Statement::DefineStruct(struct_info)) = + checked_statements.get_mut(struct_statement_idx) + && let Some(updated_struct_info) = self.structs.get(*struct_name) + { + *struct_info = updated_struct_info.clone(); + } + self.non_generic_struct_instances + .remove(&struct_name.to_compact_string()); + + continue; + } + + checked_statements.push(self.check_statement(&statement)?); } Ok(checked_statements) diff --git a/numbat/src/typechecker/substitutions.rs b/numbat/src/typechecker/substitutions.rs index 217d384c..621b16ad 100644 --- a/numbat/src/typechecker/substitutions.rs +++ b/numbat/src/typechecker/substitutions.rs @@ -85,19 +85,7 @@ impl ApplySubstitution for Type { } return_type.apply(s) } - Type::Struct(info) => { - // Apply substitution to type arguments - if let StructKind::Instance(type_args) = &mut info.kind { - for arg in type_args { - arg.apply(s)?; - } - } - // Apply substitution to field types - for (_, field_type) in info.fields.values_mut() { - field_type.apply(s)?; - } - Ok(()) - } + Type::Struct(info) => info.apply(s), Type::List(element_type) => element_type.apply(s), } } @@ -165,6 +153,12 @@ impl ApplySubstitution for StructInfo { for (_, field_type) in self.fields.values_mut() { field_type.apply(s)?; } + for method_info in self.methods.values_mut() { + if let Some(operator_impl) = &mut method_info.operator_impl { + operator_impl.rhs_type.apply(s)?; + operator_impl.output_type.apply(s)?; + } + } Ok(()) } } @@ -201,6 +195,16 @@ impl ApplySubstitution for Expression<'_> { rhs.apply(s)?; type_scheme.apply(s) } + Expression::DeferredBinaryOperator { + lhs, + rhs, + type_scheme, + .. + } => { + lhs.apply(s)?; + rhs.apply(s)?; + type_scheme.apply(s) + } Expression::FunctionCall { args, type_scheme, .. } => { @@ -253,6 +257,18 @@ impl ApplySubstitution for Expression<'_> { struct_type.apply(s)?; field_type.apply(s) } + Expression::IndexCall { + receiver, + args, + type_scheme, + .. + } => { + receiver.apply(s)?; + for arg in args { + arg.apply(s)?; + } + type_scheme.apply(s) + } Expression::List { elements, type_scheme, @@ -263,6 +279,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 +311,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..5dac3b7b 100644 --- a/numbat/src/typechecker/tests/type_checking.rs +++ b/numbat/src/typechecker/tests/type_checking.rs @@ -771,6 +771,524 @@ 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_successful_typecheck( + " + struct Point { + x: A, + y: A, + + @add + fn add(self, rhs: Self) -> Self = + Point { x: self.x + rhs.x, y: self.y + rhs.y } + } + + let sum: Point = Point { x: 1 a, y: 2 a } + Point { x: 3 a, y: 4 a } + let x: A = sum.x + ", + ); + + assert_successful_typecheck( + " + struct Point { + x: A, + y: A, + + @add + fn add(self, rhs: Self) -> Self = + Point { x: self.x + rhs.x, y: self.y + rhs.y } + + @add + fn add_scalar(self, rhs: Scalar) -> Self = + Point { x: self.x + rhs * 1 a, y: self.y + rhs * 1 a } + } + + let sum_points: Point = Point { x: 1 a, y: 2 a } + Point { x: 3 a, y: 4 a } + let shifted: Point = Point { x: 1 a, y: 2 a } + 3 + let x1: A = sum_points.x + let x2: A = shifted.x + ", + ); + + assert_successful_typecheck( + " + struct Point { + x: A, + y: A, + } + + struct Shift { + amount: A, + + @radd + fn add_to_point(self, lhs: Point) -> Point = + Point { x: lhs.x + self.amount, y: lhs.y + self.amount } + } + + let shifted: Point = Point { x: 1 a, y: 2 a } + Shift { amount: 3 a } + let x: A = shifted.x + ", + ); + + assert_successful_typecheck( + " + struct Vec2 { + x: D, + y: D, + + @add + fn add(self, rhs: Self) -> Self = + Self { x: self.x + rhs.x, y: self.y + rhs.y } + + @mul + fn scale(self, factor: Scalar) -> Self = + Self { x: self.x * factor, y: self.y * factor } + + @rmul + fn scale_from_left(self, lhs: Scalar) -> Self = + Self { x: lhs * self.x, y: lhs * self.y } + } + + let sum: Vec2 = Vec2 { x: 1 a, y: 2 a } + Vec2 { x: 3 a, y: 4 a } + let scaled: Vec2 = Vec2 { x: 1 a, y: 2 a } * 3 + let reverse_scaled: Vec2 = 3 * Vec2 { x: 1 a, y: 2 a } + let x1: A = sum.x + let x2: A = scaled.x + let x3: A = reverse_scaled.x + ", + ); + + assert_successful_typecheck( + " + struct Pair { + left: Scalar, + right: Scalar, + + @index + fn get(self, i: Scalar) -> Scalar = + if i == 0 then self.left else self.right + } + + struct Grid2 { + a: Scalar, + b: Scalar, + c: Scalar, + d: Scalar, + + @index + fn get(self, row: Scalar, col: Scalar) -> Scalar = + if row == 0 then + if col == 0 then self.a else self.b + else + if col == 0 then self.c else self.d + } + + let first: Scalar = Pair { left: 10, right: 20 }[0] + let last: Scalar = Pair { left: 10, right: 20 }[1] + let bottom_right: Scalar = Grid2 { a: 1, b: 2, c: 3, d: 4 }[1, 1] + ", + ); + + assert_successful_typecheck( + " + let x: Scalar = [10, 20, 30][1] + let y: A = [1 a, 2 a, 3 a][2] + ", + ); + + assert_successful_typecheck( + " + struct Vec2 { + x: Scalar, + y: Scalar, + + @add + fn add(self, rhs: Self) -> Self = + Vec2 { x: self.x + rhs.x, y: self.y + rhs.y } + + @index + fn get(self, i: Scalar) -> Scalar = + if i == 0 then self.x else self.y + } + + let value: Scalar = (Vec2 { x: 1, y: 2 } + Vec2 { x: 3, y: 4 })[1] + ", + ); + + assert_successful_typecheck( + " + struct Point { + x: A, + y: A, + } + + struct Offset { + amount: A, + + @rsub + fn sub_from_point(self, lhs: Point) -> Point = + Point { x: lhs.x - self.amount, y: lhs.y - self.amount } + } + + struct Scale { + factor: Scalar, + + @rmul + fn scale_point(self, lhs: Point) -> Point = + Point { x: lhs.x * self.factor, y: lhs.y * self.factor } + } + + struct Ratio { + factor: Scalar, + + @rdiv + fn div_point(self, lhs: Point) -> Point = + Point { x: lhs.x / self.factor, y: lhs.y / self.factor } + } + + let p_sub: Point = Point { x: 5 a, y: 7 a } - Offset { amount: 2 a } + let p_mul: Point = Point { x: 2 a, y: 3 a } * Scale { factor: 4 } + let p_div: Point = Point { x: 8 a, y: 6 a } / Ratio { factor: 2 } + let sub_x: A = p_sub.x + let mul_y: A = p_mul.y + let div_x: A = p_div.x + ", + ); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + @add + fn make() -> Self = Point { x: 1 a } + } + " + ), + TypeCheckError::InvalidOperatorMethodSignature(_, name) if name == "make" + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + @index + fn make(self) -> A = self.x + } + " + ), + TypeCheckError::InvalidIndexMethodSignature(_, name) if name == "make" + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { x: A } + struct Shift { + amount: A, + @radd + fn make() -> Shift = Shift { amount: 1 a } + } + " + ), + TypeCheckError::InvalidOperatorMethodSignature(_, name) if name == "make" + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + y: A, + } + + struct Shift { + amount: A, + + @radd + fn add_to_point(self, lhs: Point) -> Point = + Point { x: lhs.x + self.amount, y: lhs.y + self.amount } + } + + Shift { amount: 3 a } + Point { x: 1 a, y: 2 a } + " + ), + TypeCheckError::IncompatibleTypesInOperator(..) + )); + + assert!(matches!( + get_typecheck_error( + " + struct Point { + x: A, + @index + fn get(self, i: Scalar) -> A = self.x + } + + let y = Point { x: 1 a }[1 a] + " + ), + TypeCheckError::IndexMethodNotFound(..) + )); + + assert!(matches!( + get_typecheck_error("let x = [1, 2][0, 1]"), + TypeCheckError::InvalidListIndexArity(_, 2) + )); + + assert!(matches!( + get_typecheck_error("let x = [1, 2][1 a]"), + TypeCheckError::InvalidListIndexType(..) + )); + + 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..e73d11d5 100644 --- a/numbat/src/typed_ast.rs +++ b/numbat/src/typed_ast.rs @@ -337,6 +337,28 @@ 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, + pub operator_impl: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StructMethodOperatorInfo { + pub operator: BinaryOperator, + pub reverse: bool, + pub rhs_type: Type, + pub output_type: Type, } /// A monomorphic type (no quantifiers). @@ -607,6 +629,7 @@ impl Type { name: info.name.clone(), kind: instantiated_kind, fields: instantiated_fields, + methods: info.methods.clone(), })) } Type::List(element_type) => { @@ -621,6 +644,40 @@ impl Type { _ => false, } } + + pub(crate) fn equals_ignoring_struct_methods(&self, other: &Type) -> bool { + match (self, other) { + (Type::TVar(v1), Type::TVar(v2)) => v1 == v2, + (Type::TPar(v1), Type::TPar(v2)) => v1 == v2, + (Type::Dimension(d1), Type::Dimension(d2)) => d1 == d2, + (Type::Boolean, Type::Boolean) + | (Type::String, Type::String) + | (Type::DateTime, Type::DateTime) => true, + (Type::List(i1), Type::List(i2)) => i1.equals_ignoring_struct_methods(i2), + (Type::Fn(p1, r1), Type::Fn(p2, r2)) => { + p1.len() == p2.len() + && p1 + .iter() + .zip(p2.iter()) + .all(|(a, b)| a.equals_ignoring_struct_methods(b)) + && r1.equals_ignoring_struct_methods(r2) + } + (Type::Struct(info1), Type::Struct(info2)) if info1.name == info2.name => { + match (&info1.kind, &info2.kind) { + (StructKind::Definition(_), StructKind::Definition(_)) => true, + (StructKind::Instance(args1), StructKind::Instance(args2)) => { + args1.len() == args2.len() + && args1 + .iter() + .zip(args2.iter()) + .all(|(a, b)| a.equals_ignoring_struct_methods(b)) + } + _ => false, + } + } + _ => false, + } + } } #[derive(Debug, Clone, PartialEq)] @@ -704,6 +761,13 @@ pub enum Expression<'a> { rhs: Box>, type_scheme: TypeScheme, }, + DeferredBinaryOperator { + op_span: Option, + op: BinaryOperator, + lhs: Box>, + rhs: Box>, + type_scheme: TypeScheme, + }, /// A 'proper' function call FunctionCall { full_span: Span, @@ -740,6 +804,20 @@ 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, + }, + IndexCall { + full_span: Span, + receiver: Box>, + args: Vec>, + type_scheme: TypeScheme, + }, List { span: Span, elements: Vec>, @@ -748,6 +826,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 { @@ -773,6 +858,15 @@ impl Expression<'_> { } span } + Expression::DeferredBinaryOperator { + op_span, lhs, rhs, .. + } => { + let mut span = lhs.full_span().extend(&rhs.full_span()); + if let Some(op_span) = op_span { + span = span.extend(op_span); + } + span + } Expression::FunctionCall { full_span, .. } => *full_span, Expression::CallableCall { full_span, .. } => *full_span, Expression::Boolean(span, _) => *span, @@ -782,6 +876,8 @@ impl Expression<'_> { Expression::String(span, _) => *span, Expression::InstantiateStruct { span, .. } => *span, Expression::AccessField { full_span, .. } => *full_span, + Expression::MethodCall { full_span, .. } => *full_span, + Expression::IndexCall { full_span, .. } => *full_span, Expression::List { span, .. } => *span, Expression::TypedHole(span, _) => *span, } @@ -818,6 +914,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 +1002,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 +1111,12 @@ impl Statement<'_> { local_variables, fn_type, .. + } + | Statement::DefineMethod { + parameters, + local_variables, + fn_type, + .. } => { let mut bindings = Vec::new(); @@ -1034,6 +1162,9 @@ impl Expression<'_> { Expression::BinaryOperatorForDate { type_scheme, .. } => { type_scheme.unsafe_as_concrete() } + Expression::DeferredBinaryOperator { type_scheme, .. } => { + type_scheme.unsafe_as_concrete() + } Expression::FunctionCall { type_scheme, .. } => type_scheme.unsafe_as_concrete(), Expression::CallableCall { type_scheme, .. } => type_scheme.unsafe_as_concrete(), Expression::Boolean(_, _) => Type::Boolean, @@ -1043,9 +1174,11 @@ impl Expression<'_> { Type::Struct(Box::new(struct_info.clone())) } Expression::AccessField { field_type, .. } => field_type.unsafe_as_concrete(), + Expression::IndexCall { type_scheme, .. } => type_scheme.unsafe_as_concrete(), 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(), } } @@ -1058,6 +1191,7 @@ impl Expression<'_> { Expression::UnaryOperator { type_scheme, .. } => type_scheme.clone(), Expression::BinaryOperator { type_scheme, .. } => type_scheme.clone(), Expression::BinaryOperatorForDate { type_scheme, .. } => type_scheme.clone(), + Expression::DeferredBinaryOperator { type_scheme, .. } => type_scheme.clone(), Expression::FunctionCall { type_scheme, .. } => type_scheme.clone(), Expression::CallableCall { type_scheme, .. } => type_scheme.clone(), Expression::Boolean(_, _) => TypeScheme::make_quantified(Type::Boolean), @@ -1067,6 +1201,7 @@ impl Expression<'_> { TypeScheme::make_quantified(Type::Struct(Box::new(struct_info.clone()))) } Expression::AccessField { field_type, .. } => field_type.clone(), + Expression::IndexCall { type_scheme, .. } => type_scheme.clone(), Expression::List { type_scheme, .. } => match type_scheme { TypeScheme::Concrete(t) => TypeScheme::Concrete(Type::List(Box::new(t.clone()))), TypeScheme::Quantified(ngen, qt) => TypeScheme::Quantified( @@ -1077,6 +1212,7 @@ impl Expression<'_> { }, ), }, + Expression::MethodCall { type_scheme, .. } => type_scheme.clone(), Expression::TypedHole(_, type_) => type_.clone(), } } @@ -1159,6 +1295,22 @@ fn decorator_markup(decorators: &Vec) -> Markup { } + m::operator(")") } + Decorator::BinaryOperator { operator, reverse } => { + let decorator_name = match (operator, reverse) { + (BinaryOperator::Add, false) => "@add", + (BinaryOperator::Sub, false) => "@sub", + (BinaryOperator::Mul, false) => "@mul", + (BinaryOperator::Div, false) => "@div", + (BinaryOperator::Add, true) => "@radd", + (BinaryOperator::Sub, true) => "@rsub", + (BinaryOperator::Mul, true) => "@rmul", + (BinaryOperator::Div, true) => "@rdiv", + _ => unreachable!("unsupported binary operator decorator"), + }; + + m::decorator(decorator_name) + } + Decorator::Index => m::decorator("@index"), } + m::nl(); } @@ -1246,6 +1398,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,11 +1574,14 @@ fn with_parens(expr: &Expression) -> Markup { | Expression::String(..) | Expression::InstantiateStruct { .. } | Expression::AccessField { .. } + | Expression::MethodCall { .. } + | Expression::IndexCall { .. } | Expression::List { .. } | Expression::TypedHole(_, _) => expr.pretty_print(), Expression::UnaryOperator { .. } | Expression::BinaryOperator { .. } | Expression::BinaryOperatorForDate { .. } + | Expression::DeferredBinaryOperator { .. } | Expression::Condition { .. } => m::operator("(") + expr.pretty_print() + m::operator(")"), } } @@ -1594,6 +1759,7 @@ impl PrettyPrint for Expression<'_> { } => m::operator("!") + with_parens(expr), BinaryOperator { op, lhs, rhs, .. } => pretty_print_binop(op, lhs, rhs), BinaryOperatorForDate { op, lhs, rhs, .. } => pretty_print_binop(op, lhs, rhs), + DeferredBinaryOperator { op, lhs, rhs, .. } => pretty_print_binop(op, lhs, rhs), FunctionCall { name, args, .. } => { // Special case: render special temperature conversion functions in their sugar form: if args.len() == 1 { @@ -1726,6 +1892,36 @@ 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(")") + } + IndexCall { receiver, args, .. } => { + receiver.pretty_print() + + 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/unit_registry.rs b/numbat/src/unit_registry.rs index e5fc53bf..b2e704d1 100644 --- a/numbat/src/unit_registry.rs +++ b/numbat/src/unit_registry.rs @@ -99,3 +99,9 @@ impl UnitRegistry { base_unit_match.into_iter().chain(derived_units) } } + +impl Default for UnitRegistry { + fn default() -> Self { + Self::new() + } +} 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..7c5b681c 100644 --- a/numbat/tests/interpreter.rs +++ b/numbat/tests/interpreter.rs @@ -137,6 +137,225 @@ fn get_diagnostic_output(code: &str) -> String { } } +#[test] +fn struct_operator_methods() { + expect_output( + " + struct Point { + x: Scalar, + y: Scalar, + + @add + fn add(self, rhs: Self) -> Self = + Point { x: self.x + rhs.x, y: self.y + rhs.y } + } + + (Point { x: 1, y: 2 } + Point { x: 3, y: 4 }).x + ", + "4", + ); + + expect_output( + " + struct Point { + x: Scalar, + y: Scalar, + + @add + fn add(self, rhs: Self) -> Self = + Point { x: self.x + rhs.x, y: self.y + rhs.y } + + @add + fn add_scalar(self, rhs: Scalar) -> Self = + Point { x: self.x + rhs, y: self.y + rhs } + } + + (Point { x: 1, y: 2 } + 3).x + (Point { x: 1, y: 2 } + Point { x: 3, y: 4 }).y + ", + "10", + ); + + expect_output( + " + struct Point { + x: Scalar, + y: Scalar, + } + + struct Shift { + amount: Scalar, + + @radd + fn add_to_point(self, lhs: Point) -> Point = + Point { x: lhs.x + self.amount, y: lhs.y + self.amount } + } + + (Point { x: 1, y: 2 } + Shift { amount: 3 }).y + ", + "5", + ); + + expect_output( + " + struct Vec2 { + x: D, + y: D, + + @add + fn add(self, rhs: Self) -> Self = + Self { x: self.x + rhs.x, y: self.y + rhs.y } + + @mul + fn scale(self, factor: Scalar) -> Self = + Self { x: self.x * factor, y: self.y * factor } + + @rmul + fn scale_from_left(self, lhs: Scalar) -> Self = + Self { x: lhs * self.x, y: lhs * self.y } + } + + let v = Vec2 { x: 3 m, y: 4 m } + let w = Vec2 { x: 300 cm, y: 400 cm } + + ((v + w).x -> m) + ((v * 2).y -> m) + ((2 * v).y -> m) + ", + "22 m", + ); + + expect_output( + " + struct Point { + x: Scalar, + y: Scalar, + } + + struct Shift { + amount: Scalar, + + @radd + fn add_to_point(self, lhs: Point) -> Point = + Point { x: lhs.x + self.amount, y: lhs.y + self.amount } + } + + struct Offset { + amount: Scalar, + + @rsub + fn sub_from_point(self, lhs: Point) -> Point = + Point { x: lhs.x - self.amount, y: lhs.y - self.amount } + } + + struct Scale { + factor: Scalar, + + @rmul + fn scale_point(self, lhs: Point) -> Point = + Point { x: lhs.x * self.factor, y: lhs.y * self.factor } + } + + struct Ratio { + factor: Scalar, + + @rdiv + fn div_point(self, lhs: Point) -> Point = + Point { x: lhs.x / self.factor, y: lhs.y / self.factor } + } + + ((Point { x: 1, y: 2 } + Shift { amount: 3 }).x + (Point { x: 5, y: 7 } - Offset { amount: 2 }).y + (Point { x: 2, y: 3 } * Scale { factor: 4 }).x + (Point { x: 8, y: 6 } / Ratio { factor: 2 }).y) + ", + "20", + ); + + expect_output( + " + struct Pair { + left: Scalar, + right: Scalar, + + @index + fn get(self, i: Scalar) -> Scalar = + if i == 0 then self.left else self.right + } + + struct Grid2 { + a: Scalar, + b: Scalar, + c: Scalar, + d: Scalar, + + @index + fn get(self, row: Scalar, col: Scalar) -> Scalar = + if row == 0 then + if col == 0 then self.a else self.b + else + if col == 0 then self.c else self.d + } + + Pair { left: 10, right: 20 }[1] + Grid2 { a: 1, b: 2, c: 3, d: 4 }[1, 0] + ", + "23", + ); + + expect_output("[10, 20, 30][1]", "20"); + expect_output("[1 m, 2 m, 3 m][2]", "3 m"); + + expect_output( + " + struct Vec2 { + x: Scalar, + y: Scalar, + + @add + fn add(self, rhs: Self) -> Self = + Vec2 { x: self.x + rhs.x, y: self.y + rhs.y } + + @index + fn get(self, i: Scalar) -> Scalar = + if i == 0 then self.x else self.y + } + + (Vec2 { x: 1, y: 2 } + Vec2 { x: 3, y: 4 })[1] + ", + "6", + ); + + expect_failure( + "[10, 20, 30][-1]", + "List index must be a non-negative integer", + ); + expect_failure( + "[10, 20, 30][1.5]", + "List index must be a non-negative integer", + ); + expect_failure("[10, 20, 30][3]", "out of bounds"); + + expect_output( + " + struct Point { + x: Scalar, + y: Scalar, + + @add + fn add(self, rhs: Self) -> Self = + Point { x: self.x + rhs.x, y: self.y + rhs.y } + } + + ( + (Point { x: 1, y: 2 } + Point { x: 3, y: 4 }).x + + 1 + ) + ", + "5", + ); + + let _ = succeed( + " + fn id(x) = x + id(now()) + 4 days + ", + ); +} + #[test] fn simple_value() { expect_output("0", "0"); @@ -1243,6 +1462,367 @@ 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_struct_field_defaults() { + expect_output( + " + struct Point { + x: Length = 2 m, + y: Length, + } + + let p = Point { y: 3 m } + p.x + p.y + ", + "5 m", + ); + + expect_output( + " + struct Counter { + n: Scalar = 10, + fn new() -> Self = Counter {} + } + + Counter::new().n + ", + "10", + ); + + expect_failure( + " + struct Point { + x: Length = 2 m, + y: Length, + } + + Point {} + ", + "Missing field", + ); +} + +#[test] +fn test_struct_features_comprehensive() { + expect_output( + " + use prelude + + struct Vec2 { + x: D = 0 m, + y: D = 0 m, + + fn new(x: D, y: D) -> Self = Vec2 { x: x, y: y } + fn with_y(y: D) -> Self = Vec2 { y: y } + fn translate(self, dx: D, dy: D) -> Self = Vec2 { x: self.x + dx, y: self.y + dy } + fn dot(self, other: Vec2) -> D * E = self.x * other.x + self.y * other.y + } + + let v1 = Vec2::new(1 m, 2 m) + let v2 = Vec2::with_y(4 m).translate(3 m, 0 m) + let v3 = Vec2 { y: 5 m } + let v4 = Vec2 { x: 300 cm, y: 400 cm } + + v1.dot(v2) + v3.dot(v4) + sqrt(4) m^2 + ", + "330_000 cm²", + ); +} + +#[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__matrix@matrix.nbt.snap b/numbat/tests/snapshots/example_snapshots__matrix@matrix.nbt.snap new file mode 100644 index 00000000..216a4fab --- /dev/null +++ b/numbat/tests/snapshots/example_snapshots__matrix@matrix.nbt.snap @@ -0,0 +1,10 @@ +--- +source: numbat/tests/example_snapshots.rs +assertion_line: 53 +expression: output +input_file: examples/matrix.nbt +--- +2x2 identity: Matrix { rows: 2, cols: 2, data: [[1, 0], [0, 1]] } +2x3 matrix: Matrix { rows: 2, cols: 3, data: [[1, 2, 3], [4, 5, 6]] } +3x2 matrix: Matrix { rows: 3, cols: 2, data: [[1, 2], [3, 4], [5, 6]] } +product: Matrix { rows: 2, cols: 2, data: [[22, 28], [49, 64]] } 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 diff --git a/numbat/tests/snapshots/prelude_and_examples__parse_error_snapshots@missing_closing_paren3.nbt.snap b/numbat/tests/snapshots/prelude_and_examples__parse_error_snapshots@missing_closing_paren3.nbt.snap index e625faf3..2113b0a0 100644 --- a/numbat/tests/snapshots/prelude_and_examples__parse_error_snapshots@missing_closing_paren3.nbt.snap +++ b/numbat/tests/snapshots/prelude_and_examples__parse_error_snapshots@missing_closing_paren3.nbt.snap @@ -4,9 +4,7 @@ expression: output input_file: examples/parse_error/missing_closing_paren3.nbt --- error: while parsing - ┌─ :1:5 - │ -1 │ (2+3 - │ ╭────^ -2 │ │ - │ ╰^ Missing closing parenthesis ')' + ┌─ :2:1 + │ +2 │ + │ ^ Missing closing parenthesis ')'