authorgravatar for 4678790+dweiller@users.noreply.github.comDominic <4678790+dweiller@users.noreply.github.com> 2023-02-02 16:20:03+11:00
committergravatar for 4678790+dweiller@users.noreply.github.comDominic <4678790+dweiller@users.noreply.github.com> 2023-02-20 09:09:06+11:00
loga180fcc93d3eff8add7b0344ddc77241dcf78f1e
tree9a20403eec7a5ac79a76aad176c8fabe3be83df1
parent5723291444116419440a187adcfa5ecb9557544e

std.compress.zstandard: add `ZstandardStream`

1 files changed, 153 insertions(+), 0 deletions(-)

lib/std/compress/zstandard.zig+153
......@@ -1,8 +1,158 @@
11const std = @import("std");
2const Allocator = std.mem.Allocator;
23
4const types = @import("zstandard/types.zig");
5
6const RingBuffer = @import("zstandard/RingBuffer.zig");
37pub const decompress = @import("zstandard/decompress.zig");
48pub usingnamespace @import("zstandard/types.zig");
59
10pub fn ZstandardStream(comptime ReaderType: type, comptime verify_checksum: bool, comptime window_size_max: usize) type {
11 return struct {
12 const Self = @This();
13
14 allocator: Allocator,
15 in_reader: ReaderType,
16 decode_state: decompress.DecodeState,
17 frame_context: decompress.FrameContext,
18 buffer: RingBuffer,
19 last_block: bool,
20 literal_fse_buffer: []types.compressed_block.Table.Fse,
21 match_fse_buffer: []types.compressed_block.Table.Fse,
22 offset_fse_buffer: []types.compressed_block.Table.Fse,
23 literals_buffer: []u8,
24 sequence_buffer: []u8,
25 checksum: if (verify_checksum) ?u32 else void,
26
27 pub const Error = ReaderType.Error || error{ MalformedBlock, MalformedFrame, EndOfStream };
28
29 pub const Reader = std.io.Reader(*Self, Error, read);
30
31 pub fn init(allocator: Allocator, source: ReaderType) !Self {
32 switch (try decompress.decodeFrameType(source)) {
33 .skippable => return error.SkippableFrame,
34 .zstandard => {
35 const frame_context = context: {
36 const frame_header = try decompress.decodeZStandardHeader(source);
37 break :context try decompress.FrameContext.init(frame_header, window_size_max, verify_checksum);
38 };
39
40 const literal_fse_buffer = try allocator.alloc(types.compressed_block.Table.Fse, types.compressed_block.table_size_max.literal);
41 errdefer allocator.free(literal_fse_buffer);
42 const match_fse_buffer = try allocator.alloc(types.compressed_block.Table.Fse, types.compressed_block.table_size_max.match);
43 errdefer allocator.free(match_fse_buffer);
44 const offset_fse_buffer = try allocator.alloc(types.compressed_block.Table.Fse, types.compressed_block.table_size_max.offset);
45 errdefer allocator.free(offset_fse_buffer);
46
47 const decode_state = decompress.DecodeState.init(literal_fse_buffer, match_fse_buffer, offset_fse_buffer);
48 const buffer = try RingBuffer.init(allocator, frame_context.window_size);
49
50 const literals_data = try allocator.alloc(u8, window_size_max);
51 errdefer allocator.free(literals_data);
52 const sequence_data = try allocator.alloc(u8, window_size_max);
53 errdefer allocator.free(sequence_data);
54
55 return Self{
56 .allocator = allocator,
57 .in_reader = source,
58 .decode_state = decode_state,
59 .frame_context = frame_context,
60 .buffer = buffer,
61 .checksum = if (verify_checksum) null else {},
62 .last_block = false,
63 .literal_fse_buffer = literal_fse_buffer,
64 .match_fse_buffer = match_fse_buffer,
65 .offset_fse_buffer = offset_fse_buffer,
66 .literals_buffer = literals_data,
67 .sequence_buffer = sequence_data,
68 };
69 },
70 }
71 }
72
73 pub fn deinit(self: *Self) void {
74 self.allocator.free(self.decode_state.literal_fse_buffer);
75 self.allocator.free(self.decode_state.match_fse_buffer);
76 self.allocator.free(self.decode_state.offset_fse_buffer);
77 self.allocator.free(self.literals_buffer);
78 self.allocator.free(self.sequence_buffer);
79 self.buffer.deinit(self.allocator);
80 }
81
82 pub fn reader(self: *Self) Reader {
83 return .{ .context = self };
84 }
85
86 pub fn read(self: *Self, buffer: []u8) Error!usize {
87 if (buffer.len == 0) return 0;
88
89 if (self.buffer.isEmpty() and !self.last_block) {
90 const header_bytes = try self.in_reader.readBytesNoEof(3);
91 const block_header = decompress.decodeBlockHeader(&header_bytes);
92
93 decompress.decodeBlockReader(
94 &self.buffer,
95 self.in_reader,
96 block_header,
97 &self.decode_state,
98 self.frame_context.block_size_max,
99 self.literals_buffer,
100 self.sequence_buffer,
101 ) catch
102 return error.MalformedBlock;
103
104 self.last_block = block_header.last_block;
105 if (self.frame_context.hasher_opt) |*hasher| {
106 const written_slice = self.buffer.sliceLast(self.buffer.len());
107 hasher.update(written_slice.first);
108 hasher.update(written_slice.second);
109 }
110 if (block_header.last_block and self.frame_context.has_checksum) {
111 const checksum = self.in_reader.readIntLittle(u32) catch return error.MalformedFrame;
112 if (verify_checksum) self.checksum = checksum;
113 }
114 }
115
116 const decoded_data_len = self.buffer.len();
117 var written_count: usize = 0;
118 while (written_count < decoded_data_len and written_count < buffer.len) : (written_count += 1) {
119 buffer[written_count] = self.buffer.read().?;
120 }
121 return written_count;
122 }
123
124 pub fn verifyChecksum(self: *Self) !bool {
125 if (verify_checksum) {
126 if (self.checksum) |checksum| {
127 if (self.frame_context.hasher_opt) |*hasher| {
128 return checksum == decompress.computeChecksum(hasher);
129 }
130 }
131 }
132 return true;
133 }
134 };
135}
136
137pub fn zstandardStream(allocator: Allocator, reader: anytype) !ZstandardStream(@TypeOf(reader), true, 8 * (1 << 20)) {
138 return ZstandardStream(@TypeOf(reader), true, 8 * (1 << 20)).init(allocator, reader);
139}
140
141fn testDecompress(data: []const u8) ![]u8 {
142 var in_stream = std.io.fixedBufferStream(data);
143 var stream = try zstandardStream(std.testing.allocator, in_stream.reader());
144 defer stream.deinit();
145 const result = stream.reader().readAllAlloc(std.testing.allocator, std.math.maxInt(usize));
146 try std.testing.expect(try stream.verifyChecksum());
147 return result;
148}
149
150fn testReader(data: []const u8, comptime expected: []const u8) !void {
151 const buf = try testDecompress(data);
152 defer std.testing.allocator.free(buf);
153 try std.testing.expectEqualSlices(u8, expected, buf);
154}
155
6156test "decompression" {
7157 const uncompressed = @embedFile("testdata/rfc8478.txt");
8158 const compressed3 = @embedFile("testdata/rfc8478.txt.zst.3");
......@@ -20,4 +170,7 @@ test "decompression" {
20170 try std.testing.expectEqual(compressed19.len, res19.read_count);
21171 try std.testing.expectEqual(uncompressed.len, res19.write_count);
22172 try std.testing.expectEqualSlices(u8, uncompressed, buffer);
173
174 try testReader(compressed3, uncompressed);
175 try testReader(compressed19, uncompressed);
23176}