Zig FlatBuffers read-only

flatbuffers.zig

///! This package provides a skeleton for read-only access to FlatBuffer encoded binary data.
const std = @import("std");
const assert = std.debug.assert;
const t = std.testing;

/// Header represents the data at offset 0
/// it consists of an unsigned offset to the root table
/// and optionally a 4-byte file identifier
const Header = packed struct {
    offset: u32, // root table offset
    identifier: u32, // optional file identifier

    fn getRootTable(self: *const Header) *const Table {
        var addr = @ptrToInt(self);
        addr += self.offset;
        return @intToPtr(*const Table, addr);
    }
};

/// Slice is used to represent a list of some table type
/// with zero-allocation access to the tables.
fn Slice(comptime T: type) type {
    comptime assert(@hasField(T, "table"));
    return struct {
        const Self = @This();
        const Type = T;
        items: []const Offset(*const Table),

        pub fn at(self: *const Self, index: usize) T {
            return T{ .table = self.items[index].get() };
        }

        pub fn len(self: *const Self) usize {
            return self.items.len;
        }

        /// toOwnedSlice allocates and returns a slice of the declared table type
        pub fn toOwnedSlice(self: *const Self, allocator: *std.mem.Allocator) ![]T {
            var slice = try allocator.allocAdvanced(T, null, self.items.len, .exact);
            for (self.items) |*item, i| {
                slice[i] = T{ .table = item.get() };
            }
            return slice;
        }
    };
}

/// Table represents an encoded table.
/// It holds a signed offset to the virtual table.
/// After the offset is the actual data values of the fields
/// described by the virtual table and any required padding.
const Table = packed struct {
    offset: i32, // signed offset to vtable
    // table field data + padding

    fn getVTable(self: *const Table) *const VTable {
        var addr = @ptrToInt(self);
        if (self.offset > 0) {
            addr -= @intCast(u32, self.offset);
        } else {
            addr += @intCast(u32, (std.math.absInt(self.offset) catch unreachable));
        }
        return @intToPtr(*const VTable, addr);
    }

    fn getField(self: *const Table, offset: u16, comptime FieldType: type) FieldType {
        if (comptime @typeInfo(FieldType) == .Struct and @hasField(FieldType, "table")) {
            // we have a table
            return FieldType{ .table = self.getField(offset, *const Offset(*const Table)).get() };
        }
        if (comptime @typeInfo(FieldType) == .Struct and @hasField(FieldType, "items")) {
            // we have a list of tables
            return FieldType{ .items = self.getField(offset, []const Offset(*const Table)) };
        }
        var addr = @ptrToInt(self);
        addr += offset;
        return switch (@typeInfo(FieldType)) {
            .Pointer => |ptr| switch (ptr.size) {
                .One, .C => @intToPtr(FieldType, addr),
                .Slice, .Many => blk: {
                    const vecOffset = @intToPtr(*const u32, addr).*;
                    addr += vecOffset;
                    const count = @intToPtr(*const u32, addr).*;
                    addr += @sizeOf(u32);
                    const slice = @intToPtr([*]const ptr.child, addr);
                    break :blk slice[0..count];
                },
            },
            .Enum => |e| @intToEnum(FieldType, @intToPtr(*const e.tag_type, addr).*),
            else => @intToPtr(*const FieldType, addr).*,
        };
    }

    fn getNthField(self: *const Table, index: usize, comptime FieldType: type) ?FieldType {
        const offset = self.getVTable().getFieldOffset(index);
        if (offset == 0) {
            return null;
        }
        return switch (@typeInfo(FieldType)) {
            .Union => |u| blk: {
                // first we get the tag field which is always 1 less than the union value
                // then we get the value using the tag to find the type
                if (self.getNthField(index - 1, u.tag_type.?)) |tag| {
                    const tagInt = @intCast(usize, @enumToInt(tag));
                    // using an inline for to unroll all possible values of tag
                    // this is required because unionInit requires all parameters to be known at compile time
                    // tagInt above is only known at runtime, so this unrolled series of if statements
                    // takes the runtime value and makes it compile time.
                    inline for (u.fields) |field, i| {
                        if (i == tagInt) {
                            if (comptime @typeInfo(field.field_type) == .Void) {
                                // special case void for "none"
                                break :blk @unionInit(FieldType, "none", .{});
                            } else {
                                break :blk @unionInit(FieldType, field.name, self.getField(offset, field.field_type));
                            }
                        }
                    }
                }
                break :blk @unionInit(FieldType, "none", .{});
            },
            else => self.getField(offset, FieldType),
        };
    }
};

/// Offset is used for the extra layer of indirection
/// used to encode sub-table fields.
/// for example, a field that is a vector of tables
/// is encoded as a vector of Offset(*const Table)
/// aka a vector of offsets to the tables
fn Offset(comptime T: type) type {
    return packed struct {
        const Self = @This();
        offset: u32,

        fn get(self: *const Self) T {
            var addr = @ptrToInt(self);
            addr += self.offset;
            return @intToPtr(T, addr);
        }
    };
}

/// VTable represents the encoded virtual table which holds
/// an unsigned 16-bit length of the virtual table itself
/// as well as an unsigned 16-bit length of the table data.
/// Following those 4 bytes are unsigned 16-bit offsets for each
/// of the fields of the table. The offsets point to the
/// data values of the fields in the table.
/// VTables might be shared between tables of different types.
const VTable = packed struct {
    vlen: u16, // vtable length
    tlen: u16, // table length
    // field offsets

    fn numFields(self: *const VTable) usize {
        return (self.vlen - @sizeOf(VTable)) / @sizeOf(u16);
    }

    fn getFieldOffset(self: *const VTable, index: usize) u16 {
        if (index >= self.numFields()) return 0;
        //assert(index < self.numFields());
        var addr = @ptrToInt(self);
        addr += @sizeOf(VTable);
        addr += (index * @sizeOf(u16));
        return @intToPtr(*u16, addr).*;
    }

    fn debugDump(self: *const VTable) void {
        std.debug.print("vtable: {}\n", .{self});
        var i: usize = 0;
        while (i < self.numFields()) : (i += 1) {
            std.debug.print("  i: {} offset: {}\n", .{ i, self.getFieldOffset(i) });
        }
    }
};

/// DefineTable takes a union(enum) where each field of the union represents a field of the table
/// and an anonymous struct with default values assigned for the table fields.
/// It returns a struct which holds a Table pointer.
/// The get function takes a union field to access an individual field of the table.
/// Example:
/// ```const Weapon = DefineTable(union(enum) {
///     name: []const u8,
///     damage: u16,
/// }, .{ .damage = 1 });
/// ...
/// weapon.get(.damage).?
/// ```
fn DefineTable(comptime definition: type, defaults: anytype) type {
    return struct {
        const Self = @This();
        const FieldType = std.meta.Tag(definition);
        const Union = @typeInfo(definition).Union;
        const Defaults = defaults;

        table: *const Table,

        /// get takes a union(enum) tag name from FieldType
        /// and return the value of that field in the encoded flatbuffer
        /// fields with default values return non-optional types
        /// fields without default values return optional types
        fn get(self: *const Self, f: anytype) if (@hasField(@TypeOf(Self.Defaults), Self.Union.fields[@enumToInt(@as(FieldType, f))].name)) Self.Union.fields[@enumToInt(@as(FieldType, f))].field_type else ?Self.Union.fields[@enumToInt(@as(FieldType, f))].field_type {
            comptime const fieldIdx = @enumToInt(@as(FieldType, f));
            comptime const fieldType = Self.Union.fields[fieldIdx].field_type;
            const val = self.table.getNthField(fieldIdx, fieldType);
            if (val) |v| {
                return v;
            }
            comptime const fieldName = Self.Union.fields[fieldIdx].name;
            return comptime if (@hasField(@TypeOf(Self.Defaults), fieldName)) @field(Self.Defaults, fieldName) else null;
        }
    };
}

test "monster type" {
    // first define the types we need to represent the monster.fbs schema in Zig
    const Color = enum(u8) { Red, Green, Blue };
    const Weapon = DefineTable(union(enum) {
        name: []const u8,
        damage: u16,
    }, .{});
    const Equipment = union(enum) {
        const Self = @This();
        none: void,
        weapon: Weapon,
    };
    const Vec3 = packed struct { x: f32, y: f32, z: f32 };
    const Monster = DefineTable(union(enum) {
        pos: Vec3,
        mana: u16,
        hp: u16,
        name: []const u8,
        friendly: bool,
        inventory: []const u8,
        color: Color,
        weapons: Slice(Weapon),
        equipped_type: std.meta.Tag(Equipment),
        equipped: Equipment,
        path: []const Vec3,
    }, .{
        .mana = 150,
        .hp = 100,
        .friendly = false,
        .color = Color.Blue,
        .equipped = Equipment.none, // XXX: always provide none for union default
    });

    const embed = @embedFile("./monster1.bin");
    // work around for alignment of embedFile data https://github.com/ziglang/zig/issues/4680
    var data: [embed.len]u8 = undefined;
    @memcpy(data[0..data.len], embed, embed.len);
    const hdr = @ptrCast(*Header, &data[0]);
    try t.expectEqual(@intCast(u32, 32), hdr.offset);
    const monster = Monster{ .table = hdr.getRootTable() };
    try t.expectEqual(Vec3{ .x = 1.0, .y = 32.5, .z = -1.0 }, monster.get(.pos).?);
    try t.expectEqual(@as(u16, 150), monster.get(.mana));
    try t.expectEqual(@as(u16, 100), monster.get(.hp));
    try t.expectEqualStrings("hippo", monster.get(.name).?);
    try t.expectEqual(Color.Blue, monster.get(.color));
    const equipped = monster.get(.equipped);
    try t.expect(equipped == .weapon);
    try t.expectEqualStrings("harness", equipped.weapon.get(.name).?);
    // 0 encodes to null, field is missing from vtable
    try t.expect(equipped.weapon.get(.damage) == null);

    {
        const weapons = monster.get(.weapons).?;
        try t.expect(2 == weapons.len());
        try t.expectEqualStrings("teeth", weapons.at(0).get(.name).?);
    }

    {
        const alloc = t.allocator;
        const weapons = try monster.get(.weapons).?.toOwnedSlice(alloc);
        defer alloc.free(weapons);

        const expectedWeapons = [_]struct {
            name: []const u8,
            damage: u16,
        }{
            .{ .name = "teeth", .damage = 6666 },
            .{ .name = "feet", .damage = 9999 },
        };
        for (weapons) |*weap, i| {
            try t.expectEqualStrings(expectedWeapons[i].name, weap.get(.name).?);
            try t.expectEqual(@intCast(u16, expectedWeapons[i].damage), weap.get(.damage).?);
        }
    }
}

monster.fbs

// Example IDL file for our monster's schema.

namespace MyGame.Sample;

enum Color:byte { Red = 0, Green, Blue = 2 }

union Equipment { Weapon } // Optionally add more tables.

struct Vec3 {
  x:float;
  y:float;
  z:float;
}

table Monster {
  pos:Vec3;
  mana:short = 150;
  hp:short = 100;
  name:string;
  friendly:bool = false (deprecated);
  inventory:[ubyte];
  color:Color = Blue;
  weapons:[Weapon];
  equipped:Equipment;
  path:[Vec3];
}

table Weapon {
  name:string;
  damage:short;
}

root_type Monster;

monster1.json

{
	"pos": {"x": 1.0, "y": 32.5, "z": -1.0},
	"name": "hippo",
	"color": "Blue",
	"weapons": [
		{"name": "teeth", "damage": 6666},
		{"name": "feet", "damage": 9999}
	],
	"equipped_type": "Weapon",
	"equipped": {"name": "harness", "damage": 0},
	"path": [
		{"x": 0.5, "y": 0.5, "z": 1.0},
		{"x": 40, "y": 66, "z": 2.0}
	]
}