authorgravatar for 4678790+dweiller@users.noreply.github.comDominic <4678790+dweiller@users.noreply.github.com> 2024-02-15 16:19:55+11:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-02-23 02:37:11-07:00
logac1b957e79e6a150c0dc8f835fe956d6bb3d3372
tree252e022fbe2bf6ac4ab7894895e2219a71143685
parent73f6d3afb5de4ee223d61aa1f400a5c58b5dbea1

std.compress.zstd: remove allocation from DecompressStream

4 files changed, 61 insertions(+), 87 deletions(-)

lib/std/compress/zstandard.zig+51-82
......@@ -1,5 +1,4 @@
11const std = @import("std");
2const Allocator = std.mem.Allocator;
32const RingBuffer = std.RingBuffer;
43
54const types = @import("zstandard/types.zig");
......@@ -10,7 +9,9 @@ pub const decompress = @import("zstandard/decompress.zig");
109
1110pub const DecompressStreamOptions = struct {
1211 verify_checksum: bool = true,
13 window_size_max: usize = 1 << 23, // 8MiB default maximum window size
12 window_size_max: usize = default_window_size_max,
13
14 pub const default_window_size_max = 1 << 23; // 8MiB default maximum window size
1415};
1516
1617pub fn DecompressStream(
......@@ -20,20 +21,29 @@ pub fn DecompressStream(
2021 return struct {
2122 const Self = @This();
2223
23 allocator: Allocator,
24 pub const window_size_max = options.window_size_max;
25
26 const table_size_max = types.compressed_block.table_size_max;
27
2428 source: std.io.CountingReader(ReaderType),
2529 state: enum { NewFrame, InFrame, LastBlock },
2630 decode_state: decompress.block.DecodeState,
2731 frame_context: decompress.FrameContext,
28 buffer: RingBuffer,
29 literal_fse_buffer: []types.compressed_block.Table.Fse,
30 match_fse_buffer: []types.compressed_block.Table.Fse,
31 offset_fse_buffer: []types.compressed_block.Table.Fse,
32 literals_buffer: []u8,
33 sequence_buffer: []u8,
32 buffer: WindowBuffer,
33 literal_fse_buffer: [table_size_max.literal]types.compressed_block.Table.Fse,
34 match_fse_buffer: [table_size_max.match]types.compressed_block.Table.Fse,
35 offset_fse_buffer: [table_size_max.offset]types.compressed_block.Table.Fse,
36 literals_buffer: [types.block_size_max]u8,
37 sequence_buffer: [types.block_size_max]u8,
3438 checksum: if (options.verify_checksum) ?u32 else void,
3539 current_frame_decompressed_size: usize,
3640
41 const WindowBuffer = struct {
42 data: *[options.window_size_max]u8 = undefined,
43 read_index: usize = 0,
44 write_index: usize = 0,
45 };
46
3747 pub const Error = ReaderType.Error || error{
3848 ChecksumFailure,
3949 DictionaryIdFlagUnsupported,
......@@ -44,14 +54,13 @@ pub fn DecompressStream(
4454
4555 pub const Reader = std.io.Reader(*Self, Error, read);
4656
47 pub fn init(allocator: Allocator, source: ReaderType) Self {
57 pub fn init(source: ReaderType, window_buffer: *[options.window_size_max]u8) Self {
4858 return Self{
49 .allocator = allocator,
5059 .source = std.io.countingReader(source),
5160 .state = .NewFrame,
5261 .decode_state = undefined,
5362 .frame_context = undefined,
54 .buffer = undefined,
63 .buffer = .{ .data = window_buffer },
5564 .literal_fse_buffer = undefined,
5665 .match_fse_buffer = undefined,
5766 .offset_fse_buffer = undefined,
......@@ -76,44 +85,11 @@ pub fn DecompressStream(
7685 options.verify_checksum,
7786 );
7887
79 const literal_fse_buffer = try self.allocator.alloc(
80 types.compressed_block.Table.Fse,
81 types.compressed_block.table_size_max.literal,
82 );
83 errdefer self.allocator.free(literal_fse_buffer);
84
85 const match_fse_buffer = try self.allocator.alloc(
86 types.compressed_block.Table.Fse,
87 types.compressed_block.table_size_max.match,
88 );
89 errdefer self.allocator.free(match_fse_buffer);
90
91 const offset_fse_buffer = try self.allocator.alloc(
92 types.compressed_block.Table.Fse,
93 types.compressed_block.table_size_max.offset,
94 );
95 errdefer self.allocator.free(offset_fse_buffer);
96
9788 const decode_state = decompress.block.DecodeState.init(
98 literal_fse_buffer,
99 match_fse_buffer,
100 offset_fse_buffer,
89 &self.literal_fse_buffer,
90 &self.match_fse_buffer,
91 &self.offset_fse_buffer,
10192 );
102 const buffer = try RingBuffer.init(self.allocator, frame_context.window_size);
103
104 const literals_data = try self.allocator.alloc(u8, frame_context.block_size_max);
105 errdefer self.allocator.free(literals_data);
106
107 const sequence_data = try self.allocator.alloc(u8, frame_context.block_size_max);
108 errdefer self.allocator.free(sequence_data);
109
110 self.literal_fse_buffer = literal_fse_buffer;
111 self.match_fse_buffer = match_fse_buffer;
112 self.offset_fse_buffer = offset_fse_buffer;
113 self.literals_buffer = literals_data;
114 self.sequence_buffer = sequence_data;
115
116 self.buffer = buffer;
11793
11894 self.decode_state = decode_state;
11995 self.frame_context = frame_context;
......@@ -126,16 +102,6 @@ pub fn DecompressStream(
126102 }
127103 }
128104
129 pub fn deinit(self: *Self) void {
130 if (self.state == .NewFrame) return;
131 self.allocator.free(self.decode_state.literal_fse_buffer);
132 self.allocator.free(self.decode_state.match_fse_buffer);
133 self.allocator.free(self.decode_state.offset_fse_buffer);
134 self.allocator.free(self.literals_buffer);
135 self.allocator.free(self.sequence_buffer);
136 self.buffer.deinit(self.allocator);
137 }
138
139105 pub fn reader(self: *Self) Reader {
140106 return .{ .context = self };
141107 }
......@@ -153,7 +119,6 @@ pub fn DecompressStream(
153119 0
154120 else
155121 error.MalformedFrame,
156 error.OutOfMemory => return error.OutOfMemory,
157122 else => return error.MalformedFrame,
158123 };
159124 }
......@@ -165,20 +130,30 @@ pub fn DecompressStream(
165130 fn readInner(self: *Self, buffer: []u8) Error!usize {
166131 std.debug.assert(self.state != .NewFrame);
167132
133 var ring_buffer = RingBuffer{
134 .data = self.buffer.data,
135 .read_index = self.buffer.read_index,
136 .write_index = self.buffer.write_index,
137 };
138 defer {
139 self.buffer.read_index = ring_buffer.read_index;
140 self.buffer.write_index = ring_buffer.write_index;
141 }
142
168143 const source_reader = self.source.reader();
169 while (self.buffer.isEmpty() and self.state != .LastBlock) {
144 while (ring_buffer.isEmpty() and self.state != .LastBlock) {
170145 const header_bytes = source_reader.readBytesNoEof(3) catch
171146 return error.MalformedFrame;
172147 const block_header = decompress.block.decodeBlockHeader(&header_bytes);
173148
174149 decompress.block.decodeBlockReader(
175 &self.buffer,
150 &ring_buffer,
176151 source_reader,
177152 block_header,
178153 &self.decode_state,
179154 self.frame_context.block_size_max,
180 self.literals_buffer,
181 self.sequence_buffer,
155 &self.literals_buffer,
156 &self.sequence_buffer,
182157 ) catch
183158 return error.MalformedBlock;
184159
......@@ -186,12 +161,12 @@ pub fn DecompressStream(
186161 if (self.current_frame_decompressed_size > size) return error.MalformedFrame;
187162 }
188163
189 const size = self.buffer.len();
164 const size = ring_buffer.len();
190165 self.current_frame_decompressed_size += size;
191166
192167 if (self.frame_context.hasher_opt) |*hasher| {
193168 if (size > 0) {
194 const written_slice = self.buffer.sliceLast(size);
169 const written_slice = ring_buffer.sliceLast(size);
195170 hasher.update(written_slice.first);
196171 hasher.update(written_slice.second);
197172 }
......@@ -216,18 +191,12 @@ pub fn DecompressStream(
216191 }
217192 }
218193
219 const size = @min(self.buffer.len(), buffer.len);
194 const size = @min(ring_buffer.len(), buffer.len);
220195 if (size > 0) {
221 self.buffer.readFirstAssumeLength(buffer, size);
196 ring_buffer.readFirstAssumeLength(buffer, size);
222197 }
223 if (self.state == .LastBlock and self.buffer.len() == 0) {
198 if (self.state == .LastBlock and ring_buffer.len() == 0) {
224199 self.state = .NewFrame;
225 self.allocator.free(self.literal_fse_buffer);
226 self.allocator.free(self.match_fse_buffer);
227 self.allocator.free(self.offset_fse_buffer);
228 self.allocator.free(self.literals_buffer);
229 self.allocator.free(self.sequence_buffer);
230 self.buffer.deinit(self.allocator);
231200 }
232201 return size;
233202 }
......@@ -235,24 +204,24 @@ pub fn DecompressStream(
235204}
236205
237206pub fn decompressStreamOptions(
238 allocator: Allocator,
239207 reader: anytype,
240208 comptime options: DecompressStreamOptions,
209 window_buffer: *[options.window_size_max]u8,
241210) DecompressStream(@TypeOf(reader), options) {
242 return DecompressStream(@TypeOf(reader), options).init(allocator, reader);
211 return DecompressStream(@TypeOf(reader), options).init(reader, window_buffer);
243212}
244213
245214pub fn decompressStream(
246 allocator: Allocator,
247215 reader: anytype,
216 window_buffer: *[DecompressStreamOptions.default_window_size_max]u8,
248217) DecompressStream(@TypeOf(reader), .{}) {
249 return DecompressStream(@TypeOf(reader), .{}).init(allocator, reader);
218 return DecompressStream(@TypeOf(reader), .{}).init(reader, window_buffer);
250219}
251220
252221fn testDecompress(data: []const u8) ![]u8 {
222 var window_buffer: [DecompressStreamOptions.default_window_size_max]u8 = undefined;
253223 var in_stream = std.io.fixedBufferStream(data);
254 var zstd_stream = decompressStream(std.testing.allocator, in_stream.reader());
255 defer zstd_stream.deinit();
224 var zstd_stream = decompressStream(in_stream.reader(), &window_buffer);
256225 const result = zstd_stream.reader().readAllAlloc(std.testing.allocator, std.math.maxInt(usize));
257226 return result;
258227}
......@@ -301,9 +270,9 @@ fn expectEqualDecoded(expected: []const u8, input: []const u8) !void {
301270 }
302271
303272 {
273 var window_buffer: [DecompressStreamOptions.default_window_size_max]u8 = undefined;
304274 var in_stream = std.io.fixedBufferStream(input);
305 var stream = decompressStream(allocator, in_stream.reader());
306 defer stream.deinit();
275 var stream = decompressStream(in_stream.reader(), &window_buffer);
307276
308277 const result = try stream.reader().readAllAlloc(allocator, std.math.maxInt(usize));
309278 defer allocator.free(result);
lib/std/compress/zstandard/decompress.zig+1-1
......@@ -409,7 +409,7 @@ pub const FrameContext = struct {
409409 .hasher_opt = if (should_compute_checksum) std.hash.XxHash64.init(0) else null,
410410 .window_size = window_size,
411411 .has_checksum = frame_header.descriptor.content_checksum_flag,
412 .block_size_max = @min(1 << 17, window_size),
412 .block_size_max = @min(types.block_size_max, window_size),
413413 .content_size = content_size,
414414 };
415415 }
lib/std/compress/zstandard/types.zig+2
......@@ -1,3 +1,5 @@
1pub const block_size_max = 1 << 17;
2
13pub const frame = struct {
24 pub const Kind = enum { zstandard, skippable };
35
src/Package/Fetch.zig+7-4
......@@ -1126,11 +1126,14 @@ fn unpackResource(
11261126// wrapped for generic use in unpackTarballCompressed: see github.com/ziglang/zig/issues/14739
11271127const ZstdWrapper = struct {
11281128 fn DecompressType(comptime T: type) type {
1129 return error{}!std.compress.zstd.DecompressStream(T, .{});
1129 return Allocator.Error!std.compress.zstd.DecompressStream(T, .{});
11301130 }
11311131
11321132 fn decompress(allocator: Allocator, reader: anytype) DecompressType(@TypeOf(reader)) {
1133 return std.compress.zstd.decompressStream(allocator, reader);
1133 const window_size = std.compress.zstd.DecompressStreamOptions.default_window_size_max;
1134 const window_buffer = try allocator.create([window_size]u8);
1135 defer allocator.destroy(window_buffer);
1136 return std.compress.zstd.decompressStream(reader, window_buffer);
11341137 }
11351138};
11361139
......@@ -1138,7 +1141,7 @@ fn unpackTarballCompressed(
11381141 f: *Fetch,
11391142 out_dir: fs.Dir,
11401143 resource: *Resource,
1141 comptime Compression: type,
1144 Compression: anytype,
11421145) RunError!void {
11431146 const gpa = f.arena.child_allocator;
11441147 const eb = &f.error_bundle;
......@@ -1151,7 +1154,7 @@ fn unpackTarballCompressed(
11511154 .{@errorName(err)},
11521155 ));
11531156 };
1154 defer decompress.deinit();
1157 defer if (@hasDecl(Compression, "deinit")) decompress.deinit();
11551158
11561159 return unpackTarball(f, out_dir, decompress.reader());
11571160}