Open multiple dispatch in zig

Published 2020-04-28

The Zig stdlib often uses open single dispatch eg:

// in stdlib

pub fn serialize(self: *Serializer, value: var) !void {
    const T = comptime @TypeOf(value);
    if (comptime std.meta.trait.hasFn("serialize")(T)) {
        return T.serialize(value, self);
    } else {
        self.defaultSerialize(value);
    }
}

// in my code

struct Foo {
    foo: usize,

    fn serialize(self: Foo, serialize: *Serializer) !void {
        ...
    }
}

There are a bunch of things going on here:

The result is a kind of duck-typed meta-programming that is type-checked at specialization-time.

This is really powerful, but there are two things I would like to improve on:

  1. It only allows choosing the function based on the type of the first argument (single dispatch)
  2. The associated functions all live in one namespace, so it's easy to have accidental collisions

Julia has a very similar compilation model (aside from also allowing runtime compilation). But Julia's functions allow multiple dispatch and belong to the namespace they were declared in, rather than the namespaces of all the types they are declared on. This enabled a fantastically composable ecosystem of libraries. Can Zig do that too?


We can certainly do closed multiple dispatch:

fn add_impl(comptime a: type, comptime b: type) var {
    if (a == Vec2 and b == Vec2) {
        return add_vec_vec;
    }
    if (a == Mat2 and b == Mat2) {
        return add_mat_mat;
    }
    if (a == Mat2 and b == Diag2) {
        return add_mat_diag;
    }
    panic("No impl for add({}, {})", .{a, b});
}

fn add(a: var, b: var) var {
    const impl = comptime add_impl(@TypeOf(a), @TypeOf(b));
    return impl(a,b);
}

Here add_impl can run arbitrary logic at compile time to decide which method it wants to use. But other packages can't extend this logic to support new types - that's what 'closed' means. We'll get to 'open' in a few more steps.

What if we had two versions of add_impl and we wanted to combine them?

fn add_impl1(comptime a: type, comptime b: type) var {
    if (a == Vec2 and b == Vec2) {
        return add_vec_vec;
    }
    return null;
}

fn add_impl2(comptime a: type, comptime b: type) var {
    if (a == Mat2 and b == Mat2) {
        return add_mat_mat;
    }
    if (a == Mat2 and b == Diag2) {
        return add_mat_diag;
    }
    return null;
}

fn dispatch(comptime name: []const u8, comptime impls: var, a: type, b: type) var {
    for (@typeInfo(@TypeOf(impls)).Struct.fields) |field| {
        if (@field(impls, field.name)(a, b)) |impl| {
            return impl;
        }
    }
    std.debug.panic("no impl for {}({}, {})", .{name, a, b});
}

fn add(a: var, b: var) var {
    const impl = comptime dispatch("add", .{ add_impl1, add_impl2 }, @TypeOf(a), @TypeOf(b) );
    return impl(a,b);
}

What's happening here is that .{ add_impl1, add_impl2 } defines a tuple and dispatch iterates over the fields of the tuple and calls each in turn until one of them returns a function.

(We could choose other logic here eg returning an error if more than one impl returns a function. This is just the easiest dispatch algorithm to demonstrate.)

But the list of impls is still hard-coded in add. The final step is move it out of the library.

fn add(a: var, b: var) var {
    const impl = comptime dispatch("add", @import("root").add_impls, @TypeOf(a), @TypeOf(b) );
    return impl(a,b);
}

@import("root") imports the top-most file in the compilation unit. So when writing an application that uses this math library, you can just put:

const add_impls = .{
  math.add_impl,
  my_lib.add_impl,
  some_other_lib.add_impl,
};

The application author gets to decide in which order to try dispatching, and can add their own dispatch functions at the top of the list to override the choices of other libraries. Like Julia, we probably still want to avoid wanton type piracy but at least we have the tools to fix it if it happens.


Ok, so I lied a little. The code above uses return type inference which hasn't been implemented yet. This code actually works in zig today by adding a second layer of boilerplate to choose the return types. If compiled with zig build-exe test.zig --release-fast -femit-llvm-ir we can open test.ll and confirm that not only was the dispatch compiled away, the math got constant folded away and all that is left is the code to print the answer.


This solves problem 1 (multiple dispatch), but only partially addresses problem 2 (name collisions). We can could go full C and make the top level name something like my_cool_math_lib_add_impls but that still causes problems if we want to import two different versions of the same library.

One neat way to address this would be replace @import("root") with variables that can only be set once, and only at compile-time. Something like:

// in math lib

once add_impls;

// in root

const math1 = @import("v1/math.zig");
const math2 = @import("v2/math.zig");

once math1.add_impls = .{ .. };
once math2.add_impls = .{ .. };