Merge pull request #24559 from ziglang/zstd

std: rework zstd for new I/O API

Author: andrewrk

lib/std/Io/DeprecatedReader.zig

@@ -373,11 +373,11 @@ pub fn discard(self: Self) anyerror!u64 {
 }
 
 /// Helper for bridging to the new `Reader` API while upgrading.
-pub fn adaptToNewApi(self: *const Self) Adapter {
+pub fn adaptToNewApi(self: *const Self, buffer: []u8) Adapter {
     return .{
         .derp_reader = self.*,
         .new_interface = .{
-            .buffer = &.{},
+            .buffer = buffer,
             .vtable = &.{ .stream = Adapter.stream },
             .seek = 0,
             .end = 0,

lib/std/Io/Reader.zig

@@ -185,6 +185,32 @@ pub fn streamExact64(r: *Reader, w: *Writer, n: u64) StreamError!void {
     while (remaining != 0) remaining -= try r.stream(w, .limited64(remaining));
 }
 
+/// "Pump" exactly `n` bytes from the reader to the writer.
+///
+/// When draining `w`, ensures that at least `preserve_len` bytes remain
+/// buffered.
+///
+/// Asserts `Writer.buffer` capacity exceeds `preserve_len`.
+pub fn streamExactPreserve(r: *Reader, w: *Writer, preserve_len: usize, n: usize) StreamError!void {
+    if (w.end + n <= w.buffer.len) {
+        @branchHint(.likely);
+        return streamExact(r, w, n);
+    }
+    // If `n` is large, we can ignore `preserve_len` up to a point.
+    var remaining = n;
+    while (remaining > preserve_len) {
+        assert(remaining != 0);
+        remaining -= try r.stream(w, .limited(remaining - preserve_len));
+        if (w.end + remaining <= w.buffer.len) return streamExact(r, w, remaining);
+    }
+    // All the next bytes received must be preserved.
+    if (preserve_len < w.end) {
+        @memmove(w.buffer[0..preserve_len], w.buffer[w.end - preserve_len ..][0..preserve_len]);
+        w.end = preserve_len;
+    }
+    return streamExact(r, w, remaining);
+}
+
 /// "Pump" data from the reader to the writer, handling `error.EndOfStream` as
 /// a success case.
 ///
@@ -240,7 +266,7 @@ pub fn allocRemaining(r: *Reader, gpa: Allocator, limit: Limit) LimitedAllocErro
 /// such case, the next byte that would be read will be the first one to exceed
 /// `limit`, and all preceeding bytes have been appended to `list`.
 ///
-/// Asserts `buffer` has nonzero capacity.
+/// If `limit` is not `Limit.unlimited`, asserts `buffer` has nonzero capacity.
 ///
 /// See also:
 /// * `allocRemaining`
@@ -251,7 +277,7 @@ pub fn appendRemaining(
     list: *std.ArrayListAlignedUnmanaged(u8, alignment),
     limit: Limit,
 ) LimitedAllocError!void {
-    assert(r.buffer.len != 0); // Needed to detect limit exceeded without losing data.
+    if (limit != .unlimited) assert(r.buffer.len != 0); // Needed to detect limit exceeded without losing data.
     const buffer_contents = r.buffer[r.seek..r.end];
     const copy_len = limit.minInt(buffer_contents.len);
     try list.appendSlice(gpa, r.buffer[0..copy_len]);

lib/std/Io/Writer.zig

@@ -256,10 +256,10 @@ test "fixed buffer flush" {
     try testing.expectEqual(10, buffer[0]);
 }
 
-/// Calls `VTable.drain` but hides the last `preserve_length` bytes from the
+/// Calls `VTable.drain` but hides the last `preserve_len` bytes from the
 /// implementation, keeping them buffered.
-pub fn drainPreserve(w: *Writer, preserve_length: usize) Error!void {
-    const temp_end = w.end -| preserve_length;
+pub fn drainPreserve(w: *Writer, preserve_len: usize) Error!void {
+    const temp_end = w.end -| preserve_len;
     const preserved = w.buffer[temp_end..w.end];
     w.end = temp_end;
     defer w.end += preserved.len;
@@ -310,24 +310,38 @@ pub fn writableSliceGreedy(w: *Writer, minimum_length: usize) Error![]u8 {
 }
 
 /// Asserts the provided buffer has total capacity enough for `minimum_length`
-/// and `preserve_length` combined.
+/// and `preserve_len` combined.
 ///
 /// Does not `advance` the buffer end position.
 ///
-/// When draining the buffer, ensures that at least `preserve_length` bytes
+/// When draining the buffer, ensures that at least `preserve_len` bytes
 /// remain buffered.
 ///
-/// If `preserve_length` is zero, this is equivalent to `writableSliceGreedy`.
-pub fn writableSliceGreedyPreserve(w: *Writer, preserve_length: usize, minimum_length: usize) Error![]u8 {
-    assert(w.buffer.len >= preserve_length + minimum_length);
+/// If `preserve_len` is zero, this is equivalent to `writableSliceGreedy`.
+pub fn writableSliceGreedyPreserve(w: *Writer, preserve_len: usize, minimum_length: usize) Error![]u8 {
+    assert(w.buffer.len >= preserve_len + minimum_length);
     while (w.buffer.len - w.end < minimum_length) {
-        try drainPreserve(w, preserve_length);
+        try drainPreserve(w, preserve_len);
     } else {
         @branchHint(.likely);
         return w.buffer[w.end..];
     }
 }
 
+/// Asserts the provided buffer has total capacity enough for `len`.
+///
+/// Advances the buffer end position by `len`.
+///
+/// When draining the buffer, ensures that at least `preserve_len` bytes
+/// remain buffered.
+///
+/// If `preserve_len` is zero, this is equivalent to `writableSlice`.
+pub fn writableSlicePreserve(w: *Writer, preserve_len: usize, len: usize) Error![]u8 {
+    const big_slice = try w.writableSliceGreedyPreserve(preserve_len, len);
+    advance(w, len);
+    return big_slice[0..len];
+}
+
 pub const WritableVectorIterator = struct {
     first: []u8,
     middle: []const []u8 = &.{},
@@ -523,16 +537,16 @@ pub fn write(w: *Writer, bytes: []const u8) Error!usize {
     return w.vtable.drain(w, &.{bytes}, 1);
 }
 
-/// Asserts `buffer` capacity exceeds `preserve_length`.
-pub fn writePreserve(w: *Writer, preserve_length: usize, bytes: []const u8) Error!usize {
-    assert(preserve_length <= w.buffer.len);
+/// Asserts `buffer` capacity exceeds `preserve_len`.
+pub fn writePreserve(w: *Writer, preserve_len: usize, bytes: []const u8) Error!usize {
+    assert(preserve_len <= w.buffer.len);
     if (w.end + bytes.len <= w.buffer.len) {
         @branchHint(.likely);
         @memcpy(w.buffer[w.end..][0..bytes.len], bytes);
         w.end += bytes.len;
         return bytes.len;
     }
-    const temp_end = w.end -| preserve_length;
+    const temp_end = w.end -| preserve_len;
     const preserved = w.buffer[temp_end..w.end];
     w.end = temp_end;
     defer w.end += preserved.len;
@@ -552,13 +566,13 @@ pub fn writeAll(w: *Writer, bytes: []const u8) Error!void {
 /// Calls `drain` as many times as necessary such that all of `bytes` are
 /// transferred.
 ///
-/// When draining the buffer, ensures that at least `preserve_length` bytes
+/// When draining the buffer, ensures that at least `preserve_len` bytes
 /// remain buffered.
 ///
-/// Asserts `buffer` capacity exceeds `preserve_length`.
-pub fn writeAllPreserve(w: *Writer, preserve_length: usize, bytes: []const u8) Error!void {
+/// Asserts `buffer` capacity exceeds `preserve_len`.
+pub fn writeAllPreserve(w: *Writer, preserve_len: usize, bytes: []const u8) Error!void {
     var index: usize = 0;
-    while (index < bytes.len) index += try w.writePreserve(preserve_length, bytes[index..]);
+    while (index < bytes.len) index += try w.writePreserve(preserve_len, bytes[index..]);
 }
 
 /// Renders fmt string with args, calling `writer` with slices of bytes.
@@ -761,11 +775,11 @@ pub fn writeByte(w: *Writer, byte: u8) Error!void {
     }
 }
 
-/// When draining the buffer, ensures that at least `preserve_length` bytes
+/// When draining the buffer, ensures that at least `preserve_len` bytes
 /// remain buffered.
-pub fn writeBytePreserve(w: *Writer, preserve_length: usize, byte: u8) Error!void {
+pub fn writeBytePreserve(w: *Writer, preserve_len: usize, byte: u8) Error!void {
     while (w.buffer.len - w.end == 0) {
-        try drainPreserve(w, preserve_length);
+        try drainPreserve(w, preserve_len);
     } else {
         @branchHint(.likely);
         w.buffer[w.end] = byte;
@@ -788,10 +802,42 @@ test splatByteAll {
     try testing.expectEqualStrings("7" ** 45, aw.writer.buffered());
 }
 
+pub fn splatBytePreserve(w: *Writer, preserve_len: usize, byte: u8, n: usize) Error!void {
+    const new_end = w.end + n;
+    if (new_end <= w.buffer.len) {
+        @memset(w.buffer[w.end..][0..n], byte);
+        w.end = new_end;
+        return;
+    }
+    // If `n` is large, we can ignore `preserve_len` up to a point.
+    var remaining = n;
+    while (remaining > preserve_len) {
+        assert(remaining != 0);
+        remaining -= try splatByte(w, byte, remaining - preserve_len);
+        if (w.end + remaining <= w.buffer.len) {
+            @memset(w.buffer[w.end..][0..remaining], byte);
+            w.end += remaining;
+            return;
+        }
+    }
+    // All the next bytes received must be preserved.
+    if (preserve_len < w.end) {
+        @memmove(w.buffer[0..preserve_len], w.buffer[w.end - preserve_len ..][0..preserve_len]);
+        w.end = preserve_len;
+    }
+    while (remaining > 0) remaining -= try w.splatByte(byte, remaining);
+}
+
 /// Writes the same byte many times, allowing short writes.
 ///
 /// Does maximum of one underlying `VTable.drain`.
 pub fn splatByte(w: *Writer, byte: u8, n: usize) Error!usize {
+    if (w.end + n <= w.buffer.len) {
+        @branchHint(.likely);
+        @memset(w.buffer[w.end..][0..n], byte);
+        w.end += n;
+        return n;
+    }
     return writeSplat(w, &.{&.{byte}}, n);
 }
 
@@ -801,9 +847,10 @@ pub fn splatBytesAll(w: *Writer, bytes: []const u8, splat: usize) Error!void {
     var remaining_bytes: usize = bytes.len * splat;
     remaining_bytes -= try w.splatBytes(bytes, splat);
     while (remaining_bytes > 0) {
-        const leftover = remaining_bytes % bytes.len;
-        const buffers: [2][]const u8 = .{ bytes[bytes.len - leftover ..], bytes };
-        remaining_bytes -= try w.writeSplat(&buffers, splat);
+        const leftover_splat = remaining_bytes / bytes.len;
+        const leftover_bytes = remaining_bytes % bytes.len;
+        const buffers: [2][]const u8 = .{ bytes[bytes.len - leftover_bytes ..], bytes };
+        remaining_bytes -= try w.writeSplat(&buffers, leftover_splat);
     }
 }
 

lib/std/compress.zig

@@ -1,75 +1,19 @@
 //! Compression algorithms.
 
-const std = @import("std.zig");
-
 pub const flate = @import("compress/flate.zig");
 pub const gzip = @import("compress/gzip.zig");
 pub const zlib = @import("compress/zlib.zig");
 pub const lzma = @import("compress/lzma.zig");
 pub const lzma2 = @import("compress/lzma2.zig");
 pub const xz = @import("compress/xz.zig");
-pub const zstd = @import("compress/zstandard.zig");
-
-pub fn HashedReader(ReaderType: type, HasherType: type) type {
-    return struct {
-        child_reader: ReaderType,
-        hasher: HasherType,
-
-        pub const Error = ReaderType.Error;
-        pub const Reader = std.io.GenericReader(*@This(), Error, read);
-
-        pub fn read(self: *@This(), buf: []u8) Error!usize {
-            const amt = try self.child_reader.read(buf);
-            self.hasher.update(buf[0..amt]);
-            return amt;
-        }
-
-        pub fn reader(self: *@This()) Reader {
-            return .{ .context = self };
-        }
-    };
-}
-
-pub fn hashedReader(
-    reader: anytype,
-    hasher: anytype,
-) HashedReader(@TypeOf(reader), @TypeOf(hasher)) {
-    return .{ .child_reader = reader, .hasher = hasher };
-}
-
-pub fn HashedWriter(WriterType: type, HasherType: type) type {
-    return struct {
-        child_writer: WriterType,
-        hasher: HasherType,
-
-        pub const Error = WriterType.Error;
-        pub const Writer = std.io.GenericWriter(*@This(), Error, write);
-
-        pub fn write(self: *@This(), buf: []const u8) Error!usize {
-            const amt = try self.child_writer.write(buf);
-            self.hasher.update(buf[0..amt]);
-            return amt;
-        }
-
-        pub fn writer(self: *@This()) Writer {
-            return .{ .context = self };
-        }
-    };
-}
-
-pub fn hashedWriter(
-    writer: anytype,
-    hasher: anytype,
-) HashedWriter(@TypeOf(writer), @TypeOf(hasher)) {
-    return .{ .child_writer = writer, .hasher = hasher };
-}
+pub const zstd = @import("compress/zstd.zig");
 
 test {
+    _ = flate;
     _ = lzma;
     _ = lzma2;
     _ = xz;
     _ = zstd;
-    _ = flate;
     _ = gzip;
     _ = zlib;
 }

lib/std/compress/xz.zig

@@ -47,7 +47,7 @@ pub fn Decompress(comptime ReaderType: type) type {
 
             var check: Check = undefined;
             const hash_a = blk: {
-                var hasher = std.compress.hashedReader(source, Crc32.init());
+                var hasher = hashedReader(source, Crc32.init());
                 try readStreamFlags(hasher.reader(), &check);
                 break :blk hasher.hasher.final();
             };
@@ -80,7 +80,7 @@ pub fn Decompress(comptime ReaderType: type) type {
                 return r;
 
             const index_size = blk: {
-                var hasher = std.compress.hashedReader(self.in_reader, Crc32.init());
+                var hasher = hashedReader(self.in_reader, Crc32.init());
                 hasher.hasher.update(&[1]u8{0x00});
 
                 var counter = std.io.countingReader(hasher.reader());
@@ -115,7 +115,7 @@ pub fn Decompress(comptime ReaderType: type) type {
             const hash_a = try self.in_reader.readInt(u32, .little);
 
             const hash_b = blk: {
-                var hasher = std.compress.hashedReader(self.in_reader, Crc32.init());
+                var hasher = hashedReader(self.in_reader, Crc32.init());
                 const hashed_reader = hasher.reader();
 
                 const backward_size = (@as(u64, try hashed_reader.readInt(u32, .little)) + 1) * 4;
@@ -140,6 +140,33 @@ pub fn Decompress(comptime ReaderType: type) type {
     };
 }
 
+pub fn HashedReader(ReaderType: type, HasherType: type) type {
+    return struct {
+        child_reader: ReaderType,
+        hasher: HasherType,
+
+        pub const Error = ReaderType.Error;
+        pub const Reader = std.io.GenericReader(*@This(), Error, read);
+
+        pub fn read(self: *@This(), buf: []u8) Error!usize {
+            const amt = try self.child_reader.read(buf);
+            self.hasher.update(buf[0..amt]);
+            return amt;
+        }
+
+        pub fn reader(self: *@This()) Reader {
+            return .{ .context = self };
+        }
+    };
+}
+
+pub fn hashedReader(
+    reader: anytype,
+    hasher: anytype,
+) HashedReader(@TypeOf(reader), @TypeOf(hasher)) {
+    return .{ .child_reader = reader, .hasher = hasher };
+}
+
 test {
     _ = @import("xz/test.zig");
 }

lib/std/compress/xz/block.zig

@@ -91,7 +91,7 @@ pub fn Decoder(comptime ReaderType: type) type {
 
             // Block Header
             {
-                var header_hasher = std.compress.hashedReader(block_reader, Crc32.init());
+                var header_hasher = xz.hashedReader(block_reader, Crc32.init());
                 const header_reader = header_hasher.reader();
 
                 const header_size = @as(u64, try header_reader.readByte()) * 4;