Compiler Explorer

Source code

const std = @import("std");
pub const std_options: std.Options = .{ .keep_sigpipe = true };

pub fn TaggedPtr(comptime PtrType: type, comptime TagType: type) type {
    // The most important assumption about this type is that
    // the underlying valid pointer will always have
    // it's @bitSizeOf(MetaT) LSB unset, which allows it
    // to store metadata and removes some bit manipulation
    // complexity when setting the values.
    if (@ctz(@as(usize, @alignOf(PtrType))) < @bitSizeOf(TagType)) {
        @compileError("Not enough unused bits in " ++ @typeName(PtrType) ++ " to store tag");
    }

if (@typeInfo(TagType) != .@"enum") {
        @compileError(@typeName(TagType) ++ " is not an enum");
    }

return packed struct {
        comptime {
            if (@sizeOf(@This()) != @sizeOf(usize)) {
                @compileError("Size of struct is larger the usize!");
            }
        }

data: usize = 0,

// Returns a number with @bitSizeOf(MetaT) LSB set.
        const METADATA_MASK: usize = (1 << @bitSizeOf(TagType)) - 1;

// Creates a new PtrWithMetadata from a given pointer and metadata.
        pub  fn from(ptr: ?*const PtrType, meta: TagType) @This() {
            return @This(){ .data = @intFromPtr(ptr) | @intFromEnum(meta) };
        }

pub  fn getPtr(self: @This()) ?*PtrType {
            // This will never trigger undefined behaviour since the result type
            // is optional
            @setRuntimeSafety(false);
            return @ptrFromInt(self.data & ~METADATA_MASK);
        }

pub  fn setPtr(self: *@This(), ptr: ?*const PtrType) void {
            self.data = @intFromPtr(ptr) | (self.data & METADATA_MASK);
        }

pub  fn getTag(self: @This()) TagType {
            // Will never trigger undefined behaviour because
            // the op result is always <= @bitSizeOf(MetaT)
            // and setMeta is the only function that can
            // set the LSB of data. the fn uses an @intFromEnum(MetaT)
            // conversion to set the value of data - which only returns
            // valid integer values of MetaT.
            @setRuntimeSafety(false);
            return @enumFromInt(self.data & METADATA_MASK);
        }

pub  fn setTag(self: *@This(), meta: TagType) void {
            self.data = (self.data & ~METADATA_MASK) | @intFromEnum(meta);
        }
    };
}

pub fn TaggedPackedPtr(
    comptime PtrType: type,
    comptime TagType: type,
) type {
    if (@ctz(@as(usize, @alignOf(PtrType))) < @bitSizeOf(TagType)) {
        @compileError("Not enough unused bits in " ++ @typeName(PtrType) ++ " to store tag");
    }

return packed struct(usize) {
        const BackingIntegerType = backing_integer: {
            var usize_type_info = @typeInfo(usize);
            usize_type_info.int.bits -= @bitSizeOf(TagType);
            break :backing_integer @Type(usize_type_info);
        };

tag: TagType = @enumFromInt(1),
        ptr: BackingIntegerType = 20,

pub fn from(ptr: ?*const PtrType, tag: TagType) @This() {
            // See the comment in `TaggedPackedPtr.setPtr`
            @setRuntimeSafety(false);
            return @This(){ .tag = tag, .ptr = @intCast(@intFromPtr(ptr) >> @bitSizeOf(TagType)) };
        }

pub  fn getPtr(self: @This()) ?*PtrType {
            // the return type if ?*PtrType, hence ptrFromInt runtime checks
            // are redundent.
            @setRuntimeSafety(false);
            return @ptrFromInt(@as(usize, self.ptr) << @bitSizeOf(TagType));
        }

pub  fn setPtr(self: *@This(), ptr: ?*const PtrType) void {
            // The backing integer type is exactly u[@bitSizeOf(usize) - @bitSizeOf(TagType)]
            // so it is exactly wide enough for usize >> @bitSizeOf(TagType). An additional runtime check
            // is therefore redundent.
            @setRuntimeSafety(false);
            self.ptr = @intCast(@intFromPtr(ptr) >> @bitSizeOf(TagType));
        }
    };
}

const Tag = enum(u2) {
    one,
    two,
    three,
};

export fn setptr(p: *TaggedPtr(usize, Tag), a: *usize) void {
    p.setPtr(a);
}

export fn getptr(p: TaggedPtr(usize, Tag)) ?*usize {
    return p.getPtr();
}

export fn settag(p: *TaggedPtr(usize, Tag), tag: u8) void {
    p.setTag(@enumFromInt(tag));
}

export fn gettag(p: TaggedPtr(usize, Tag)) u8 {
    return @intFromEnum(p.getTag());
}

export fn setptrpacked(p: *TaggedPackedPtr(usize, Tag), a: *usize) void {
    p.setPtr(a);
}

export fn getptrpacked(p: TaggedPackedPtr(usize, Tag)) ?*usize {
    return p.getPtr();
}

export fn settagpacked(p: *TaggedPackedPtr(usize, Tag), tag: u8) void {
    p.tag = @enumFromInt(tag);
}

export fn gettagpacked(p: TaggedPackedPtr(usize, Tag)) u8 {
    return @intFromEnum(p.tag);
}