Zig Comptime Forth

forth.zig

const std = @import("std");

test "forwarding a value" {
    const result = forth("value", .{
        .value = @as(u32, 42),
    });

    std.testing.expectEqual(@as(u32, 42), result);
}

test "basic addition" {
    const result = forth("a b +", .{
        .a = @as(u32, 10),
        .b = @as(u32, 20),
    });

    std.testing.expectEqual(@as(u32, 30), result);
}

test "peer type resolution" {
    const result = forth("a b +", .{
        .a = @as(u32, 10),
        .b = @as(u8, 20),
    });

    std.testing.expectEqual(@as(u32, 30), result);
}

test "multiple additions with peer type" {
    const result = forth("a b c + +", .{
        .a = @as(u24, 10),
        .b = @as(u7, 20),
        .c = @as(u9, 30),
    });

    std.testing.expectEqual(@as(u24, 60), result);
}

test "basic subtraction" {
    const result = forth("a b -", .{
        .a = @as(u32, 20),
        .b = @as(u32, 10),
    });

    std.testing.expectEqual(@as(u32, 10), result);
}

test "output" {
    forth("msg printstr a b - .", .{
        .msg = "20 - 7 = ",
        .a = @as(u32, 20),
        .b = @as(u32, 7),
    });
}

pub fn Forth(comptime code: []const u8, environment: anytype) type {
    const env = @TypeOf(environment);

    const type_stack = MakeTypeStack(.{});

    const return_type = blk: { // phase 1: analyzing a type set
        comptime var iterator = std.mem.tokenize(code, "\r\n\t ");
        inline while (comptime iterator.next()) |word| {
            if (comptime std.mem.eql(u8, word, "+")) {
                const rhs = comptime type_stack.pop(.{});
                const lhs = comptime type_stack.pop(.{});

                const ptr = @TypeOf(@as(lhs, 0) + @as(rhs, 0));

                comptime type_stack.push(.{ptr});
            } else if (comptime std.mem.eql(u8, word, "-")) {
                const rhs = comptime type_stack.pop(.{});
                const lhs = comptime type_stack.pop(.{});

                const ptr = @TypeOf(@as(lhs, 0) - @as(rhs, 0));

                comptime type_stack.push(.{ptr});
            } else if (comptime std.mem.eql(u8, word, ".")) {
                _ = comptime type_stack.pop(.{});
            } else if (comptime std.mem.eql(u8, word, "printstr")) {
                _ = comptime type_stack.pop(.{});
            } else if (comptime @hasField(env, word)) {
                comptime type_stack.push(.{@TypeOf(@field(environment, word))});
            } else {
                @compileError("Unknown word '" ++ word ++ "'!");
            }
        }

        break :blk switch (comptime type_stack.size(.{})) {
            0 => void,
            1 => comptime type_stack.pop(.{}),
            else => @compileError("Stack imbalance!"),
        };
    };

    const stack_slot = blk: {
        comptime var fields: []const std.builtin.TypeInfo.UnionField = &[_]std.builtin.TypeInfo.UnionField{};

        inline for (type_stack.types(.{})) |t, i| {
            fields = fields ++ &[_]std.builtin.TypeInfo.UnionField{
                std.builtin.TypeInfo.UnionField{
                    .name = std.fmt.comptimePrint("{}", .{i}),
                    .field_type = t,
                    .alignment = @alignOf(t),
                },
            };
        }

        break :blk @Type(.{
            .Union = .{
                .layout = .Auto,
                .tag_type = null,
                .fields = fields,
                .decls = &[_]std.builtin.TypeInfo.Declaration{},
            },
        });
    };

    return struct {
        const TypeStack = type_stack;
        const ReturnType = return_type;
        const StackSlot = stack_slot;
        const Environment = env;
    };
}

pub fn forth(comptime code: []const u8, environment: anytype) Forth(code, environment).ReturnType {
    const Instance = Forth(code, environment);
    const TypeStack = Instance.TypeStack;

    var stack: [TypeStack.max_size(.{})]Instance.StackSlot = undefined;

    // phase 2: running the actual code

    // @compileLog("run", code);

    comptime var stack_balance: usize = 0;
    comptime var iterator = std.mem.tokenize(code, "\r\n\t ");
    inline while (comptime iterator.next()) |word| {
        // @compileLog(word, stack_balance);
        if (comptime std.mem.eql(u8, word, "+")) {
            const RHS = comptime TypeStack.pop(.{});
            const LHS = comptime TypeStack.pop(.{});

            comptime stack_balance -= 2;

            const lhs = @field(stack[stack_balance + 0], comptime Instance.TypeStack.field(.{LHS}));
            const rhs = @field(stack[stack_balance + 1], comptime Instance.TypeStack.field(.{RHS}));

            const result = lhs + rhs;
            const T = @TypeOf(result);

            stack[stack_balance] = @unionInit(Instance.StackSlot, comptime Instance.TypeStack.field(.{T}), result);

            comptime stack_balance += 1;

            comptime TypeStack.push(.{T});
        } else if (comptime std.mem.eql(u8, word, "-")) {
            const RHS = comptime TypeStack.pop(.{});
            const LHS = comptime TypeStack.pop(.{});

            comptime stack_balance -= 2;

            const rhs = @field(stack[stack_balance + 1], comptime Instance.TypeStack.field(.{RHS}));
            const lhs = @field(stack[stack_balance + 0], comptime Instance.TypeStack.field(.{LHS}));

            const result = lhs - rhs;
            const T = @TypeOf(result);

            stack[stack_balance] = @unionInit(Instance.StackSlot, comptime Instance.TypeStack.field(.{T}), result);

            comptime stack_balance += 1;

            comptime TypeStack.push(.{T});
        } else if (comptime std.mem.eql(u8, word, ".")) {
            const VAL = comptime TypeStack.pop(.{});

            comptime stack_balance -= 1;

            const value = @field(stack[stack_balance], comptime Instance.TypeStack.field(.{VAL}));

            std.debug.print("{any}\n", .{value});
        } else if (comptime std.mem.eql(u8, word, "printstr")) {
            const VAL = comptime TypeStack.pop(.{});

            comptime stack_balance -= 1;

            const value = @field(stack[stack_balance], comptime Instance.TypeStack.field(.{VAL}));

            std.debug.print("{s}", .{value});
        } else if (comptime @hasField(Instance.Environment, word)) {
            const value = @field(environment, word);
            const T = @TypeOf(value);

            comptime TypeStack.push(.{T});
            stack[stack_balance] = @unionInit(Instance.StackSlot, comptime Instance.TypeStack.field(.{T}), value);
            comptime stack_balance += 1;
        } else {
            @compileError("Unknown word '" ++ word ++ "'!");
        }
    }
    if (Instance.ReturnType == void)
        return
    else
        return @field(stack[0], comptime Instance.TypeStack.field(.{Instance.ReturnType}));
}

fn MakeTypeStack(tag: anytype) type {
    comptime var values: []const type = &[_]type{};
    comptime var contained_types: []const type = &[_]type{};
    comptime var max_depth: usize = 0;
    return struct {
        pub fn size(_: anytype) usize {
            return values.len;
        }

        pub fn max_size(_: anytype) usize {
            return max_depth;
        }

        pub fn types(_: anytype) []const type {
            return contained_types;
        }

        pub fn indexOf(comptime tup: anytype) usize {
            const T: type = tup[0];
            inline for (contained_types) |t, i| {
                if (t == T)
                    return i;
            }
            @compileError("Type " ++ @typeName(T) ++ " not in stack!");
        }

        pub fn field(comptime tup: anytype) []const u8 {
            return std.fmt.comptimePrint("{}", .{indexOf(tup)});
        }

        pub fn push(tup: anytype) void {
            const T: type = tup[0];
            values = values ++ &[_]type{T};
            max_depth = std.math.max(max_depth, values.len);

            const is_contained = inline for (contained_types) |t| {
                if (t == T)
                    break true;
            } else false;
            if (!is_contained) {
                contained_types = contained_types ++ &[_]type{T};
            }
        }

        pub fn pop(_: anytype) type {
            if (values.len == 0) {
                @compileError("Stack underflow");
            } else {
                const limit = values.len - 1;
                const T = values[limit];
                values = values[0..limit];
                return T;
            }
        }
    };
}

test "TypeStack" {
    const stack = MakeTypeStack(.{});

    std.testing.expectEqual(@as(usize, 0), comptime stack.size(.{}));

    comptime stack.push(.{[]const u8});
    comptime stack.push(.{u32});
    comptime stack.push(.{f32});
    comptime stack.push(.{u32});
    comptime stack.push(.{u32});

    std.testing.expectEqual(@as(usize, 5), comptime stack.size(.{}));

    comptime stack.push(.{u8});

    std.testing.expectEqual(@as(usize, 6), comptime stack.size(.{}));

    comptime std.testing.expectEqual(u8, comptime stack.pop(.{}));
    comptime std.testing.expectEqual(u32, comptime stack.pop(.{}));
    comptime std.testing.expectEqual(u32, comptime stack.pop(.{}));
    comptime std.testing.expectEqual(f32, comptime stack.pop(.{}));
    comptime std.testing.expectEqual(u32, comptime stack.pop(.{}));
    comptime std.testing.expectEqual([]const u8, comptime stack.pop(.{}));

    std.testing.expectEqual(@as(usize, 6), comptime stack.max_size(.{}));

    comptime std.testing.expectEqualSlices(type, &[_]type{ []const u8, u32, f32, u8 }, comptime stack.types(.{}));

    std.testing.expectEqual(@as(usize, 0), comptime stack.indexOf(.{[]const u8}));
    std.testing.expectEqual(@as(usize, 1), comptime stack.indexOf(.{u32}));
    std.testing.expectEqual(@as(usize, 2), comptime stack.indexOf(.{f32}));
    std.testing.expectEqual(@as(usize, 3), comptime stack.indexOf(.{u8}));

    std.testing.expectEqualStrings("0", comptime stack.field(.{[]const u8}));
    std.testing.expectEqualStrings("1", comptime stack.field(.{u32}));
    std.testing.expectEqualStrings("2", comptime stack.field(.{f32}));
    std.testing.expectEqualStrings("3", comptime stack.field(.{u8}));
}