Something went wrong. Try again.
Something went wrong. Try again.
4.1 kB · 124 lines
Zig
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125const std = @import("std");const assert = @import("../quirks.zig").inlineAssert;const Allocator = std.mem.Allocator;
/// Creates a union that can be used to accommodate data that fit within an array,/// are a stable pointer, or require deallocation. This is helpful for thread/// messaging utilities.pub fn MessageData(comptime Elem: type, comptime small_size: comptime_int) type { return union(enum) { pub const Self = @This();
pub const Small = struct { pub const Max = small_size; pub const Array = [Max]Elem; pub const Len = std.math.IntFittingRange(0, small_size); data: Array = undefined, len: Len = 0, };
pub const Alloc = struct { alloc: Allocator, data: []Elem, };
pub const Stable = []const Elem;
/// A small write where the data fits into this union size. small: Small,
/// A stable pointer so we can just pass the slice directly through. /// This is useful i.e. for const data. stable: Stable,
/// Allocated and must be freed with the provided allocator. This /// should be rarely used. alloc: Alloc,
/// Initializes the union for a given data type. This will /// attempt to fit into a small value if possible, otherwise /// will allocate and put into alloc. /// /// This can't and will never detect stable pointers. pub fn init(alloc: Allocator, data: anytype) !Self { switch (@typeInfo(@TypeOf(data))) { .pointer => |info| { assert(info.size == .slice); assert(info.child == Elem);
// If it fits in our small request, do that. if (data.len <= Small.Max) { var buf: Small.Array = undefined; @memcpy(buf[0..data.len], data); return Self{ .small = .{ .data = buf, .len = @intCast(data.len), }, }; }
// Otherwise, allocate const buf = try alloc.dupe(Elem, data); errdefer alloc.free(buf); return Self{ .alloc = .{ .alloc = alloc, .data = buf, }, }; },
else => unreachable, } }
pub fn deinit(self: Self) void { switch (self) { .small, .stable => {}, .alloc => |v| v.alloc.free(v.data), } }
/// Returns a const slice of the data pointed to by this request. pub fn slice(self: *const Self) []const Elem { return switch (self.*) { .small => |*v| v.data[0..v.len], .stable => |v| v, .alloc => |v| v.data, }; } };}
test "MessageData init small" { const testing = std.testing; const alloc = testing.allocator;
const Data = MessageData(u8, 10); const input = "hello!"; const io = try Data.init(alloc, @as([]const u8, input)); try testing.expect(io == .small);}
test "MessageData init alloc" { const testing = std.testing; const alloc = testing.allocator;
const Data = MessageData(u8, 10); const input = "hello! " ** 100; const io = try Data.init(alloc, @as([]const u8, input)); try testing.expect(io == .alloc); io.alloc.alloc.free(io.alloc.data);}
test "MessageData small fits non-u8 sized data" { const testing = std.testing; const alloc = testing.allocator;
const len = 500; const Data = MessageData(u8, len); const input: []const u8 = "X" ** len; const io = try Data.init(alloc, input); try testing.expect(io == .small);}