authorgravatar for 4678790+dweiller@users.noreply.github.comDominic <4678790+dweiller@users.noreply.github.com> 2023-02-04 13:49:53+11:00
committergravatar for 4678790+dweiller@users.noreply.github.comDominic <4678790+dweiller@users.noreply.github.com> 2023-02-20 09:09:06+11:00
loga9c8376305d4c99591432e0ed267fad665bb4f5f
tree91d055c7966263b911f01374b865d03b8cffddc6
parent06ab5a2cd21ecd972477f067bb1d595a9ebef483

std.compress.zstandard: make ZstandardStream decode multiple frames

1 files changed, 97 insertions(+), 53 deletions(-)

lib/std/compress/zstandard.zig+97-53
......@@ -13,6 +13,7 @@ pub fn ZstandardStream(comptime ReaderType: type, comptime verify_checksum: bool
1313
1414 allocator: Allocator,
1515 in_reader: ReaderType,
16 state: enum { NewFrame, InFrame },
1617 decode_state: decompress.block.DecodeState,
1718 frame_context: decompress.FrameContext,
1819 buffer: RingBuffer,
......@@ -24,16 +25,43 @@ pub fn ZstandardStream(comptime ReaderType: type, comptime verify_checksum: bool
2425 sequence_buffer: []u8,
2526 checksum: if (verify_checksum) ?u32 else void,
2627
27 pub const Error = ReaderType.Error || error{ MalformedBlock, MalformedFrame };
28 pub const Error = ReaderType.Error || error{ ChecksumFailure, MalformedBlock, MalformedFrame, OutOfMemory };
2829
2930 pub const Reader = std.io.Reader(*Self, Error, read);
3031
3132 pub fn init(allocator: Allocator, source: ReaderType) !Self {
32 switch (try decompress.decodeFrameType(source)) {
33 .skippable => return error.SkippableFrame,
33 return Self{
34 .allocator = allocator,
35 .in_reader = source,
36 .state = .NewFrame,
37 .decode_state = undefined,
38 .frame_context = undefined,
39 .buffer = undefined,
40 .last_block = undefined,
41 .literal_fse_buffer = undefined,
42 .match_fse_buffer = undefined,
43 .offset_fse_buffer = undefined,
44 .literals_buffer = undefined,
45 .sequence_buffer = undefined,
46 .checksum = undefined,
47 };
48 }
49
50 fn frameInit(self: *Self) !void {
51 var bytes: [4]u8 = undefined;
52 const bytes_read = try self.in_reader.readAll(&bytes);
53 if (bytes_read == 0) return error.NoBytes;
54 if (bytes_read < 4) return error.EndOfStream;
55 const frame_type = try decompress.frameType(std.mem.readIntLittle(u32, &bytes));
56 switch (frame_type) {
57 .skippable => {
58 const size = try self.in_reader.readIntLittle(u32);
59 try self.in_reader.skipBytes(size, .{});
60 self.state = .NewFrame;
61 },
3462 .zstandard => {
3563 const frame_context = context: {
36 const frame_header = try decompress.decodeZstandardHeader(source);
64 const frame_header = try decompress.decodeZstandardHeader(self.in_reader);
3765 break :context try decompress.FrameContext.init(
3866 frame_header,
3967 window_size_max,
......@@ -41,56 +69,58 @@ pub fn ZstandardStream(comptime ReaderType: type, comptime verify_checksum: bool
4169 );
4270 };
4371
44 const literal_fse_buffer = try allocator.alloc(
72 const literal_fse_buffer = try self.allocator.alloc(
4573 types.compressed_block.Table.Fse,
4674 types.compressed_block.table_size_max.literal,
4775 );
48 errdefer allocator.free(literal_fse_buffer);
76 errdefer self.allocator.free(literal_fse_buffer);
4977
50 const match_fse_buffer = try allocator.alloc(
78 const match_fse_buffer = try self.allocator.alloc(
5179 types.compressed_block.Table.Fse,
5280 types.compressed_block.table_size_max.match,
5381 );
54 errdefer allocator.free(match_fse_buffer);
82 errdefer self.allocator.free(match_fse_buffer);
5583
56 const offset_fse_buffer = try allocator.alloc(
84 const offset_fse_buffer = try self.allocator.alloc(
5785 types.compressed_block.Table.Fse,
5886 types.compressed_block.table_size_max.offset,
5987 );
60 errdefer allocator.free(offset_fse_buffer);
88 errdefer self.allocator.free(offset_fse_buffer);
6189
6290 const decode_state = decompress.block.DecodeState.init(
6391 literal_fse_buffer,
6492 match_fse_buffer,
6593 offset_fse_buffer,
6694 );
67 const buffer = try RingBuffer.init(allocator, frame_context.window_size);
68
69 const literals_data = try allocator.alloc(u8, window_size_max);
70 errdefer allocator.free(literals_data);
71
72 const sequence_data = try allocator.alloc(u8, window_size_max);
73 errdefer allocator.free(sequence_data);
74
75 return Self{
76 .allocator = allocator,
77 .in_reader = source,
78 .decode_state = decode_state,
79 .frame_context = frame_context,
80 .buffer = buffer,
81 .checksum = if (verify_checksum) null else {},
82 .last_block = false,
83 .literal_fse_buffer = literal_fse_buffer,
84 .match_fse_buffer = match_fse_buffer,
85 .offset_fse_buffer = offset_fse_buffer,
86 .literals_buffer = literals_data,
87 .sequence_buffer = sequence_data,
88 };
95 const buffer = try RingBuffer.init(self.allocator, frame_context.window_size);
96
97 const literals_data = try self.allocator.alloc(u8, window_size_max);
98 errdefer self.allocator.free(literals_data);
99
100 const sequence_data = try self.allocator.alloc(u8, window_size_max);
101 errdefer self.allocator.free(sequence_data);
102
103 self.literal_fse_buffer = literal_fse_buffer;
104 self.match_fse_buffer = match_fse_buffer;
105 self.offset_fse_buffer = offset_fse_buffer;
106 self.literals_buffer = literals_data;
107 self.sequence_buffer = sequence_data;
108
109 self.buffer = buffer;
110
111 self.decode_state = decode_state;
112 self.frame_context = frame_context;
113
114 self.checksum = if (verify_checksum) null else {};
115 self.last_block = false;
116
117 self.state = .InFrame;
89118 },
90119 }
91120 }
92121
93122 pub fn deinit(self: *Self) void {
123 if (self.state == .NewFrame) return;
94124 self.allocator.free(self.decode_state.literal_fse_buffer);
95125 self.allocator.free(self.decode_state.match_fse_buffer);
96126 self.allocator.free(self.decode_state.offset_fse_buffer);
......@@ -105,6 +135,19 @@ pub fn ZstandardStream(comptime ReaderType: type, comptime verify_checksum: bool
105135
106136 pub fn read(self: *Self, buffer: []u8) Error!usize {
107137 if (buffer.len == 0) return 0;
138 while (self.state == .NewFrame) {
139 self.frameInit() catch |err| switch (err) {
140 error.NoBytes => return 0,
141 error.OutOfMemory => return error.OutOfMemory,
142 else => return error.MalformedFrame,
143 };
144 }
145
146 return self.readInner(buffer);
147 }
148
149 fn readInner(self: *Self, buffer: []u8) Error!usize {
150 std.debug.assert(self.state == .InFrame);
108151
109152 if (self.buffer.isEmpty() and !self.last_block) {
110153 const header_bytes = self.in_reader.readBytesNoEof(3) catch return error.MalformedFrame;
......@@ -127,9 +170,15 @@ pub fn ZstandardStream(comptime ReaderType: type, comptime verify_checksum: bool
127170 hasher.update(written_slice.first);
128171 hasher.update(written_slice.second);
129172 }
130 if (block_header.last_block and self.frame_context.has_checksum) {
131 const checksum = self.in_reader.readIntLittle(u32) catch return error.MalformedFrame;
132 if (verify_checksum) self.checksum = checksum;
173 if (block_header.last_block) {
174 if (self.frame_context.has_checksum) {
175 const checksum = self.in_reader.readIntLittle(u32) catch return error.MalformedFrame;
176 if (comptime verify_checksum) {
177 if (self.frame_context.hasher_opt) |*hasher| {
178 if (checksum != decompress.computeChecksum(hasher)) return error.ChecksumFailure;
179 }
180 }
181 }
133182 }
134183 }
135184
......@@ -138,18 +187,16 @@ pub fn ZstandardStream(comptime ReaderType: type, comptime verify_checksum: bool
138187 while (written_count < decoded_data_len and written_count < buffer.len) : (written_count += 1) {
139188 buffer[written_count] = self.buffer.read().?;
140189 }
141 return written_count;
142 }
143
144 pub fn verifyChecksum(self: *Self) !bool {
145 if (verify_checksum) {
146 if (self.checksum) |checksum| {
147 if (self.frame_context.hasher_opt) |*hasher| {
148 return checksum == decompress.computeChecksum(hasher);
149 }
150 }
190 if (self.buffer.len() == 0) {
191 self.state = .NewFrame;
192 self.allocator.free(self.literal_fse_buffer);
193 self.allocator.free(self.match_fse_buffer);
194 self.allocator.free(self.offset_fse_buffer);
195 self.allocator.free(self.literals_buffer);
196 self.allocator.free(self.sequence_buffer);
197 self.buffer.deinit(self.allocator);
151198 }
152 return true;
199 return written_count;
153200 }
154201 };
155202}
......@@ -163,7 +210,6 @@ fn testDecompress(data: []const u8) ![]u8 {
163210 var stream = try zstandardStream(std.testing.allocator, in_stream.reader());
164211 defer stream.deinit();
165212 const result = stream.reader().readAllAlloc(std.testing.allocator, std.math.maxInt(usize));
166 try std.testing.expect(try stream.verifyChecksum());
167213 return result;
168214}
169215
......@@ -181,14 +227,12 @@ test "decompression" {
181227 var buffer = try std.testing.allocator.alloc(u8, uncompressed.len);
182228 defer std.testing.allocator.free(buffer);
183229
184 const res3 = try decompress.decodeFrame(buffer, compressed3, true);
185 try std.testing.expectEqual(compressed3.len, res3.read_count);
186 try std.testing.expectEqual(uncompressed.len, res3.write_count);
230 const res3 = try decompress.decode(buffer, compressed3, true);
231 try std.testing.expectEqual(uncompressed.len, res3);
187232 try std.testing.expectEqualSlices(u8, uncompressed, buffer);
188233
189 const res19 = try decompress.decodeFrame(buffer, compressed19, true);
190 try std.testing.expectEqual(compressed19.len, res19.read_count);
191 try std.testing.expectEqual(uncompressed.len, res19.write_count);
234 const res19 = try decompress.decode(buffer, compressed19, true);
235 try std.testing.expectEqual(uncompressed.len, res19);
192236 try std.testing.expectEqualSlices(u8, uncompressed, buffer);
193237
194238 try testReader(compressed3, uncompressed);