diff --git a/src/array.zig b/src/array.zig index b2a20fa..ab18919 100644 --- a/src/array.zig +++ b/src/array.zig @@ -7,7 +7,7 @@ const isOptional = @import("utils.zig").isOptional; const maybePackNull = @import("null.zig").maybePackNull; const maybeUnpackNull = @import("null.zig").maybeUnpackNull; -const packIntValue = @import("int.zig").packIntValue; +const packHeaderAndInt = @import("utils.zig").packHeaderAndInt; const unpackIntValue = @import("int.zig").unpackIntValue; const packAny = @import("any.zig").packAny; @@ -33,11 +33,9 @@ pub fn packArrayHeader(writer: *std.Io.Writer, len: usize) !void { if (len <= hdrs.FIXARRAY_MAX - hdrs.FIXARRAY_MIN) { try writer.writeByte(hdrs.FIXARRAY_MIN + @as(u8, @intCast(len))); } else if (len <= std.math.maxInt(u16)) { - try writer.writeByte(hdrs.ARRAY16); - try packIntValue(writer, u16, @intCast(len)); + try packHeaderAndInt(writer, hdrs.ARRAY16, u16, @intCast(len)); } else if (len <= std.math.maxInt(u32)) { - try writer.writeByte(hdrs.ARRAY32); - try packIntValue(writer, u32, @intCast(len)); + try packHeaderAndInt(writer, hdrs.ARRAY32, u32, @intCast(len)); } else { return error.ArrayTooLong; } diff --git a/src/binary.zig b/src/binary.zig index d5aa17c..013256b 100644 --- a/src/binary.zig +++ b/src/binary.zig @@ -7,7 +7,7 @@ const NonOptional = @import("utils.zig").NonOptional; const maybePackNull = @import("null.zig").maybePackNull; const maybeUnpackNull = @import("null.zig").maybeUnpackNull; -const packIntValue = @import("int.zig").packIntValue; +const packHeaderAndInt = @import("utils.zig").packHeaderAndInt; const unpackIntValue = @import("int.zig").unpackIntValue; const unpackShortIntValue = @import("int.zig").unpackShortIntValue; @@ -29,14 +29,11 @@ pub fn sizeOfPackedBinary(len: usize) !usize { pub fn packBinaryHeader(writer: *std.Io.Writer, len: usize) !void { if (len <= std.math.maxInt(u8)) { - try writer.writeByte(hdrs.BIN8); - try packIntValue(writer, u8, @intCast(len)); + try packHeaderAndInt(writer, hdrs.BIN8, u8, @intCast(len)); } else if (len <= std.math.maxInt(u16)) { - try writer.writeByte(hdrs.BIN16); - try packIntValue(writer, u16, @intCast(len)); + try packHeaderAndInt(writer, hdrs.BIN16, u16, @intCast(len)); } else if (len <= std.math.maxInt(u32)) { - try writer.writeByte(hdrs.BIN32); - try packIntValue(writer, u32, @intCast(len)); + try packHeaderAndInt(writer, hdrs.BIN32, u32, @intCast(len)); } else { return error.BinaryTooLong; } diff --git a/src/float.zig b/src/float.zig index 6179ee4..4858bbc 100644 --- a/src/float.zig +++ b/src/float.zig @@ -3,6 +3,7 @@ const hdrs = @import("headers.zig"); const NonOptional = @import("utils.zig").NonOptional; const maybePackNull = @import("null.zig").maybePackNull; +const packHeaderAndInt = @import("utils.zig").packHeaderAndInt; const maybeUnpackNull = @import("null.zig").maybeUnpackNull; inline fn assertFloatType(comptime T: type) type { @@ -30,26 +31,21 @@ pub fn packFloat(writer: *std.Io.Writer, comptime T: type, value_or_maybe_null: const Type = assertFloatType(T); const value: Type = try maybePackNull(writer, T, value_or_maybe_null) orelse return; - comptime var TargetType: type = undefined; const type_info = @typeInfo(Type); - switch (type_info.float.bits) { - 0...32 => { - try writer.writeByte(hdrs.FLOAT32); - TargetType = f32; - }, - 33...64 => { - try writer.writeByte(hdrs.FLOAT64); - TargetType = f64; - }, + const TargetType = switch (type_info.float.bits) { + 0...32 => f32, + 33...64 => f64, else => @compileError("Unsupported float size"), - } + }; + const header = switch (type_info.float.bits) { + 0...32 => hdrs.FLOAT32, + 33...64 => hdrs.FLOAT64, + else => unreachable, + }; const IntType = std.meta.Int(.unsigned, @bitSizeOf(TargetType)); const int_value = @as(IntType, @bitCast(@as(TargetType, @floatCast(value)))); - - var buf: [@sizeOf(IntType)]u8 = undefined; - std.mem.writeInt(IntType, buf[0..], int_value, .big); - try writer.writeAll(buf[0..]); + return packHeaderAndInt(writer, header, IntType, int_value); } pub fn readFloatValue(reader: *std.Io.Reader, comptime SourceFloat: type, comptime TargetFloat: type) !TargetFloat { diff --git a/src/int.zig b/src/int.zig index 65a6eda..1b61ec8 100644 --- a/src/int.zig +++ b/src/int.zig @@ -3,6 +3,7 @@ const hdrs = @import("headers.zig"); const NonOptional = @import("utils.zig").NonOptional; const maybePackNull = @import("null.zig").maybePackNull; +const packHeaderAndInt = @import("utils.zig").packHeaderAndInt; const maybeUnpackNull = @import("null.zig").maybeUnpackNull; inline fn assertIntType(comptime T: type) type { @@ -66,12 +67,6 @@ pub fn getIntSize(comptime T: type, value: T) usize { } } -pub fn packIntValue(writer: *std.Io.Writer, comptime T: type, value: T) !void { - var buf: [@sizeOf(T)]u8 = undefined; - std.mem.writeInt(T, buf[0..], value, .big); - try writer.writeAll(buf[0..]); -} - pub fn packInt(writer: *std.Io.Writer, comptime T: type, value_or_maybe_null: T) !void { const Type = assertIntType(T); const value: Type = try maybePackNull(writer, T, value_or_maybe_null) orelse return; @@ -122,8 +117,7 @@ pub fn packInt(writer: *std.Io.Writer, comptime T: type, value_or_maybe_null: T) } pub fn packFixedSizeInt(writer: *std.Io.Writer, comptime T: type, value: T) !void { - try writer.writeByte(resolveFixedSizeIntHeader(T)); - try packIntValue(writer, T, value); + return packHeaderAndInt(writer, resolveFixedSizeIntHeader(T), T, value); } inline fn resolveFixedSizeIntHeader(comptime T: type) u8 { diff --git a/src/map.zig b/src/map.zig index 8e876bd..4d329df 100644 --- a/src/map.zig +++ b/src/map.zig @@ -8,7 +8,7 @@ const isOptional = @import("utils.zig").isOptional; const maybePackNull = @import("null.zig").maybePackNull; const maybeUnpackNull = @import("null.zig").maybeUnpackNull; -const packIntValue = @import("int.zig").packIntValue; +const packHeaderAndInt = @import("utils.zig").packHeaderAndInt; const unpackIntValue = @import("int.zig").unpackIntValue; const packAny = @import("any.zig").packAny; @@ -34,11 +34,9 @@ pub fn packMapHeader(writer: *std.Io.Writer, len: usize) !void { if (len <= hdrs.FIXMAP_MAX - hdrs.FIXMAP_MIN) { try writer.writeByte(hdrs.FIXMAP_MIN + @as(u8, @intCast(len))); } else if (len <= std.math.maxInt(u16)) { - try writer.writeByte(hdrs.MAP16); - try packIntValue(writer, u16, @intCast(len)); + try packHeaderAndInt(writer, hdrs.MAP16, u16, @intCast(len)); } else if (len <= std.math.maxInt(u32)) { - try writer.writeByte(hdrs.MAP32); - try packIntValue(writer, u32, @intCast(len)); + try packHeaderAndInt(writer, hdrs.MAP32, u32, @intCast(len)); } else { return error.MapTooLong; } diff --git a/src/msgpack.zig b/src/msgpack.zig index 804f931..50aebdf 100644 --- a/src/msgpack.zig +++ b/src/msgpack.zig @@ -15,7 +15,6 @@ pub const unpackBool = @import("bool.zig").unpackBool; pub const getIntSize = @import("int.zig").getIntSize; pub const getMaxIntSize = @import("int.zig").getMaxIntSize; pub const packInt = @import("int.zig").packInt; -pub const packIntValue = @import("int.zig").packIntValue; pub const unpackInt = @import("int.zig").unpackInt; pub const getFloatSize = @import("float.zig").getFloatSize; @@ -40,6 +39,7 @@ pub const sizeOfPackedString = @import("string.zig").sizeOfPackedString; pub const sizeOfPackedStringHeader = @import("string.zig").sizeOfPackedStringHeader; pub const packStringHeader = @import("string.zig").packStringHeader; pub const packString = @import("string.zig").packString; +pub const packStringLiteral = @import("string.zig").packStringLiteral; pub const unpackStringHeader = @import("string.zig").unpackStringHeader; pub const unpackString = @import("string.zig").unpackString; pub const unpackStringInto = @import("string.zig").unpackStringInto; diff --git a/src/string.zig b/src/string.zig index d8f0201..7331b29 100644 --- a/src/string.zig +++ b/src/string.zig @@ -7,7 +7,8 @@ const NonOptional = @import("utils.zig").NonOptional; const maybePackNull = @import("null.zig").maybePackNull; const maybeUnpackNull = @import("null.zig").maybeUnpackNull; -const packIntValue = @import("int.zig").packIntValue; +const packHeaderAndInt = @import("utils.zig").packHeaderAndInt; +const reserveArray = @import("utils.zig").reserveArray; const unpackIntValue = @import("int.zig").unpackIntValue; const unpackShortIntValue = @import("int.zig").unpackShortIntValue; @@ -33,14 +34,11 @@ pub fn packStringHeader(writer: *std.Io.Writer, len: usize) !void { if (len <= hdrs.FIXSTR_MAX - hdrs.FIXSTR_MIN) { try writer.writeByte(hdrs.FIXSTR_MIN + @as(u8, @intCast(len))); } else if (len <= std.math.maxInt(u8)) { - try writer.writeByte(hdrs.STR8); - try packIntValue(writer, u8, @intCast(len)); + try packHeaderAndInt(writer, hdrs.STR8, u8, @intCast(len)); } else if (len <= std.math.maxInt(u16)) { - try writer.writeByte(hdrs.STR16); - try packIntValue(writer, u16, @intCast(len)); + try packHeaderAndInt(writer, hdrs.STR16, u16, @intCast(len)); } else if (len <= std.math.maxInt(u32)) { - try writer.writeByte(hdrs.STR32); - try packIntValue(writer, u32, @intCast(len)); + try packHeaderAndInt(writer, hdrs.STR32, u32, @intCast(len)); } else { return error.StringTooLong; } @@ -64,6 +62,35 @@ pub fn packString(writer: *std.Io.Writer, value_or_maybe_null: ?[]const u8) !voi try writer.writeAll(value); } +/// Packs a string whose contents are known at comptime, such as a struct field +/// name used as a map key. The header and the bytes are one contiguous store of +/// comptime-known length, so this compiles to immediate stores with no `memcpy`. +pub fn packStringLiteral(writer: *std.Io.Writer, comptime value: []const u8) !void { + const bytes = comptime blk: { + if (value.len <= hdrs.FIXSTR_MAX - hdrs.FIXSTR_MIN) { + break :blk [_]u8{hdrs.FIXSTR_MIN + @as(u8, value.len)} ++ value[0..].*; + } else if (value.len <= std.math.maxInt(u8)) { + break :blk [_]u8{ hdrs.STR8, @intCast(value.len) } ++ value[0..].*; + } else if (value.len <= std.math.maxInt(u16)) { + var len_bytes: [2]u8 = undefined; + std.mem.writeInt(u16, &len_bytes, @intCast(value.len), .big); + break :blk [_]u8{hdrs.STR16} ++ len_bytes ++ value[0..].*; + } else if (value.len <= std.math.maxInt(u32)) { + var len_bytes: [4]u8 = undefined; + std.mem.writeInt(u32, &len_bytes, @intCast(value.len), .big); + break :blk [_]u8{hdrs.STR32} ++ len_bytes ++ value[0..].*; + } else { + @compileError("String literal too long: " ++ value); + } + }; + + if (reserveArray(writer, bytes.len)) |dst| { + dst.* = bytes; + return; + } + try writer.writeAll(&bytes); +} + pub fn unpackString(reader: *std.Io.Reader, allocator: std.mem.Allocator) ![]u8 { const len = try unpackStringHeader(reader, u32); diff --git a/src/struct.zig b/src/struct.zig index 2e9e1c6..4793385 100644 --- a/src/struct.zig +++ b/src/struct.zig @@ -16,7 +16,7 @@ const unpackMapHeader = @import("map.zig").unpackMapHeader; const packInt = @import("int.zig").packInt; const unpackInt = @import("int.zig").unpackInt; -const packString = @import("string.zig").packString; +const packStringLiteral = @import("string.zig").packStringLiteral; const unpackStringInto = @import("string.zig").unpackStringInto; const packArrayHeader = @import("array.zig").packArrayHeader; @@ -103,10 +103,10 @@ pub fn packStructAsMap(writer: *std.Io.Writer, comptime T: type, value: T, compt try packInt(writer, u16, i); }, .field_name => { - try packString(writer, field.name); + try packStringLiteral(writer, field.name); }, .field_name_prefix => |prefix| { - try packString(writer, strPrefix(field.name, prefix)); + try packStringLiteral(writer, comptime strPrefix(field.name, prefix)); }, .custom => { const key = comptime T.msgpackFieldKey(@field(FieldEnum, field.name)); diff --git a/src/union.zig b/src/union.zig index d473614..1f8cc77 100644 --- a/src/union.zig +++ b/src/union.zig @@ -15,7 +15,7 @@ const unpackMapHeader = @import("map.zig").unpackMapHeader; const packInt = @import("int.zig").packInt; const unpackInt = @import("int.zig").unpackInt; -const packString = @import("string.zig").packString; +const packStringLiteral = @import("string.zig").packStringLiteral; const unpackStringInto = @import("string.zig").unpackStringInto; const packArrayHeader = @import("array.zig").packArrayHeader; @@ -62,7 +62,7 @@ fn strPrefix(src: []const u8, len: usize) []const u8 { return src[0..@min(src.len, len)]; } -pub fn packUnionAsMap(writer: *std.Io.Writer, comptime T: type, value: T, opts: UnionAsMapOptions) !void { +pub fn packUnionAsMap(writer: *std.Io.Writer, comptime T: type, value: T, comptime opts: UnionAsMapOptions) !void { const type_info = @typeInfo(T); const fields = type_info.@"union".fields; @@ -77,10 +77,10 @@ pub fn packUnionAsMap(writer: *std.Io.Writer, comptime T: type, value: T, opts: try packInt(writer, u16, i); }, .field_name => { - try packString(writer, field.name); + try packStringLiteral(writer, field.name); }, .field_name_prefix => |prefix| { - try packString(writer, strPrefix(field.name, prefix)); + try packStringLiteral(writer, comptime strPrefix(field.name, prefix)); }, } try packAny(writer, @field(value, field.name)); @@ -88,7 +88,7 @@ pub fn packUnionAsMap(writer: *std.Io.Writer, comptime T: type, value: T, opts: } } -pub fn packUnionAsTagged(writer: *std.Io.Writer, comptime T: type, value: T, opts: UnionAsTaggedOptions) !void { +pub fn packUnionAsTagged(writer: *std.Io.Writer, comptime T: type, value: T, comptime opts: UnionAsTaggedOptions) !void { const type_info = @typeInfo(T); const fields = type_info.@"union".fields; @@ -103,22 +103,22 @@ pub fn packUnionAsTagged(writer: *std.Io.Writer, comptime T: type, value: T, opt try packMapHeader(writer, field_count + 1); - try packString(writer, opts.tag_field); + try packStringLiteral(writer, opts.tag_field); switch (opts.tag_value) { .field_index => { try packInt(writer, u16, i); }, .field_name => { - try packString(writer, field.name); + try packStringLiteral(writer, field.name); }, .field_name_prefix => |prefix| { - try packString(writer, strPrefix(field.name, prefix)); + try packStringLiteral(writer, comptime strPrefix(field.name, prefix)); }, } if (field_type_info == .@"struct") { inline for (field_type_info.@"struct".fields) |struct_field| { - try packString(writer, struct_field.name); + try packStringLiteral(writer, struct_field.name); try packAny(writer, @field(field_value, struct_field.name)); } } else if (field.type != void) { @@ -139,7 +139,7 @@ pub fn packUnion(writer: *std.Io.Writer, comptime T: type, value_or_maybe_null: @compileError("Expected union type"); } - const format = if (std.meta.hasFn(Type, "msgpackFormat")) Type.msgpackFormat() else default_union_format; + const format = comptime if (std.meta.hasFn(Type, "msgpackFormat")) Type.msgpackFormat() else default_union_format; switch (format) { .as_map => |opts| { return packUnionAsMap(writer, Type, value, opts); diff --git a/src/utils.zig b/src/utils.zig index 0deb086..bcd2cdb 100644 --- a/src/utils.zig +++ b/src/utils.zig @@ -22,6 +22,35 @@ test isOptional { try std.testing.expect(!isOptional(u32)); } +/// Reserves `len` bytes of the writer's buffer, or returns null if it does not +/// fit. Keeps `len` comptime-known so stores through the result lower to +/// immediate stores instead of a `memcpy` call. Callers must fill the result. +pub inline fn reserveArray(writer: *std.Io.Writer, comptime len: usize) ?*[len]u8 { + if (writer.end + len <= writer.buffer.len) { + @branchHint(.likely); + const dst = writer.buffer[writer.end..][0..len]; + writer.end += len; + return dst; + } + return null; +} + +/// Writes a header byte followed by a big-endian `T`, as a single contiguous +/// store when the writer's buffer has room, and as a single `writeAll` +/// otherwise. Never splits the header from its payload across a drain. +pub fn packHeaderAndInt(writer: *std.Io.Writer, header: u8, comptime T: type, value: T) !void { + const size = @sizeOf(T); + if (reserveArray(writer, 1 + size)) |dst| { + dst[0] = header; + std.mem.writeInt(T, dst[1..][0..size], value, .big); + return; + } + var buf: [1 + size]u8 = undefined; + buf[0] = header; + std.mem.writeInt(T, buf[1..][0..size], value, .big); + try writer.writeAll(&buf); +} + var no_allocator_dummy: u8 = 0; pub const NoAllocator = struct {