Ad Hoc Polymorphism in Zig

· ·

  • Fixed bugs in the code: wrong @typeName arguments, missing semicolons, a wrong signature.
  • Explained generic return types, and that argument types must match exactly.
  • Added the maths.
  • Fixed typos.
  • Written for Zig 0.14; the type reflection API has changed since.

LetsLet's say we have a function such as the following.

fn dot(args: anytype) i32 {
    // compute dot product...
}

We want this function to operate on two Point2D or two Point3D structs. But we do not want to expose 2 different functions for this, i.ei.e., something like the following.

pub fn dot2d(p1: Point2D, p2: Point2D) i32 { // ... }
pub fn dot2d(p1: Point2D, p2: Point2D) i32 {
    // ...
}

pub fn dot3d(p1: Point3D, p2: Point3D) i32 { // ... }
pub fn dot3d(p1: Point3D, p2: Point3D) i32 {
    // ...
}

Both compute the same thing, just for a different number of dimensions.

p⋅q=∑i=1npiqi,n∈{2,3}\mathbf{p} \cdot \mathbf{q} = \sum_{i=1}^{n} p_i q_i, \qquad n \in \{2, 3\}

Rather, we would like our function dot to select the correct implementation based on the type and length of args. The more technical term for this is "Ad hoc polymorphism". In type theory terms, dot should have both function types at once, an intersection type:

dot:(P2×P2→Z)∩(P3×P3→Z)\mathrm{dot} : (P_2 \times P_2 \to \mathbb{Z}) \cap (P_3 \times P_3 \to \mathbb{Z})

Zig does not support ad hoc polymorphism, but using comptime and compile time type reflection, we can sort of mimickmimic this. What we essentially want is a function that takes two arguments, the arguments we are calling our function with, and a list of functions which are our different implementations of this method. The function signature should look like something likesomething like this

fn adHocPoly(args: anytype, fns: anytype) fn () typed.ReturnType(fns[0])
fn adHocPoly(args: anytype, fns: anytype) typed.ReturnType(fns[0])

typed.ReturnType is a utility type part of a small library I made, I have provided the implementation below if you want to use it in your own project.

const Type = @import("std").builtin.Type;
pub inline fn ReturnType(comptime func: anytype) type {
    comptime {
        const type_info: Type = @typeInfo(@TypeOf(func));

        return type_info.@"fn".return_type orelse {
            @compileError("Function has no return type. This should not be possible but is possible due to the current Zig spec.");
            // A null return type means it is generic: it depends on the arguments,
            // e.g. `fn f(x: anytype) @TypeOf(x)`, so it is unknown until the call.
            @compileError("function has a generic return type, it depends on its arguments");
        };
    }
}

You might notice a limitation on our function signature, we assume that every implementation has the same return type. LetsLet's lay out what we want this function to do. Given some arguments of any type, we want to pick the function in the list of functions supplied, whose parameters match the arguments, in order, type and length. If no functions parameters is a match, we want to error (specifially compile error).If no function's parameters match, we want a compile error. Put formally, given functions f1,…,fmf_1, \dots, f_m and arguments a1,…,ana_1, \dots, a_n, we call the first fkf_k whose parameter types are exactly the argument types:

k=min⁡{ j∣fj:T1×⋯×Tn→R,  ai:Ti for all i }k = \min \{\, j \mid f_j : T_1 \times \dots \times T_n \to R,\ \ a_i : T_i \text{ for all } i \,\}

At this point I would recommend the reader to try to implement this function themselves as an exercise, nevertheless, below is the implementation of this function.

pub fn adHocPoly(args: anytype, fns: anytype) typed.ReturnType(fns[0]) {
    const RT = typed.ReturnType(fns[0]);
    comptime for (fns) |func| {
        if (typed.ReturnType(func) != RT) {
            @compileError("all functions must have the same return type");
        }
    };

    const args_ty = @typeInfo(@TypeOf(args));
    const found_idx = blk: switch (args_ty) {
        inline .@"struct" => |s| {
            const fields = s.fields;
            const fns_ty = @typeInfo(@TypeOf(fns));
            switch (fns_ty) {
                inline .@"struct" => |fns_struc| {
                    comptime for (fns_struc.fields, 0..) |func, fn_idx| {
                        const fn_ty = @typeInfo(func.type);
                        switch (fn_ty) {
                            inline .@"fn" => |fn_ty_info| {
                                if (fn_ty_info.params.len != fields.len) {
                                    continue;
                                }

                                var found = true;
                                for (fn_ty_info.params, 0..) |param, i| {
                                    if (param.type) |param_ty| {
                                        if (param_ty != fields[i].type) {
                                            found = false;
                                            break;
                                        }
                                    } else {
                                        @compileError("generic fn parameter of anytype found, can not deduce which fn to call");
                                    }
                                }
                                if (found) {
                                    break :blk fn_idx;
                                }
                            },
                            else => @compileError("functions params was not an array of functions, was " ++ @typeName(@Type(fn_ty))),
                            else => @compileError("fns must be a tuple of functions, found " ++ @typeName(func.type)),
                        }
                    };

                    @compileError("no function which fulfills type signature of args was provided");
                },
                else => @compileError("functions was not a struct of fn, was " ++ @typeName(@Type(fns_ty))),
                else => @compileError("fns must be a tuple of functions, was " ++ @typeName(@TypeOf(fns))),
            }
        },
        else => @compileError("expected struct for args, was " ++ @typeName(@Type(args))),
        else => @compileError("expected a tuple for args, was " ++ @typeName(@TypeOf(args))),
    };

    return @call(.auto, fns[found_idx], args);
}

Now we can implement our ad hoc polymorphic dot function.

const Point2D = struct {
    x: i32,
    y: i32,
};

fn dot2d(p1: Point2D, p2: Point2D) i32 {
    return p1.x * p2.x + p1.y * p2.y;
}

const Point3D = struct {
    x: i32,
    y: i32,
    z: i32,
};

fn dot3d(p1: Point3D, p2: Point3D) i32 {
    return p1.x * p2.x + p1.y * p2.y + p1.z * p2.z;
}

fn dot(args: anytype) i32 {
    return adHocPoly(
        args,
        .{ dot2d, dot3d },
    );
}

test "ad_hoc_polymorphism" {
    try testing.expectEqual(11, dot(.{
        Point2D{ .x = 1, .y = 2 },
        Point2D{ .x = 3, .y = 4 },
    }));

    try testing.expectEqual(32, dot(.{
        Point3D{ .x = 1, .y = 2, .z = 3 },
        Point3D{ .x = 4, .y = 5, .z = 6 },
    }));
}

Note that the parameter types must match the argument types exactly, no coercion happens. Integer literals have the type comptime_int, so this fails to compile with "no function which fulfills type signature of args was provided", even though calling add(1, 2) directly works.

fn add(a: i32, b: i32) i32 {
    return a + b;
}

const sum = adHocPoly(.{ 1, 2 }, .{add}); // compile error
const sum = adHocPoly(.{ @as(i32, 1), @as(i32, 2) }, .{add}); // ok

Suppose we would like to also support multiple return types as well. To do that we would have to modify the API of our polymorphic function to take the expected return type as an argument.

pub fn adHocPoly(args: anytype, fns: anytype) typed.ReturnType(fns[0]) {
pub fn adHocPoly(comptime T: type, args: anytype, fns: anytype) T {
  const RT = typed.ReturnType(fns[0]);
  comptime for (fns) |func| {
      if (typed.ReturnType(func) != RT) {
          @compileError("all functions must have the same return type");
      }
  }

  const args_ty = @typeInfo(@TypeOf(args));
  const found_idx = blk: switch (args_ty) {
      inline .@"struct" => |s| {
          const fields = s.fields;
          const fns_ty = @typeInfo(@TypeOf(fns));
          switch (fns_ty) {
              inline .@"struct" => |fns_struc| {
                  comptime for (fns_struc.fields, 0..) |func, fn_idx| {
                      const fn_ty = @typeInfo(func.type);
                      switch (fn_ty) {
                          inline .@"fn" => |fn_ty_info| {
                              if (fn_ty_info.params.len != fields.len) {
                                continue;
                              }
                              if (fn_ty_info.return_type) |RT| {
                                if (RT != T) {
                                  continue;
                                }
                              } else {
                                unreachable;
                                @compileError("generic return type found, can not compare it to T");
                              }

                              var found = true;
                              for (fn_ty_info.params, 0..) |param, i| {
                                  if (param.type) |param_ty| {
                                      if (param_ty != fields[i].type) {
                                          found = false;
                                          break;
                                      }
                                  } else {
                                      @compileError("generic fn parameter of anytype found, can not deduce which fn to call");
                                  }
                              }
                              if (found) {
                                  break :blk fn_idx;
                              }
                          },
                          else => @compileError("functions params was not an array of functions, was " ++ @typeName(@Type(fn_ty))),
                          else => @compileError("fns must be a tuple of functions, found " ++ @typeName(func.type)),
                      }
                  }
                  };

                  @compileError("no function which fulfills type signature of args was provided");
              },
              else => @compileError("functions was not a struct of fn, was " ++ @typeName(@Type(fns_ty))),
              else => @compileError("fns must be a tuple of functions, was " ++ @typeName(@TypeOf(fns))),
          }
      },
      else => @compileError("expected struct for args, was " ++ @typeName(@Type(args))),
      else => @compileError("expected a tuple for args, was " ++ @typeName(@TypeOf(args))),
  }
  };

  return @call(.auto, fns[found_idx], args);
}

The return type is now passed in by the caller, so dot becomes

fn dot(args: anytype) i32 {
    return adHocPoly(i32, args, .{ dot2d, dot3d });
}

You can find this entire exercise on my github for reference.