We have traits at home

I have had more time recently to write code in Zig[1], and it has been good for my soul. I wanted to share some insights I have arrived at regarding abstractions in Zig, including some techniques that I have found useful but not seen employed more broadly.

Considerable[2] ink[3] has[4] been[5] spilled[6] on Zig Interfaces, and justifiably so. While not a language feature, interfaces are the de-facto way of implementing runtime polymorphism in Zig. Yet[7] more[8] has been said about zig's comptime, and how it may be leveraged to avoid the performance woes of dynamic dispatch and write slick, data-oriented code.

To put it briefly, polymorphism in Zig is conventionally obtained in one of two ways:

This state of affairs has proven to be quite effective so far, but it is not without its tradeoffs.

Interfaces offer the sort of type contracts that we are used to from other languages. A function that accepts an interface has a guarantee of its method signatures, but not its implementation. However because interfaces are, by design, modular at runtime, they also introduce a hard barrier for the compiler and thus a performance overhead. There is enough indirection that it often cannot effectively apply the optimizations that it would to monomorphic code.

For cases of performance critical code, Zig code typically relies on comptime duck typing. A function argument is either anytype or parameterized by a type argument T, and the function implementation will either perform compile-time reflection or make assumptions about the type's methods that can be checked at compile-time. Implementations can then be monomorphised, similar to C++ templating or C macros. This is slightly more safe than runtime duck typing, because any interface violations will be caught at compile time rather than runtime, but the caller and callee must still be careful to document and follow the contract of what types may be passed.

If you are coming from another strongly typed language, you would be forgiven for finding this arrangement rather simplistic. This is perhaps by design -- Zig aims to be simple. However if you have grown accustomed to Rust's traits or Haskell typeclasses, there appears to be no ready equivalent in Zig.

I will argue that, with a little adaptation, it is perfectly possible to express similar abstractions over types in Zig userland. I would propose a new convention for defining compile-time type contracts as Traits in Zig. It doesn't require any libraries or machinery - in fact if fits remarkably well into vanilla Zig. It is inspired by the functional programming world in general, and OCaml in particular.

Definition: a Trait is a compile-time struct which contains related functions over a given type. [9]

So as not to bury the lead, let me begin with a brief example of how we might define a trait for folding (reducing) on a type.

/// T forms a Reducer under an empty value and a combine operation fn Reducer(T: type) type { return struct { empty: T, combine: fn (T, T) T,

fn reduce(self: @This(), xs: []const T) T { var accum = self.empty; for (xs) |x| accum = self.combine(accum, x); return accum; } }; }

/// i32 forms a sum reducer under 0 and addition const sum_trait = impl(Reducer(i32), struct { pub const empty = 0; pub fn combine(x: T, y: T) T { return x + y; } });

/// i32 forms a product reducer under 1 and multiplication const product_trait = impl(Reducer(i32), struct { pub const empty = 1; pub fn combine(x: T, y: T) T { return x * y; } });

export fn run(buf: [*c]i32, n: usize) i32 { const sum = sum_trait.reduce(buf[0..n]); const product = product_trait.reduce(buf[0..n]); return sum + product; }

Ugh. Isn't that just buttery smooth?[10] The contract is clear, LSP autocompletion picks it up, the compiler even knows to inline it and unroll the loop.

Traits are just bags of functions and types that tell the receiving function how to execute subtasks that it's not responsible for defining. They behave a lot like vtables, except that they are evaluated at comptime rather than runtime.

To justify traits, let's take a tour through Zig polymorphism with a common motivating example: arithmetic operations. I've seen a few complaints over the years about how awkward is it to do math in Zig over user-defined types. In its pursuit of explicit control flow, Zig has eschewed operator overloading, and this can be jarring if you're coming from other typed languages.

Let's imagine we have some user-defined types like Rational and Complex, and we want to define their basic arithmentic operations. We would probably end up with something like the following.

const Rational = struct { numerator: i32, denominator: i32,

fn add(lhs: Rational, rhs: Rational) Rational {...} fn sub(lhs: Rational, rhs: Rational) Rational {...} fn mul(lhs: Rational, rhs: Rational) Rational {...} fn div(lhs: Rational, rhs: Rational) Rational {...} };

const Complex = struct { real: f32, imaginary: f32,

fn add(lhs: Complex, rhs: Complex) Complex {...} fn sub(lhs: Complex, rhs: Complex) Complex {...} fn mul(lhs: Complex, rhs: Complex) Complex {...} fn div(lhs: Complex, rhs: Complex) Complex {...} };

For the sake of having something to do with them, let's say that we want to compute the Möbius transformation: (a * x + b) / (c * x + d), which is useful for complex numbers as well as rationals. To keep things simple, we'll assume that all arguments to the function are of the same type.

If we were to implement the Mobius transformation naively, we could just use built-in operators.

fn mobius(T: type, x: T, a: T, b: T, c: T, d: T) T { return (a * x + b) / (c * x + d); }

This would work fine, as long as the inputs were all floats of the same type. Division would break for integers, and of course it wouldn't work with our user-defined types, but it's a starting point.

We could of course notice that both of our user-defined types have arithmetic functions with matching names and similar signatures, which would lend themselves to duck typing as well.

fn mobius(T: type, x: T, a: T, b: T, c: T, d: T) T { const numerator = a.mul(x).add(b); const denominator = c.mul(x).add(d); return numerator.div(denominator); }

This will work for our user-defined types, but obviously not with Zig's primitive types. We would have to wrap those in our own user-defined types with their own functions. More troublingly, this sort of ties us to a sort of object-oriented pattern: the de-facto way to obey this abstraction contract is to define methods on the struct.

The good news is that it's all monomorphised at compile time, allowing the compiler to simplify and inline these function calls.

If we wanted this to work dynamically at runtime, we could perhaps declare a Num interface that uses a VTable and opaque pointer to dynamically dispatch these math operations.

const Num = struct { value: [32]u8, vtable: *const VTable,

const VTable = struct { add: const fn([]u8, []u8) Num, sub: const fn([]u8, []u8) Num, mul: const fn([]u8, []u8) Num, div: const fn([]u8, []u8) Num, };

fn add(lhs: Num, rhs: Num) Num { ... } fn sub(lhs: Num, rhs: Num) Num { ... } fn mul(lhs: Num, rhs: Num) Num { ... } fn div(lhs: Num, rhs: Num) Num { ... }

fn mobius(x: Num, a: Num, b: Num, c: Num, d: Num) { const numerator = a.mul(x).add(b); const denominator = c.mul(x).add(d); return numerator.div(denominator); } }

We would then need to wire up Complex and Rational to expose fn num(self: *Self) Num (or some such) that wraps our arithmentic functions, which is a little too tedious to reproduce here, but it would work fine on the happy path. You would of course need to be a bit careful: there's nothing stopping you from mixing complex and rational numbers in your function calls, and this would lead to undefined behavior. You could probably add some saftey checks with a bit of overengineering, but we won't get into that here.

Not only is this clunkier than comptime duck typing, it will also be slower. We are effectively demanding dynamic dispatch whenever we want to do a math operation, which will trip up the compiler and have a significant performance impact if we're doing a lot of math.

When we define the mobius function, we really want to say "I don't know what type I'm dealing with, but it needs to have arithmetic operations." In this sort of case, I think it makes the most sense to just pass the functions it needs to use at compile-time. This would remove any ambiguity, at the cost of being rather explicit.

fn mobius( comptime T: type, comptime add: fn (T, T) T, comptime mul: fn (T, T) T, comptime div: fn (T, T) T, x: T, a: T, b: T, c: T, d: T, ) T { const numerator = add(mul(a, x), b); const denominator = add(mul(c, x), d); return div(numerator, denominator); }

This is effectively compile-time dependency injection at a really fine-grained level. The signature is rather verbose, but there's actually a lot less boilerplate to write for Complex or Rational compared to runtime interfaces.

mobius(Complex, Complex.add, Complex.mul, Complex.div, x, a, b, c, d); mobius(Real, Real.add, Real.mul, Real.div, x, a, b, c, d);

From here, it's not a huge leap to traits. We're effectively just cleaning up the function signature by gathering together the arithmetic functions into one struct.[11]

fn Num(T: type) type { return struct { add: fn (T, T) T, sub: fn (T, T) T, mul: fn (T, T) T, div: fn (T, T) T, }; }

fn mobius(T: type, num: Num(T), x: T, a: T, b: T, c: T, d: T) T { const numerator = num.add(num.mul(a, x), b); const denominator = num.add(num.mul(c, x), d); return num.div(numerator, denominator); }

From here, I usually just declare the trait implementation as part of my type.

const Complex = struct { // fields... const num = Num(Complex) { .add = add, .sub = sub, .mul = mul, .div = div }; // functions (add, sub, mul, div...) };

mobius(Complex, Complex.num, x, a, b, c, d);

What I really like about this approach is how flexible it is, while remaining very clear about the contract between the function and type. For example, imagine that we wanted to extend our definition of Num to the primitive types like i32 and f32. There are a bunch of ways we could do this, but I would probably just define a function that produces something appropriate for integer and float types.

fn builtinNum(T: type) Num(T) { const inner = struct { fn add(lhs: T, rhs: T) T { return lhs + rhs; } fn sub(lhs: T, rhs: T) T { return lhs - rhs; } fn mul(lhs: T, rhs: T) T { return lhs * rhs; } }; const div = switch (@typeInfo(T)) { .int, .comptime_int => struct { fn div(lhs: T, rhs: T) T { return @divTrunc(lhs, rhs); } }.div, .float, .comptime_float => struct { fn div(lhs: T, rhs: T) T { return lhs / rhs; } }.div, else => @compileError("unsupported type: " ++ @typeName(T)), }; return .{ .add = inner.add, .sub = inner.sub, .mul = inner.mul, .div = div }; }

mobius(f32, builtinNum(f32), x, a, b, c, d);

Traits are also transparent to the compiler, which will happily monomorphize and inline these functions. I prefer this over runtime polymorphism most of the time, especially in hot loops or for library functions that are unopinionated on data layout.

You'll notice that traits have to be explicitly conditioned on a known type. Unlike interfaces, there's no way to say "give me a number, I don't care what kind." For some applications, it would be nice if we could have an associated type.

fn Iter(Elem: type) type { return struct { T: type, next: fn (T) ?Elem, }; }

Unfortunately, the above isn't valid Zig, and it's unclear how that would even work at a language level. Instead, traits insist that the types they are operating on are known in advance.

fn Iter(T: type, Elem: type) type { return struct { next: fn (T) ?Elem, }; }

I have not yet found an ergonomic way around this. Traits do not play very well with type erasure.

Before proceeding, I want to introduce some small helper functions that have s