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 @@...@@ -1,5 +1,4 @@
1const std = @import("std");1const std = @import("std");
2const Allocator = std.mem.Allocator;
3const RingBuffer = std.RingBuffer;2const RingBuffer = std.RingBuffer;
43
5const types = @import("zstandard/types.zig");4const types = @import("zstandard/types.zig");
...@@ -10,7 +9,9 @@ pub const decompress = @import("zstandard/decompress.zig");...@@ -10,7 +9,9 @@ pub const decompress = @import("zstandard/decompress.zig");
109
11pub const DecompressStreamOptions = struct {10pub const DecompressStreamOptions = struct {
12 verify_checksum: bool = true,11 verify_checksum: bool = true,
13 window_size_max: usize = 1 << 23, // 8MiB default maximum window size12 window_size_max: usize = default_window_size_max,
13
14 pub const default_window_size_max = 1 << 23; // 8MiB default maximum window size
14};15};
1516
16pub fn DecompressStream(17pub fn DecompressStream(
...@@ -20,20 +21,29 @@ pub fn DecompressStream(...@@ -20,20 +21,29 @@ pub fn DecompressStream(
20 return struct {21 return struct {
21 const Self = @This();22 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
24 source: std.io.CountingReader(ReaderType),28 source: std.io.CountingReader(ReaderType),
25 state: enum { NewFrame, InFrame, LastBlock },29 state: enum { NewFrame, InFrame, LastBlock },
26 decode_state: decompress.block.DecodeState,30 decode_state: decompress.block.DecodeState,
27 frame_context: decompress.FrameContext,31 frame_context: decompress.FrameContext,
28 buffer: RingBuffer,32 buffer: WindowBuffer,
29 literal_fse_buffer: []types.compressed_block.Table.Fse,33 literal_fse_buffer: [table_size_max.literal]types.compressed_block.Table.Fse,
30 match_fse_buffer: []types.compressed_block.Table.Fse,34 match_fse_buffer: [table_size_max.match]types.compressed_block.Table.Fse,
31 offset_fse_buffer: []types.compressed_block.Table.Fse,35 offset_fse_buffer: [table_size_max.offset]types.compressed_block.Table.Fse,
32 literals_buffer: []u8,36 literals_buffer: [types.block_size_max]u8,
33 sequence_buffer: []u8,37 sequence_buffer: [types.block_size_max]u8,
34 checksum: if (options.verify_checksum) ?u32 else void,38 checksum: if (options.verify_checksum) ?u32 else void,
35 current_frame_decompressed_size: usize,39 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
37 pub const Error = ReaderType.Error || error{47 pub const Error = ReaderType.Error || error{
38 ChecksumFailure,48 ChecksumFailure,
39 DictionaryIdFlagUnsupported,49 DictionaryIdFlagUnsupported,
...@@ -44,14 +54,13 @@ pub fn DecompressStream(...@@ -44,14 +54,13 @@ pub fn DecompressStream(
4454
45 pub const Reader = std.io.Reader(*Self, Error, read);55 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 {
48 return Self{58 return Self{
49 .allocator = allocator,
50 .source = std.io.countingReader(source),59 .source = std.io.countingReader(source),
51 .state = .NewFrame,60 .state = .NewFrame,
52 .decode_state = undefined,61 .decode_state = undefined,
53 .frame_context = undefined,62 .frame_context = undefined,
54 .buffer = undefined,63 .buffer = .{ .data = window_buffer },
55 .literal_fse_buffer = undefined,64 .literal_fse_buffer = undefined,
56 .match_fse_buffer = undefined,65 .match_fse_buffer = undefined,
57 .offset_fse_buffer = undefined,66 .offset_fse_buffer = undefined,
...@@ -76,44 +85,11 @@ pub fn DecompressStream(...@@ -76,44 +85,11 @@ pub fn DecompressStream(
76 options.verify_checksum,85 options.verify_checksum,
77 );86 );
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
97 const decode_state = decompress.block.DecodeState.init(88 const decode_state = decompress.block.DecodeState.init(
98 literal_fse_buffer,89 &self.literal_fse_buffer,
99 match_fse_buffer,90 &self.match_fse_buffer,
100 offset_fse_buffer,91 &self.offset_fse_buffer,
101 );92 );
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
118 self.decode_state = decode_state;94 self.decode_state = decode_state;
119 self.frame_context = frame_context;95 self.frame_context = frame_context;
...@@ -126,16 +102,6 @@ pub fn DecompressStream(...@@ -126,16 +102,6 @@ pub fn DecompressStream(
126 }102 }
127 }103 }
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
139 pub fn reader(self: *Self) Reader {105 pub fn reader(self: *Self) Reader {
140 return .{ .context = self };106 return .{ .context = self };
141 }107 }
...@@ -153,7 +119,6 @@ pub fn DecompressStream(...@@ -153,7 +119,6 @@ pub fn DecompressStream(
153 0119 0
154 else120 else
155 error.MalformedFrame,121 error.MalformedFrame,
156 error.OutOfMemory => return error.OutOfMemory,
157 else => return error.MalformedFrame,122 else => return error.MalformedFrame,
158 };123 };
159 }124 }
...@@ -165,20 +130,30 @@ pub fn DecompressStream(...@@ -165,20 +130,30 @@ pub fn DecompressStream(
165 fn readInner(self: *Self, buffer: []u8) Error!usize {130 fn readInner(self: *Self, buffer: []u8) Error!usize {
166 std.debug.assert(self.state != .NewFrame);131 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
168 const source_reader = self.source.reader();143 const source_reader = self.source.reader();
169 while (self.buffer.isEmpty() and self.state != .LastBlock) {144 while (ring_buffer.isEmpty() and self.state != .LastBlock) {
170 const header_bytes = source_reader.readBytesNoEof(3) catch145 const header_bytes = source_reader.readBytesNoEof(3) catch
171 return error.MalformedFrame;146 return error.MalformedFrame;
172 const block_header = decompress.block.decodeBlockHeader(&header_bytes);147 const block_header = decompress.block.decodeBlockHeader(&header_bytes);
173148
174 decompress.block.decodeBlockReader(149 decompress.block.decodeBlockReader(
175 &self.buffer,150 &ring_buffer,
176 source_reader,151 source_reader,
177 block_header,152 block_header,
178 &self.decode_state,153 &self.decode_state,
179 self.frame_context.block_size_max,154 self.frame_context.block_size_max,
180 self.literals_buffer,155 &self.literals_buffer,
181 self.sequence_buffer,156 &self.sequence_buffer,
182 ) catch157 ) catch
183 return error.MalformedBlock;158 return error.MalformedBlock;
184159
...@@ -186,12 +161,12 @@ pub fn DecompressStream(...@@ -186,12 +161,12 @@ pub fn DecompressStream(
186 if (self.current_frame_decompressed_size > size) return error.MalformedFrame;161 if (self.current_frame_decompressed_size > size) return error.MalformedFrame;
187 }162 }
188163
189 const size = self.buffer.len();164 const size = ring_buffer.len();
190 self.current_frame_decompressed_size += size;165 self.current_frame_decompressed_size += size;
191166
192 if (self.frame_context.hasher_opt) |*hasher| {167 if (self.frame_context.hasher_opt) |*hasher| {
193 if (size > 0) {168 if (size > 0) {
194 const written_slice = self.buffer.sliceLast(size);169 const written_slice = ring_buffer.sliceLast(size);
195 hasher.update(written_slice.first);170 hasher.update(written_slice.first);
196 hasher.update(written_slice.second);171 hasher.update(written_slice.second);
197 }172 }
...@@ -216,18 +191,12 @@ pub fn DecompressStream(...@@ -216,18 +191,12 @@ pub fn DecompressStream(
216 }191 }
217 }192 }
218193
219 const size = @min(self.buffer.len(), buffer.len);194 const size = @min(ring_buffer.len(), buffer.len);
220 if (size > 0) {195 if (size > 0) {
221 self.buffer.readFirstAssumeLength(buffer, size);196 ring_buffer.readFirstAssumeLength(buffer, size);
222 }197 }
223 if (self.state == .LastBlock and self.buffer.len() == 0) {198 if (self.state == .LastBlock and ring_buffer.len() == 0) {
224 self.state = .NewFrame;199 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);
231 }200 }
232 return size;201 return size;
233 }202 }
...@@ -235,24 +204,24 @@ pub fn DecompressStream(...@@ -235,24 +204,24 @@ pub fn DecompressStream(
235}204}
236205
237pub fn decompressStreamOptions(206pub fn decompressStreamOptions(
238 allocator: Allocator,
239 reader: anytype,207 reader: anytype,
240 comptime options: DecompressStreamOptions,208 comptime options: DecompressStreamOptions,
209 window_buffer: *[options.window_size_max]u8,
241) DecompressStream(@TypeOf(reader), options) {210) DecompressStream(@TypeOf(reader), options) {
242 return DecompressStream(@TypeOf(reader), options).init(allocator, reader);211 return DecompressStream(@TypeOf(reader), options).init(reader, window_buffer);
243}212}
244213
245pub fn decompressStream(214pub fn decompressStream(
246 allocator: Allocator,
247 reader: anytype,215 reader: anytype,
216 window_buffer: *[DecompressStreamOptions.default_window_size_max]u8,
248) DecompressStream(@TypeOf(reader), .{}) {217) DecompressStream(@TypeOf(reader), .{}) {
249 return DecompressStream(@TypeOf(reader), .{}).init(allocator, reader);218 return DecompressStream(@TypeOf(reader), .{}).init(reader, window_buffer);
250}219}
251220
252fn testDecompress(data: []const u8) ![]u8 {221fn testDecompress(data: []const u8) ![]u8 {
222 var window_buffer: [DecompressStreamOptions.default_window_size_max]u8 = undefined;
253 var in_stream = std.io.fixedBufferStream(data);223 var in_stream = std.io.fixedBufferStream(data);
254 var zstd_stream = decompressStream(std.testing.allocator, in_stream.reader());224 var zstd_stream = decompressStream(in_stream.reader(), &window_buffer);
255 defer zstd_stream.deinit();
256 const result = zstd_stream.reader().readAllAlloc(std.testing.allocator, std.math.maxInt(usize));225 const result = zstd_stream.reader().readAllAlloc(std.testing.allocator, std.math.maxInt(usize));
257 return result;226 return result;
258}227}
...@@ -301,9 +270,9 @@ fn expectEqualDecoded(expected: []const u8, input: []const u8) !void {...@@ -301,9 +270,9 @@ fn expectEqualDecoded(expected: []const u8, input: []const u8) !void {
301 }270 }
302271
303 {272 {
273 var window_buffer: [DecompressStreamOptions.default_window_size_max]u8 = undefined;
304 var in_stream = std.io.fixedBufferStream(input);274 var in_stream = std.io.fixedBufferStream(input);
305 var stream = decompressStream(allocator, in_stream.reader());275 var stream = decompressStream(in_stream.reader(), &window_buffer);
306 defer stream.deinit();
307276
308 const result = try stream.reader().readAllAlloc(allocator, std.math.maxInt(usize));277 const result = try stream.reader().readAllAlloc(allocator, std.math.maxInt(usize));
309 defer allocator.free(result);278 defer allocator.free(result);
lib/std/compress/zstandard/decompress.zig+1-1
...@@ -409,7 +409,7 @@ pub const FrameContext = struct {...@@ -409,7 +409,7 @@ pub const FrameContext = struct {
409 .hasher_opt = if (should_compute_checksum) std.hash.XxHash64.init(0) else null,409 .hasher_opt = if (should_compute_checksum) std.hash.XxHash64.init(0) else null,
410 .window_size = window_size,410 .window_size = window_size,
411 .has_checksum = frame_header.descriptor.content_checksum_flag,411 .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),
413 .content_size = content_size,413 .content_size = content_size,
414 };414 };
415 }415 }
lib/std/compress/zstandard/types.zig+2
...@@ -1,3 +1,5 @@...@@ -1,3 +1,5 @@
1pub const block_size_max = 1 << 17;
2
1pub const frame = struct {3pub const frame = struct {
2 pub const Kind = enum { zstandard, skippable };4 pub const Kind = enum { zstandard, skippable };
35
src/Package/Fetch.zig+7-4
...@@ -1126,11 +1126,14 @@ fn unpackResource(...@@ -1126,11 +1126,14 @@ fn unpackResource(
1126// wrapped for generic use in unpackTarballCompressed: see github.com/ziglang/zig/issues/147391126// wrapped for generic use in unpackTarballCompressed: see github.com/ziglang/zig/issues/14739
1127const ZstdWrapper = struct {1127const ZstdWrapper = struct {
1128 fn DecompressType(comptime T: type) type {1128 fn DecompressType(comptime T: type) type {
1129 return error{}!std.compress.zstd.DecompressStream(T, .{});1129 return Allocator.Error!std.compress.zstd.DecompressStream(T, .{});
1130 }1130 }
11311131
1132 fn decompress(allocator: Allocator, reader: anytype) DecompressType(@TypeOf(reader)) {1132 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);
1134 }1137 }
1135};1138};
11361139
...@@ -1138,7 +1141,7 @@ fn unpackTarballCompressed(...@@ -1138,7 +1141,7 @@ fn unpackTarballCompressed(
1138 f: *Fetch,1141 f: *Fetch,
1139 out_dir: fs.Dir,1142 out_dir: fs.Dir,
1140 resource: *Resource,1143 resource: *Resource,
1141 comptime Compression: type,1144 Compression: anytype,
1142) RunError!void {1145) RunError!void {
1143 const gpa = f.arena.child_allocator;1146 const gpa = f.arena.child_allocator;
1144 const eb = &f.error_bundle;1147 const eb = &f.error_bundle;
...@@ -1151,7 +1154,7 @@ fn unpackTarballCompressed(...@@ -1151,7 +1154,7 @@ fn unpackTarballCompressed(
1151 .{@errorName(err)},1154 .{@errorName(err)},
1152 ));1155 ));
1153 };1156 };
1154 defer decompress.deinit();1157 defer if (@hasDecl(Compression, "deinit")) decompress.deinit();
11551158
1156 return unpackTarball(f, out_dir, decompress.reader());1159 return unpackTarball(f, out_dir, decompress.reader());
1157}1160}