lib/choir/src/serialization/binary/writer.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const binary = @import("root.zig");
3
4 const Storage = enum { allocated, borrowed };
5
6 pub fn Writer(comptime limits: binary.Limits) type {
7 return BufferWriter(limits, .allocated);
8 }
9
10 /// Has no allocator capability. Its finished bytes borrow the caller's buffer.
11 pub fn FixedWriter(comptime limits: binary.Limits) type {
12 return BufferWriter(limits, .borrowed);
13 }
14
15 fn BufferWriter(comptime limits: binary.Limits, comptime ownership: Storage) type {
16 return struct {
17 allocator: if (ownership == .allocated) std.mem.Allocator else void,
18 bytes: std.ArrayListUnmanaged(u8) = .empty,
19 total_entries: usize = 0,
20
21 const Self = @This();
22
23 pub fn init(input: if (ownership == .allocated) std.mem.Allocator else []u8) Self {
24 return if (ownership == .allocated)
25 .{ .allocator = input }
26 else
27 .{ .allocator = {}, .bytes = .initBuffer(input) };
28 }
29
30 pub fn deinit(self: *Self) void {
31 if (ownership == .allocated) self.bytes.deinit(self.allocator);
32 self.* = undefined;
33 }
34
35 pub fn finish(self: *Self) std.mem.Allocator.Error![]u8 {
36 if (ownership == .allocated) return self.bytes.toOwnedSlice(self.allocator);
37 const bytes = self.bytes.items;
38 self.bytes = .empty;
39 return bytes;
40 }
41
42 pub fn writeRaw(self: *Self, value: []const u8) binary.WriteError!void {
43 const new_length = std.math.add(usize, self.bytes.items.len, value.len) catch
44 return error.LimitExceeded;
45 if (new_length > limits.serialized_bytes) return error.LimitExceeded;
46 if (ownership == .allocated) {
47 try self.bytes.appendSlice(self.allocator, value);
48 } else {
49 if (new_length > self.bytes.capacity) return error.LimitExceeded;
50 self.bytes.appendSliceAssumeCapacity(value);
51 }
52 }
53
54 pub const writeInt = encodeInt;
55
56 pub const writeBool = encodeBool;
57
58 pub const writeTag = encodeTag;
59
60 pub fn writeCount(self: *Self, count: usize) binary.WriteError!void {
61 if (count > limits.collection_entries or count > std.math.maxInt(u32)) {
62 return error.LimitExceeded;
63 }
64 self.total_entries = std.math.add(usize, self.total_entries, count) catch
65 return error.LimitExceeded;
66 if (self.total_entries > limits.total_entries) return error.LimitExceeded;
67 try self.writeInt(u32, @intCast(count));
68 }
69
70 pub fn writeString(self: *Self, value: []const u8) binary.WriteError!void {
71 try writeLengthPrefixed(self, value, limits.string_bytes);
72 }
73
74 pub fn writeBlob(self: *Self, value: []const u8) binary.WriteError!void {
75 try writeLengthPrefixed(self, value, limits.blob_bytes);
76 }
77
78 pub const writeOptionalString = encodeOptionalString;
79
80 pub const writeOptionalU16 = encodeOptionalU16;
81 };
82 }
83
84 test "fixed binary writer borrows output and refuses capacity overflow without mutation" {
85 var storage: [8]u8 = @splat(0xa5);
86 var writer = FixedWriter(.{
87 .serialized_bytes = 16,
88 .string_bytes = 8,
89 .blob_bytes = 8,
90 .collection_entries = 8,
91 .total_entries = 8,
92 }).init(&storage);
93 defer writer.deinit();
94 try writer.writeInt(u32, 0x12345678);
95 try std.testing.expectError(error.LimitExceeded, writer.writeRaw("12345"));
96 try std.testing.expectEqual(4, writer.bytes.items.len);
97 try std.testing.expectEqualSlices(u8, &.{ 0xa5, 0xa5, 0xa5, 0xa5 }, storage[4..]);
98 try writer.writeInt(u32, 0x9abcdef0);
99 const output = try writer.finish();
100 try std.testing.expectEqual(@intFromPtr(&storage), @intFromPtr(output.ptr));
101 try std.testing.expectEqualSlices(u8, &.{
102 0x78, 0x56, 0x34, 0x12, 0xf0, 0xde, 0xbc, 0x9a,
103 }, output);
104 }
105
106 fn encodeInt(self: anytype, comptime T: type, value: T) binary.WriteError!void {
107 var storage: [@sizeOf(T)]u8 = undefined;
108 std.mem.writeInt(T, &storage, value, .little);
109 try self.writeRaw(&storage);
110 }
111
112 fn encodeBool(self: anytype, value: bool) binary.WriteError!void {
113 try self.writeInt(u8, if (value) 1 else 0);
114 }
115
116 fn encodeTag(self: anytype, value: anytype) binary.WriteError!void {
117 try self.writeInt(u8, try binary.encodeTag(value));
118 }
119
120 fn encodeOptionalString(self: anytype, value: ?[]const u8) binary.WriteError!void {
121 try self.writeBool(value != null);
122 if (value) |present| try self.writeString(present);
123 }
124
125 fn encodeOptionalU16(self: anytype, value: ?u16) binary.WriteError!void {
126 try self.writeBool(value != null);
127 if (value) |present| try self.writeInt(u16, present);
128 }
129
130 fn writeLengthPrefixed(
131 self: anytype,
132 value: []const u8,
133 maximum: usize,
134 ) binary.WriteError!void {
135 if (value.len > maximum or value.len > std.math.maxInt(u32)) return error.LimitExceeded;
136 try self.writeInt(u32, @intCast(value.len));
137 try self.writeRaw(value);
138 }