diff --git a/snappy.zig b/snappy.zig index bcc2c48..f01b0c8 100644 --- a/snappy.zig +++ b/snappy.zig @@ -261,23 +261,32 @@ pub fn decode(allocator: Allocator, src: []const u8) ![]u8 { // TODO: Split up encode and decode into separate files once I better understand modules. fn emitLiteral(dst: []u8, lit: []const u8) usize { var i: usize = 0; - const n = @as(usize, @intCast(lit.len - 1)); - switch (n) { - 0...59 => { - dst[0] = @as(u8, @intCast(n)) << 2 | tagLiteral; - i = 1; - }, - 60...255 => { - dst[0] = 60 << 2 | tagLiteral; - dst[1] = @as(u8, @intCast(n)); - i = 2; - }, - else => { - dst[0] = 61 << 2 | tagLiteral; - dst[1] = @as(u8, @truncate(n)); - dst[2] = @as(u8, @truncate(n >> 8)); - i = 3; - }, + const n = @as(u32, @intCast(lit.len - 1)); + if (n < 60) { + dst[0] = @as(u8, @intCast(n)) << 2 | tagLiteral; + i = 1; + } else if (n < (1 << 8)) { + dst[0] = 60 << 2 | tagLiteral; + dst[1] = @as(u8, @intCast(n)); + i = 2; + } else if (n < (1 << 16)) { + dst[0] = 61 << 2 | tagLiteral; + dst[1] = @as(u8, @intCast(n & 0xff)); + dst[2] = @as(u8, @intCast(n >> 8)); + i = 3; + } else if (n < (1 << 24)) { + dst[0] = 62 << 2 | tagLiteral; + dst[1] = @as(u8, @intCast(n & 0xff)); + dst[2] = @as(u8, @intCast((n >> 8) & 0xff)); + dst[3] = @as(u8, @intCast(n >> 16)); + i = 4; + } else { + dst[0] = 63 << 2 | tagLiteral; + dst[1] = @as(u8, @intCast(n & 0xff)); + dst[2] = @as(u8, @intCast((n >> 8) & 0xff)); + dst[3] = @as(u8, @intCast((n >> 16) & 0xff)); + dst[4] = @as(u8, @intCast(n >> 24)); + i = 5; } std.mem.copyForwards(u8, dst[i..], lit); @@ -499,3 +508,42 @@ test "simple decode" { try testing.expectEqualSlices(u8, decoded, "oh snap, snappy is cool!\n"); } + +test "emit literal length > 255" { + var lit: [300]u8 = undefined; + for (&lit, 0..) |*b, i| { + b.* = @as(u8, @truncate(i * 31 + 7)); + } + + var dst: [320]u8 = undefined; + const written = emitLiteral(&dst, &lit); + + const n: u32 = lit.len - 1; + try testing.expectEqual(@as(usize, 3 + lit.len), written); + try testing.expectEqual(@as(u8, 61 << 2 | tagLiteral), dst[0]); + try testing.expectEqual(@as(u8, @intCast(n & 0xff)), dst[1]); + try testing.expectEqual(@as(u8, @intCast(n >> 8)), dst[2]); + try testing.expectEqualSlices(u8, lit[0..], dst[3 .. 3 + lit.len]); +} + +test "emit literal length > 65535" { + const allocator = testing.allocator; + const lit_len: usize = 70000; + var lit = try allocator.alloc(u8, lit_len); + defer allocator.free(lit); + for (lit, 0..) |*b, i| { + b.* = @as(u8, @truncate(i * 13 + 5)); + } + + const n: u32 = @as(u32, @intCast(lit.len - 1)); + var dst = try allocator.alloc(u8, 4 + lit.len); + defer allocator.free(dst); + const written = emitLiteral(dst, lit); + + try testing.expectEqual(@as(usize, 4 + lit.len), written); + try testing.expectEqual(@as(u8, 62 << 2 | tagLiteral), dst[0]); + try testing.expectEqual(@as(u8, @intCast(n & 0xff)), dst[1]); + try testing.expectEqual(@as(u8, @intCast((n >> 8) & 0xff)), dst[2]); + try testing.expectEqual(@as(u8, @intCast(n >> 16)), dst[3]); + try testing.expectEqualSlices(u8, lit[0..], dst[4 .. 4 + lit.len]); +}