From ac1b957e79e6a150c0dc8f835fe956d6bb3d3372 Mon Sep 17 00:00:00 2001 From: dweiller <4678790+dweiller@users.noreply.github.com> Date: Thu, 15 Feb 2024 16:19:55 +1100 Subject: [PATCH] std.compress.zstd: remove allocation from DecompressStream --- lib/std/compress/zstandard.zig | 133 +++++++++------------- lib/std/compress/zstandard/decompress.zig | 2 +- lib/std/compress/zstandard/types.zig | 2 + src/Package/Fetch.zig | 11 +- 4 files changed, 61 insertions(+), 87 deletions(-) diff --git a/lib/std/compress/zstandard.zig b/lib/std/compress/zstandard.zig index 89879f88d1..352f816625 100644 --- a/lib/std/compress/zstandard.zig +++ b/lib/std/compress/zstandard.zig @@ -1,5 +1,4 @@ const std = @import("std"); -const Allocator = std.mem.Allocator; const RingBuffer = std.RingBuffer; const types = @import("zstandard/types.zig"); @@ -10,7 +9,9 @@ pub const decompress = @import("zstandard/decompress.zig"); pub const DecompressStreamOptions = struct { verify_checksum: bool = true, - window_size_max: usize = 1 << 23, // 8MiB default maximum window size + window_size_max: usize = default_window_size_max, + + pub const default_window_size_max = 1 << 23; // 8MiB default maximum window size }; pub fn DecompressStream( @@ -20,20 +21,29 @@ pub fn DecompressStream( return struct { const Self = @This(); - allocator: Allocator, + pub const window_size_max = options.window_size_max; + + const table_size_max = types.compressed_block.table_size_max; + source: std.io.CountingReader(ReaderType), state: enum { NewFrame, InFrame, LastBlock }, decode_state: decompress.block.DecodeState, frame_context: decompress.FrameContext, - buffer: RingBuffer, - literal_fse_buffer: []types.compressed_block.Table.Fse, - match_fse_buffer: []types.compressed_block.Table.Fse, - offset_fse_buffer: []types.compressed_block.Table.Fse, - literals_buffer: []u8, - sequence_buffer: []u8, + buffer: WindowBuffer, + literal_fse_buffer: [table_size_max.literal]types.compressed_block.Table.Fse, + match_fse_buffer: [table_size_max.match]types.compressed_block.Table.Fse, + offset_fse_buffer: [table_size_max.offset]types.compressed_block.Table.Fse, + literals_buffer: [types.block_size_max]u8, + sequence_buffer: [types.block_size_max]u8, checksum: if (options.verify_checksum) ?u32 else void, current_frame_decompressed_size: usize, + const WindowBuffer = struct { + data: *[options.window_size_max]u8 = undefined, + read_index: usize = 0, + write_index: usize = 0, + }; + pub const Error = ReaderType.Error || error{ ChecksumFailure, DictionaryIdFlagUnsupported, @@ -44,14 +54,13 @@ pub fn DecompressStream( pub const Reader = std.io.Reader(*Self, Error, read); - pub fn init(allocator: Allocator, source: ReaderType) Self { + pub fn init(source: ReaderType, window_buffer: *[options.window_size_max]u8) Self { return Self{ - .allocator = allocator, .source = std.io.countingReader(source), .state = .NewFrame, .decode_state = undefined, .frame_context = undefined, - .buffer = undefined, + .buffer = .{ .data = window_buffer }, .literal_fse_buffer = undefined, .match_fse_buffer = undefined, .offset_fse_buffer = undefined, @@ -76,44 +85,11 @@ pub fn DecompressStream( options.verify_checksum, ); - const literal_fse_buffer = try self.allocator.alloc( - types.compressed_block.Table.Fse, - types.compressed_block.table_size_max.literal, - ); - errdefer self.allocator.free(literal_fse_buffer); - - const match_fse_buffer = try self.allocator.alloc( - types.compressed_block.Table.Fse, - types.compressed_block.table_size_max.match, - ); - errdefer self.allocator.free(match_fse_buffer); - - const offset_fse_buffer = try self.allocator.alloc( - types.compressed_block.Table.Fse, - types.compressed_block.table_size_max.offset, - ); - errdefer self.allocator.free(offset_fse_buffer); - const decode_state = decompress.block.DecodeState.init( - literal_fse_buffer, - match_fse_buffer, - offset_fse_buffer, + &self.literal_fse_buffer, + &self.match_fse_buffer, + &self.offset_fse_buffer, ); - const buffer = try RingBuffer.init(self.allocator, frame_context.window_size); - - const literals_data = try self.allocator.alloc(u8, frame_context.block_size_max); - errdefer self.allocator.free(literals_data); - - const sequence_data = try self.allocator.alloc(u8, frame_context.block_size_max); - errdefer self.allocator.free(sequence_data); - - self.literal_fse_buffer = literal_fse_buffer; - self.match_fse_buffer = match_fse_buffer; - self.offset_fse_buffer = offset_fse_buffer; - self.literals_buffer = literals_data; - self.sequence_buffer = sequence_data; - - self.buffer = buffer; self.decode_state = decode_state; self.frame_context = frame_context; @@ -126,16 +102,6 @@ pub fn DecompressStream( } } - pub fn deinit(self: *Self) void { - if (self.state == .NewFrame) return; - self.allocator.free(self.decode_state.literal_fse_buffer); - self.allocator.free(self.decode_state.match_fse_buffer); - self.allocator.free(self.decode_state.offset_fse_buffer); - self.allocator.free(self.literals_buffer); - self.allocator.free(self.sequence_buffer); - self.buffer.deinit(self.allocator); - } - pub fn reader(self: *Self) Reader { return .{ .context = self }; } @@ -153,7 +119,6 @@ pub fn DecompressStream( 0 else error.MalformedFrame, - error.OutOfMemory => return error.OutOfMemory, else => return error.MalformedFrame, }; } @@ -165,20 +130,30 @@ pub fn DecompressStream( fn readInner(self: *Self, buffer: []u8) Error!usize { std.debug.assert(self.state != .NewFrame); + var ring_buffer = RingBuffer{ + .data = self.buffer.data, + .read_index = self.buffer.read_index, + .write_index = self.buffer.write_index, + }; + defer { + self.buffer.read_index = ring_buffer.read_index; + self.buffer.write_index = ring_buffer.write_index; + } + const source_reader = self.source.reader(); - while (self.buffer.isEmpty() and self.state != .LastBlock) { + while (ring_buffer.isEmpty() and self.state != .LastBlock) { const header_bytes = source_reader.readBytesNoEof(3) catch return error.MalformedFrame; const block_header = decompress.block.decodeBlockHeader(&header_bytes); decompress.block.decodeBlockReader( - &self.buffer, + &ring_buffer, source_reader, block_header, &self.decode_state, self.frame_context.block_size_max, - self.literals_buffer, - self.sequence_buffer, + &self.literals_buffer, + &self.sequence_buffer, ) catch return error.MalformedBlock; @@ -186,12 +161,12 @@ pub fn DecompressStream( if (self.current_frame_decompressed_size > size) return error.MalformedFrame; } - const size = self.buffer.len(); + const size = ring_buffer.len(); self.current_frame_decompressed_size += size; if (self.frame_context.hasher_opt) |*hasher| { if (size > 0) { - const written_slice = self.buffer.sliceLast(size); + const written_slice = ring_buffer.sliceLast(size); hasher.update(written_slice.first); hasher.update(written_slice.second); } @@ -216,18 +191,12 @@ pub fn DecompressStream( } } - const size = @min(self.buffer.len(), buffer.len); + const size = @min(ring_buffer.len(), buffer.len); if (size > 0) { - self.buffer.readFirstAssumeLength(buffer, size); + ring_buffer.readFirstAssumeLength(buffer, size); } - if (self.state == .LastBlock and self.buffer.len() == 0) { + if (self.state == .LastBlock and ring_buffer.len() == 0) { self.state = .NewFrame; - self.allocator.free(self.literal_fse_buffer); - self.allocator.free(self.match_fse_buffer); - self.allocator.free(self.offset_fse_buffer); - self.allocator.free(self.literals_buffer); - self.allocator.free(self.sequence_buffer); - self.buffer.deinit(self.allocator); } return size; } @@ -235,24 +204,24 @@ pub fn DecompressStream( } pub fn decompressStreamOptions( - allocator: Allocator, reader: anytype, comptime options: DecompressStreamOptions, + window_buffer: *[options.window_size_max]u8, ) DecompressStream(@TypeOf(reader), options) { - return DecompressStream(@TypeOf(reader), options).init(allocator, reader); + return DecompressStream(@TypeOf(reader), options).init(reader, window_buffer); } pub fn decompressStream( - allocator: Allocator, reader: anytype, + window_buffer: *[DecompressStreamOptions.default_window_size_max]u8, ) DecompressStream(@TypeOf(reader), .{}) { - return DecompressStream(@TypeOf(reader), .{}).init(allocator, reader); + return DecompressStream(@TypeOf(reader), .{}).init(reader, window_buffer); } fn testDecompress(data: []const u8) ![]u8 { + var window_buffer: [DecompressStreamOptions.default_window_size_max]u8 = undefined; var in_stream = std.io.fixedBufferStream(data); - var zstd_stream = decompressStream(std.testing.allocator, in_stream.reader()); - defer zstd_stream.deinit(); + var zstd_stream = decompressStream(in_stream.reader(), &window_buffer); const result = zstd_stream.reader().readAllAlloc(std.testing.allocator, std.math.maxInt(usize)); return result; } @@ -301,9 +270,9 @@ fn expectEqualDecoded(expected: []const u8, input: []const u8) !void { } { + var window_buffer: [DecompressStreamOptions.default_window_size_max]u8 = undefined; var in_stream = std.io.fixedBufferStream(input); - var stream = decompressStream(allocator, in_stream.reader()); - defer stream.deinit(); + var stream = decompressStream(in_stream.reader(), &window_buffer); const result = try stream.reader().readAllAlloc(allocator, std.math.maxInt(usize)); defer allocator.free(result); diff --git a/lib/std/compress/zstandard/decompress.zig b/lib/std/compress/zstandard/decompress.zig index a012312ab1..86be16268f 100644 --- a/lib/std/compress/zstandard/decompress.zig +++ b/lib/std/compress/zstandard/decompress.zig @@ -409,7 +409,7 @@ pub const FrameContext = struct { .hasher_opt = if (should_compute_checksum) std.hash.XxHash64.init(0) else null, .window_size = window_size, .has_checksum = frame_header.descriptor.content_checksum_flag, - .block_size_max = @min(1 << 17, window_size), + .block_size_max = @min(types.block_size_max, window_size), .content_size = content_size, }; } diff --git a/lib/std/compress/zstandard/types.zig b/lib/std/compress/zstandard/types.zig index 98cec9f866..41c3797d16 100644 --- a/lib/std/compress/zstandard/types.zig +++ b/lib/std/compress/zstandard/types.zig @@ -1,3 +1,5 @@ +pub const block_size_max = 1 << 17; + pub const frame = struct { pub const Kind = enum { zstandard, skippable }; diff --git a/src/Package/Fetch.zig b/src/Package/Fetch.zig index 2405832aef..1cb98e560b 100644 --- a/src/Package/Fetch.zig +++ b/src/Package/Fetch.zig @@ -1126,11 +1126,14 @@ fn unpackResource( // wrapped for generic use in unpackTarballCompressed: see github.com/ziglang/zig/issues/14739 const ZstdWrapper = struct { fn DecompressType(comptime T: type) type { - return error{}!std.compress.zstd.DecompressStream(T, .{}); + return Allocator.Error!std.compress.zstd.DecompressStream(T, .{}); } fn decompress(allocator: Allocator, reader: anytype) DecompressType(@TypeOf(reader)) { - return std.compress.zstd.decompressStream(allocator, reader); + const window_size = std.compress.zstd.DecompressStreamOptions.default_window_size_max; + const window_buffer = try allocator.create([window_size]u8); + defer allocator.destroy(window_buffer); + return std.compress.zstd.decompressStream(reader, window_buffer); } }; @@ -1138,7 +1141,7 @@ fn unpackTarballCompressed( f: *Fetch, out_dir: fs.Dir, resource: *Resource, - comptime Compression: type, + Compression: anytype, ) RunError!void { const gpa = f.arena.child_allocator; const eb = &f.error_bundle; @@ -1151,7 +1154,7 @@ fn unpackTarballCompressed( .{@errorName(err)}, )); }; - defer decompress.deinit(); + defer if (@hasDecl(Compression, "deinit")) decompress.deinit(); return unpackTarball(f, out_dir, decompress.reader()); }