2

Using comptime to invert bijective functions on enums

 1 year ago
source link: https://zig.news/rbino/using-comptime-to-invert-bijective-functions-on-enums-3pmk
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
Cover image for Using comptime to invert bijective functions on enums

Using comptime to invert bijective functions on enums

Hi everyone! This is my first post on Zig News, and I've decided to break the ice with this bikeshed I ended up into while working on the Advent of Code 2022.

We have the usual "Rock Paper Scissors" game, and given a Shape we need to be able to retrieve the Shape that it beats and the Shape it's beaten by.

pub const Shape = enum(u8) {
    rock,
    paper,
    scissors,

    pub fn beats(self: Shape) Shape {
        return switch (self) {
            .rock => .scissors,
            .paper => .rock,
            .scissors => .paper,
        };
    }

    pub fn beatenBy(self: Shape) Shape {
        return switch (self) {
            .rock => .paper,
            .paper => .scissors,
            .scissors => .rock,
        };
    }
}

This works, but it's very error prone because we don't have any guarantee that the second function is actually the inverse of the first one.

Moreover, if the "Rock Paper Scissors" ISO Committee decides to change the rules for RockPaperScissors23 (now with Atomics!), we have two places where our code needs to change.

Can we do better? Using the power of comptime, I think we can.

pub const Shape = enum(u8) {
    rock,
    paper,
    scissors,

    pub fn beats(self: Shape) Shape {
        return switch (self) {
            .rock => .scissors,
            .paper => .rock,
            .scissors => .paper,
        };
    }

    pub fn beatenBy(self: Shape) Shape {
        switch (self) {
            inline else => |shape| {
                const winner = comptime blk: {
                    inline for (@typeInfo(Shape).Enum.fields) |field| {
                        const other_shape = @intToEnum(Shape, field.value);
                        if (other_shape.beats() == shape) {
                            break :blk other_shape;
                        }
                    }
                    unreachable;
                };
                return winner;
            },
        }
    }
}

As you can see, now beatenBy is generated at comptime starting from beats, and this allows us to concentrate all changes only in the beats function.

In fact, we can also generalize the concept and create a comptime function that returns the inverse of an arbitrary bijective function from an enum domain T1 to an enum codomain T2:

fn invert(comptime T1: type, comptime T2: type, comptime originalFn: fn (a: T1) T2) fn (T2) T1 {
    if (@typeInfo(T1).Enum.fields.len != @typeInfo(T2).Enum.fields.len) {
        @compileError("Trying to invert a non-bijective function: domain " ++
            @typeName(T1) ++ " and codomain " ++ @typeName(T2) ++ " have different sizes");
    }

    return struct {
        fn function(self: T2) T1 {
            switch (self) {
                inline else => |input| {
                    const inverse = comptime blk: {
                        inline for (@typeInfo(T1).Enum.fields) |field| {
                            const candidate = @intToEnum(T1, field.value);
                            if (originalFn(candidate) == input) {
                                break :blk candidate;
                            }
                        }
                        @compileError("Trying to invert a non-bijective function: " ++
                            @tagName(input) ++ " is not contained in the codomain");
                    };
                    return inverse;
                },
            }
        }
    }.function;
}

This reduces the creation of beatenBy to

const beatenBy = invert(Shape, Shape, beats);

Using this also gives us an additional advantage: the comptime guarantee that the function is bijective. You can use this code to test this quickly:

pub const Shape = enum {
    square,
    triangle,
    circle,

    pub fn toColor(self: Shape) Color {
        return switch (self) {
            .square => .red,
            .triangle => .green,
            .circle => .blue,
        };
    }
};

pub const Color = enum {
    red,
    green,
    blue,

    pub const toShape = invert(Shape, Color, Shape.toColor);
};

If you make two different Shapes return the same color in toColor, the compilation will fail with:

src/main.zig:20:25: error: Trying to invert a non-bijective function: blue is not contained in the codomain
                        @compileError("Trying to invert a non-bijective function: " ++
                        ^~~~~~~~~~~~~

And if you add an extra element either to the Shape enum or to the Color enum, the compilation will fail with:

src/main.zig:5:9: error: Trying to invert a non-bijective function: domain main.Shape and codomain main.Color have different sizes
        @compileError("Trying to invert a non-bijective function: domain " ++
        ^~~~~~~~~~~~~

I'm aware that the "Rock Paper Scissors" problem could've been solved in a different way (e.g. representing the "beats" relationship with a circular buffer and looking at the element after or before you), but I took the occasion to make myself a little more comfortable with comptime.

Let me know if you have any suggestion or correction in the comments and happy Zigging!


About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK