#set document(title: "2.9 Add operator support to custom types", author: "Modular Inc. / XYZ Homework") #set page(width: 8.5in, height: auto, margin: 1in) #import "@preview/cetz:0.5.2" #set text(font: ("STIX Two Text", "Libertinus Serif", "New Computer Modern"), size: 10.5pt, lang: "en") #show math.equation: set text(font: ("STIX Two Math", "New Computer Modern Math")) #set par(justify: true, leading: 0.62em, spacing: 0.9em) #set enum(spacing: 1.1em) // room between list items so tall inline fractions don't collide #set list(spacing: 1.1em) #set table(stroke: 0.5pt + rgb("#c7ccd3")) #let BLUE = rgb("#183B6F") // brand navy — section bars + example/solution labels (white on navy 11.09:1) #let ORANGE = rgb("#A94509") // brand primary-700 — AA-safe deep orange for TEXT (5.93:1 on white; raw brand #F37021 is 2.94:1 and must never carry text) #let RED = rgb("#DC2626") // brand error-600 #let GREEN = rgb("#059669") // brand success-600 (decoration only; small green text uses green-text #007942) #show heading.where(level: 1): it => block(width: 100%, above: 0pt, below: 16pt, fill: gradient.linear(BLUE, rgb("#2C5AA0")), inset: (x: 14pt, y: 12pt), radius: 3pt, text(fill: white, weight: "bold", size: 19pt, it.body)) #show heading.where(level: 2): it => block(width: 100%, above: 18pt, below: 10pt, fill: BLUE, inset: (x: 10pt, y: 6pt), radius: 2pt, text(fill: white, weight: "bold", size: 12pt, it.body)) #show heading.where(level: 3): it => text(fill: ORANGE, weight: "bold", size: 12.5pt, it.body) #show heading.where(level: 4): it => text(fill: BLUE, weight: "bold", size: 10.5pt, it.body) #let examplebox(label, title, body) = block(width: 100%, breakable: true, fill: rgb("#EFF1F5"), stroke: 0.5pt + rgb("#CFDDF0"), radius: 4pt, inset: 10pt, above: 12pt, below: 12pt)[ #block(below: 6pt)[#box(fill: BLUE, inset: (x: 6pt, y: 2pt), radius: 2pt, text(fill: white, weight: "bold", size: 8.5pt, label)) #h(0.4em) #strong[#title]] #body] // rail = decorative left rule (raw brand token); labelcolor = AA-safe label text shade #let notebox(label, rail, labelcolor, tint, body) = block(width: 100%, breakable: true, fill: tint, stroke: (left: 3pt + rail), inset: (left: 10pt, rest: 8pt), radius: (right: 4pt), above: 11pt, below: 11pt)[ #text(fill: labelcolor, weight: "bold", size: 7.5pt, tracking: 0.5pt)[#upper(label)] #linebreak() #body] #let solutionbox(body) = block(above: 4pt, below: 8pt)[ #text(fill: BLUE, weight: "bold", size: 8.5pt)[Solution] #linebreak() #body] #let figph(msg) = block(width: 100%, height: 60pt, fill: rgb("#f6f7f9"), stroke: (paint: rgb("#c7ccd3"), dash: "dashed"), radius: 4pt, inset: 10pt)[ #align(center + horizon, text(fill: rgb("#889"), style: "italic", size: 9pt, msg))] // Standardize inlined figure sizes: measure the natural CeTZ canvas, then scale to a // consistent envelope (aspect-aware; see build_typst.py FIG_* constants). Unlike the // print preamble, dimensions are FLOORED: in an editor a user can trim a figure to a // degenerate 1-D shape (a bare line), and w/h or tw/w would then divide by zero. #let _STD_W = 3.5 #let _WIDE_W = 5.6 #let _MAX_H = 3.4 #let _ASPECT_WIDE = 2.2 #let _UPSCALE_MAX = 1.15 #let stdfig(body) = context { let m = measure(body) let w = calc.max(m.width / 1in, 0.01) let h = calc.max(m.height / 1in, 0.01) let tw = if w / h > _ASPECT_WIDE { _WIDE_W } else { _STD_W } let s = calc.min(tw / w, _MAX_H / h, _UPSCALE_MAX) align(center, box(scale(x: s * 100%, y: s * 100%, reflow: true, body))) } #show figure: set block(breakable: false) #set figure(gap: 8pt) #show figure.caption: set text(size: 8.5pt, fill: rgb("#555")) == 2.9#h(0.6em)Add operator support to custom types Each Mojo operator maps to a set of dunder methods you can add to your struct implementation. These methods let you use operator syntax instead of calling methods directly. Knowing your operators and their related methods opens the full suite of operator syntax to your custom structs. === Forward, reverse, and in-place methods Each binary operator uses up to three method forms. For example, consider the addition a + b: - #strong[Forward]: Mojo tries a.\_\_add\_\_(b) first. - #strong[Reverse]: If the forward method doesn't exist or can't handle b's type, Mojo falls back to b.\_\_radd\_\_(a). - #strong[In-place]: For a += b, Mojo calls a.\_\_iadd\_\_(b). Reversed methods exist for mixed-type expressions where the left operand doesn't know about the right operand's type: a + 5 \# calls a.\_\_add\_\_(5) 5 + a \# Int doesn't know your type, falls back \# to a.\_\_radd\_\_(5) a += 5 \# calls a.\_\_iadd\_\_(5) === Unary operators A unary operator returns the original value if unchanged, or a new value representing the result. For example, -x uses the unary negation operator: \@fieldwise\_init struct MyInt: var value: Int def \_\_neg\_\_(self) -\> Self: return Self(-self.value) If x is a MyInt, then -x returns a new instance with its value field negated. === Comparison operators and traits Operators do not require that you conform your types to traits. However, there are benefits to doing so. The Comparable trait provides defaults for \<=, \>, and \>=. You just implement \_\_lt\_\_() and \_\_eq\_\_(). Similarly, the Equatable trait provides defaults for \_\_eq\_\_() and \_\_ne\_\_() when all fields are Equatable. For types without a natural ordering (like complex numbers), only implement Equatable, and not Comparable. === Subscript operators Implement \_\_getitem\_\_() for reads and \_\_setitem\_\_() for writes. Both subscripting methods accept variadic arguments for multi-dimensional indexing. For a simple one-dimensional collection, you unlock subscripting with a simple index: struct MySeq\[T: Copyable\]: def \_\_getitem\_\_(self, idx: Int) -\> T: ... def \_\_setitem\_\_(mut self, idx: Int, value: T): ... For multi-dimensional collections, make use of variadics or multiple index arguments: struct Grid\[T: Copyable\]: \# Fixed two dimensions def \_\_getitem\_\_(self, x: Int, y: Int) -\> T: ... \# Arbitrary dimensions def \_\_getitem\_\_(self, \*indices: Int) -\> T: ... Custom subscripts can support slicing as well as indices, such as obj\[1:5\]. Implement \_\_getitem\_\_() with a #link("https://mojolang.org/docs/std/builtin/builtin_slice/Slice/")[Slice] parameter instead of Int. Each Slice has three optional fields: start, end, and step. You normalize these by calling indices(). Pass your type's size. This returns a triplet of values representing the span adjusted to your extent, resolving omitted values or negative indices into non-negative positions: struct MySeq\[T: Copyable\]: var size: Int def \_\_getitem\_\_(self, span: Slice) -\> Self: var start: Int var end: Int var step: Int start, end, step = span.indices(self.size) ... === Walkthrough: Build a Complex type The next sections incrementally build a Complex struct. This example demonstrates every category of operator implementation. In this walk-through, you'll work with unary operators, binary operators with same-type and mixed-type operands, reversed methods, in-place assignment, equality comparison, and subscript access. #notebox("Note", rgb("#8a94a6"), rgb("#556666"), rgb("#f7f8fa"))[ The standard library includes #link("https://mojolang.org/docs/std/complex/complex/ComplexSIMD/")[ComplexSIMD], a parameterized complex number type with basic arithmetic support. The Complex type in this example is independent and not based on ComplexSIMD. ] === Create the base type A complex number holds real and imaginary parts, stored in the real re and imaginary im fields: from std.math import sqrt \@fieldwise\_init struct Complex( Equatable, Writable, TrivialRegisterPassable, ): var re: Float64 var im: Float64 - Conforming to TrivialRegisterPassable gives you value semantics without needing to write special lifecycle methods. - Equatable lets you compare two instances, and Writable produces output for print() statements. ==== Convenience initializer Adding a convenience initializer lets you create instances using only the real part of your number: def \_\_init\_\_(out self, re: Float64): self.re = re self.im = 0.0 === Make your type printable Implementing Writable lets you use print() and String() directly. This custom implementation provides parentheses and separate real and imaginary output. \# Struct method def write\_to(self, mut writer: Some\[Writer\]): writer.write("(", self.re) if self.im \< 0: writer.write(" - ", -self.im) else: writer.write(" + ", self.im) writer.write("i)") ... c = Complex(3.14, -2.72) print(c) \# (3.14 - 2.72i) === Add unary operator support +c returns the value unchanged. -c negates both components: \# methods def \_\_pos\_\_(self) -\> Self: return self def \_\_neg\_\_(self) -\> Self: return Self(-self.re, -self.im) ... c = Complex(-1.2, 6.5) print(+c) \# (-1.2 + 6.5i) print(-c) \# (1.2 - 6.5i) === Support binary arithmetic Add addition, subtraction, multiplication, and division between two Complex values with dunders. Each form returns a new Complex instance: def \_\_add\_\_(self, rhs: Self) -\> Self: return Self(self.re + rhs.re, self.im + rhs.im) def \_\_sub\_\_(self, rhs: Self) -\> Self: return Self(self.re - rhs.re, self.im - rhs.im) def \_\_mul\_\_(self, rhs: Self) -\> Self: return Self( self.re \* rhs.re - self.im \* rhs.im, self.re \* rhs.im + self.im \* rhs.re, ) def \_\_truediv\_\_(self, rhs: Self) -\> Self: denom = rhs.squared\_norm() return Self( (self.re \* rhs.re + self.im \* rhs.im) / denom, (self.im \* rhs.re - self.re \* rhs.im) / denom, ) def squared\_norm(self) -\> Float64: return self.re \* self.re + self.im \* self.im def norm(self) -\> Float64: return sqrt(self.squared\_norm()) c1 = Complex(-1.2, 6.5) c2 = Complex(3.14, -2.72) print(c1 + c2) \# (1.94 + 3.78i) print(c1 \* c2) \# (13.91 + 23.67i) === Add mixed-type arithmetic with reversed methods To support expressions like 2.5 + c where Float64 is on the left, you need both overloaded forward methods and reversed methods. Without \_\_radd\_\_(), 2.5 + c would fail because Float64 doesn't know about Complex: \# Forward: Complex + Float64 def \_\_add\_\_(self, rhs: Float64) -\> Self: return Self(self.re + rhs, self.im) \# Reversed: Float64 + Complex def \_\_radd\_\_(self, lhs: Float64) -\> Self: return Self(self.re + lhs, self.im) def \_\_sub\_\_(self, rhs: Float64) -\> Self: return Self(self.re - rhs, self.im) def \_\_rsub\_\_(self, lhs: Float64) -\> Self: return Self(lhs - self.re, -self.im) def \_\_mul\_\_(self, rhs: Float64) -\> Self: return Self(self.re \* rhs, self.im \* rhs) def \_\_rmul\_\_(self, lhs: Float64) -\> Self: return Self(lhs \* self.re, lhs \* self.im) def \_\_truediv\_\_(self, rhs: Float64) -\> Self: return Self(self.re / rhs, self.im / rhs) def \_\_rtruediv\_\_(self, lhs: Float64) -\> Self: denom = self.squared\_norm() return Self( (lhs \* self.re) / denom, (-lhs \* self.im) / denom, ) Now both orderings work: c = Complex(-1.2, 6.5) print(c + 2.5) \# (1.3 + 6.5i) print(2.5 + c) \# (1.3 + 6.5i) print(2.5 \* c) \# (-3.0 + 16.25i) ==== Allow in-place assignment In-place methods modify self directly instead of returning a new value. You can overload for both Complex and Float64 operands: def \_\_iadd\_\_(mut self, rhs: Self): self.re += rhs.re self.im += rhs.im def \_\_iadd\_\_(mut self, rhs: Float64): self.re += rhs def \_\_isub\_\_(mut self, rhs: Self): self.re -= rhs.re self.im -= rhs.im def \_\_isub\_\_(mut self, rhs: Float64): self.re -= rhs def \_\_imul\_\_(mut self, rhs: Self): var new\_re = self.re \* rhs.re - self.im \* rhs.im var new\_im = self.re \* rhs.im + self.im \* rhs.re self.re = new\_re self.im = new\_im def \_\_imul\_\_(mut self, rhs: Float64): self.re \*= rhs self.im \*= rhs def \_\_itruediv\_\_(mut self, rhs: Self): var denom = rhs.squared\_norm() var new\_re = (self.re \* rhs.re + self.im \* rhs.im) / denom var new\_im = (self.im \* rhs.re - self.re \* rhs.im) / denom self.re = new\_re self.im = new\_im def \_\_itruediv\_\_(mut self, rhs: Float64): self.re /= rhs self.im /= rhs ... c = Complex(-1.0, -1.0) c += Complex(0.5, -0.5) print(c) \# (-0.5 - 1.5i) c += 2.75 print(c) \# (2.25 - 1.5i) c \*= 0.75 print(c) \# (1.6875 - 1.125i) c /= 2.0 print(c) \# (0.84375 - 0.5625i) === Support type equality checks Complex numbers have no natural ordering, so you should implement Equatable (not Comparable). This gives you == and != without implying that one complex number is "less than" another: def \_\_eq\_\_(self, other: Self) -\> Bool: return (self.re == other.re and self.im == other.im) \# Conforming to Equatable provides this automatically def \_\_ne\_\_(self, other: Self) -\> Bool: return not (self == other) c1 = Complex(-1.2, 6.5) c2 = Complex(-1.2, 6.5) c3 = Complex(3.14, -2.72) print(c1 == c2) \# True print(c1 != c3) \# True === Unlock subscript access The get and set item dunders allow you to index content within your type. For this example, the real part of the complex number is index 0, and index 1 returns the imaginary component: def \_\_getitem\_\_(self, idx: Int) raises -\> Float64: if idx == 0: return self.re if idx == 1: return self.im raise "index out of bounds" def \_\_setitem\_\_( mut self, idx: Int, value: Float64 ) raises: if idx == 0: self.re = value elif idx == 1: self.im = value else: raise "index out of bounds" ... c = Complex(3.14) print(c\[0\], c\[1\]) \# 3.14 0.0 c\[1\] = 42.0 print(c) \# (3.14 + 42.0i) === Every operator, one walkthrough This example walked you through every Mojo operator from simple arithmetic to comparisons to subscripting. Implementing the right dunders and/or conforming to the right traits enables you to use operator syntax for nearly any custom type.