authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-08-13 22:02:30-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-08-14 12:56:37-07:00
logaf7e142485d422e1fdbff1dcfc6e95e4f9126453
treea59c0e59dfe5e56ba1195a08a5aefefa53847c73
parent96e4825fbba714322ad750ee179da3d5f8463bb5

std.Io.Writer: introduce rebase to the vtable

fixes #24814

2 files changed, 103 insertions(+), 61 deletions(-)

lib/std/Io/Writer.zig+102-58
......@@ -4,7 +4,7 @@ const native_endian = builtin.target.cpu.arch.endian();
44const Writer = @This();
55const std = @import("../std.zig");
66const assert = std.debug.assert;
7const Limit = std.io.Limit;
7const Limit = std.Io.Limit;
88const File = std.fs.File;
99const testing = std.testing;
1010const Allocator = std.mem.Allocator;
......@@ -76,6 +76,14 @@ pub const VTable = struct {
7676 /// There may be subsequent calls to `drain` and `sendFile` after a `flush`
7777 /// operation.
7878 flush: *const fn (w: *Writer) Error!void = defaultFlush,
79
80 /// Ensures `capacity` more bytes can be buffered without rebasing.
81 ///
82 /// The most recent `preserve` bytes must remain buffered.
83 ///
84 /// Only called when `capacity` bytes cannot fit into the unused capacity
85 /// of `buffer`.
86 rebase: *const fn (w: *Writer, preserve: usize, capacity: usize) Error!void = defaultRebase,
7987};
8088
8189pub const Error = error{
......@@ -117,6 +125,7 @@ pub fn fixed(buffer: []u8) Writer {
117125 .vtable = &.{
118126 .drain = fixedDrain,
119127 .flush = noopFlush,
128 .rebase = failingRebase,
120129 },
121130 .buffer = buffer,
122131 };
......@@ -130,6 +139,7 @@ pub const failing: Writer = .{
130139 .vtable = &.{
131140 .drain = failingDrain,
132141 .sendFile = failingSendFile,
142 .rebase = failingRebase,
133143 },
134144};
135145
......@@ -276,7 +286,7 @@ fn writeSplatHeaderLimitFinish(
276286
277287test "writeSplatHeader splatting avoids buffer aliasing temptation" {
278288 const initial_buf = try testing.allocator.alloc(u8, 8);
279 var aw: std.io.Writer.Allocating = .initOwnedSlice(testing.allocator, initial_buf);
289 var aw: Allocating = .initOwnedSlice(testing.allocator, initial_buf);
280290 defer aw.deinit();
281291 // This test assumes 8 vector buffer in this function.
282292 const n = try aw.writer.writeSplatHeader("header which is longer than buf ", &.{
......@@ -307,24 +317,41 @@ pub fn noopFlush(w: *Writer) Error!void {
307317
308318test "fixed buffer flush" {
309319 var buffer: [1]u8 = undefined;
310 var writer: std.io.Writer = .fixed(&buffer);
320 var writer: Writer = .fixed(&buffer);
311321
312322 try writer.writeByte(10);
313323 try writer.flush();
314324 try testing.expectEqual(10, buffer[0]);
315325}
316326
317/// Calls `VTable.drain` but hides the last `preserve_len` bytes from the
318/// implementation, keeping them buffered.
319pub fn drainPreserve(w: *Writer, preserve_len: usize) Error!void {
320 const preserved_head = w.end -| preserve_len;
321 const preserved_tail = w.end;
322 const preserved_len = preserved_tail - preserved_head;
323 w.end = preserved_head;
324 defer w.end += preserved_len;
325 assert(0 == try w.vtable.drain(w, &.{""}, 1));
326 assert(w.end <= preserved_head + preserved_len);
327 @memmove(w.buffer[w.end..][0..preserved_len], w.buffer[preserved_head..preserved_tail]);
327pub fn rebase(w: *Writer, preserve: usize, unused_capacity_len: usize) Error!void {
328 if (w.buffer.len - w.end >= unused_capacity_len) {
329 @branchHint(.likely);
330 return;
331 }
332 try w.vtable.rebase(w, preserve, unused_capacity_len);
333}
334
335pub fn defaultRebase(w: *Writer, preserve: usize, minimum_len: usize) Error!void {
336 while (w.buffer.len - w.end < minimum_len) {
337 {
338 // TODO: instead of this logic that "hides" data from
339 // the implementation, introduce a seek index to Writer
340 const preserved_head = w.end -| preserve;
341 const preserved_tail = w.end;
342 const preserved_len = preserved_tail - preserved_head;
343 w.end = preserved_head;
344 defer w.end += preserved_len;
345 assert(0 == try w.vtable.drain(w, &.{""}, 1));
346 assert(w.end <= preserved_head + preserved_len);
347 @memmove(w.buffer[w.end..][0..preserved_len], w.buffer[preserved_head..preserved_tail]);
348 }
349
350 // If the loop condition was false this assertion would have passed
351 // anyway. Otherwise, give the implementation a chance to grow the
352 // buffer before asserting on the buffer length.
353 assert(w.buffer.len - preserve >= minimum_len);
354 }
328355}
329356
330357pub fn unusedCapacitySlice(w: *const Writer) []u8 {
......@@ -353,53 +380,44 @@ pub fn writableSlice(w: *Writer, len: usize) Error![]u8 {
353380 return big_slice[0..len];
354381}
355382
356/// Asserts the provided buffer has total capacity enough for `minimum_length`.
383/// Asserts the provided buffer has total capacity enough for `minimum_len`.
357384///
358385/// Does not `advance` the buffer end position.
359386///
360/// If `minimum_length` is zero, this is equivalent to `unusedCapacitySlice`.
361pub fn writableSliceGreedy(w: *Writer, minimum_length: usize) Error![]u8 {
362 while (w.buffer.len - w.end < minimum_length) {
363 assert(0 == try w.vtable.drain(w, &.{""}, 1));
364 // If the loop condition was false this assertion would have passed
365 // anyway. Otherwise, give the implementation a chance to grow the
366 // buffer before asserting on the buffer length.
367 assert(w.buffer.len >= minimum_length);
368 } else {
369 @branchHint(.likely);
370 return w.buffer[w.end..];
371 }
387/// If `minimum_len` is zero, this is equivalent to `unusedCapacitySlice`.
388pub fn writableSliceGreedy(w: *Writer, minimum_len: usize) Error![]u8 {
389 return writableSliceGreedyPreserve(w, 0, minimum_len);
372390}
373391
374/// Asserts the provided buffer has total capacity enough for `minimum_length`
375/// and `preserve_len` combined.
392/// Asserts the provided buffer has total capacity enough for `minimum_len`
393/// and `preserve` combined.
376394///
377395/// Does not `advance` the buffer end position.
378396///
379/// When draining the buffer, ensures that at least `preserve_len` bytes
397/// When draining the buffer, ensures that at least `preserve` bytes
380398/// remain buffered.
381399///
382/// If `preserve_len` is zero, this is equivalent to `writableSliceGreedy`.
383pub fn writableSliceGreedyPreserve(w: *Writer, preserve_len: usize, minimum_length: usize) Error![]u8 {
384 assert(w.buffer.len >= preserve_len + minimum_length);
385 while (w.buffer.len - w.end < minimum_length) {
386 try drainPreserve(w, preserve_len);
387 } else {
400/// If `preserve` is zero, this is equivalent to `writableSliceGreedy`.
401pub fn writableSliceGreedyPreserve(w: *Writer, preserve: usize, minimum_len: usize) Error![]u8 {
402 if (w.buffer.len - w.end >= minimum_len) {
388403 @branchHint(.likely);
389404 return w.buffer[w.end..];
390405 }
406 try rebase(w, preserve, minimum_len);
407 assert(w.buffer.len >= preserve + minimum_len);
408 return w.buffer[w.end..];
391409}
392410
393411/// Asserts the provided buffer has total capacity enough for `len`.
394412///
395413/// Advances the buffer end position by `len`.
396414///
397/// When draining the buffer, ensures that at least `preserve_len` bytes
415/// When draining the buffer, ensures that at least `preserve` bytes
398416/// remain buffered.
399417///
400/// If `preserve_len` is zero, this is equivalent to `writableSlice`.
401pub fn writableSlicePreserve(w: *Writer, preserve_len: usize, len: usize) Error![]u8 {
402 const big_slice = try w.writableSliceGreedyPreserve(preserve_len, len);
418/// If `preserve` is zero, this is equivalent to `writableSlice`.
419pub fn writableSlicePreserve(w: *Writer, preserve: usize, len: usize) Error![]u8 {
420 const big_slice = try w.writableSliceGreedyPreserve(preserve, len);
403421 advance(w, len);
404422 return big_slice[0..len];
405423}
......@@ -708,16 +726,18 @@ pub fn writeByte(w: *Writer, byte: u8) Error!void {
708726 }
709727}
710728
711/// When draining the buffer, ensures that at least `preserve_len` bytes
729/// When draining the buffer, ensures that at least `preserve` bytes
712730/// remain buffered.
713pub fn writeBytePreserve(w: *Writer, preserve_len: usize, byte: u8) Error!void {
714 while (w.buffer.len - w.end == 0) {
715 try drainPreserve(w, preserve_len);
716 } else {
731pub fn writeBytePreserve(w: *Writer, preserve: usize, byte: u8) Error!void {
732 if (w.buffer.len - w.end != 0) {
717733 @branchHint(.likely);
718734 w.buffer[w.end] = byte;
719735 w.end += 1;
736 return;
720737 }
738 try w.vtable.rebase(w, preserve, 1);
739 w.buffer[w.end] = byte;
740 w.end += 1;
721741}
722742
723743/// Writes the same byte many times, performing the underlying write call as
......@@ -735,18 +755,18 @@ test splatByteAll {
735755 try testing.expectEqualStrings("7" ** 45, aw.writer.buffered());
736756}
737757
738pub fn splatBytePreserve(w: *Writer, preserve_len: usize, byte: u8, n: usize) Error!void {
758pub fn splatBytePreserve(w: *Writer, preserve: usize, byte: u8, n: usize) Error!void {
739759 const new_end = w.end + n;
740760 if (new_end <= w.buffer.len) {
741761 @memset(w.buffer[w.end..][0..n], byte);
742762 w.end = new_end;
743763 return;
744764 }
745 // If `n` is large, we can ignore `preserve_len` up to a point.
765 // If `n` is large, we can ignore `preserve` up to a point.
746766 var remaining = n;
747 while (remaining > preserve_len) {
767 while (remaining > preserve) {
748768 assert(remaining != 0);
749 remaining -= try splatByte(w, byte, remaining - preserve_len);
769 remaining -= try splatByte(w, byte, remaining - preserve);
750770 if (w.end + remaining <= w.buffer.len) {
751771 @memset(w.buffer[w.end..][0..remaining], byte);
752772 w.end += remaining;
......@@ -754,9 +774,9 @@ pub fn splatBytePreserve(w: *Writer, preserve_len: usize, byte: u8, n: usize) Er
754774 }
755775 }
756776 // All the next bytes received must be preserved.
757 if (preserve_len < w.end) {
758 @memmove(w.buffer[0..preserve_len], w.buffer[w.end - preserve_len ..][0..preserve_len]);
759 w.end = preserve_len;
777 if (preserve < w.end) {
778 @memmove(w.buffer[0..preserve], w.buffer[w.end - preserve ..][0..preserve]);
779 w.end = preserve;
760780 }
761781 while (remaining > 0) remaining -= try w.splatByte(byte, remaining);
762782}
......@@ -1667,7 +1687,7 @@ pub const ByteSizeUnits = enum {
16671687
16681688/// Format option `precision` is ignored when `value` is less than 1kB
16691689pub fn printByteSize(
1670 w: *std.io.Writer,
1690 w: *Writer,
16711691 value: u64,
16721692 comptime units: ByteSizeUnits,
16731693 options: std.fmt.Options,
......@@ -2169,7 +2189,7 @@ test "fixed output" {
21692189
21702190test "writeSplat 0 len splat larger than capacity" {
21712191 var buf: [8]u8 = undefined;
2172 var w: std.io.Writer = .fixed(&buf);
2192 var w: Writer = .fixed(&buf);
21732193 const n = try w.writeSplat(&.{"something that overflows buf"}, 0);
21742194 try testing.expectEqual(0, n);
21752195}
......@@ -2188,6 +2208,13 @@ pub fn failingSendFile(w: *Writer, file_reader: *File.Reader, limit: Limit) File
21882208 return error.WriteFailed;
21892209}
21902210
2211pub fn failingRebase(w: *Writer, preserve: usize, capacity: usize) Error!void {
2212 _ = w;
2213 _ = preserve;
2214 _ = capacity;
2215 return error.WriteFailed;
2216}
2217
21912218pub const Discarding = struct {
21922219 count: u64,
21932220 writer: Writer,
......@@ -2455,7 +2482,7 @@ pub fn Hashing(comptime Hasher: type) type {
24552482/// Maintains `Writer` state such that it writes to the unused capacity of an
24562483/// array list, filling it up completely before making a call through the
24572484/// vtable, causing a resize. Consequently, the same, optimized, non-generic
2458/// machine code that uses `std.io.Reader`, such as formatted printing, takes
2485/// machine code that uses `std.Io.Reader`, such as formatted printing, takes
24592486/// the hot paths when using this API.
24602487///
24612488/// When using this API, it is not necessary to call `flush`.
......@@ -2514,6 +2541,7 @@ pub const Allocating = struct {
25142541 .drain = Allocating.drain,
25152542 .sendFile = Allocating.sendFile,
25162543 .flush = noopFlush,
2544 .rebase = growingRebase,
25172545 };
25182546
25192547 pub fn deinit(a: *Allocating) void {
......@@ -2595,7 +2623,7 @@ pub const Allocating = struct {
25952623 return list.items.len - start_len;
25962624 }
25972625
2598 fn sendFile(w: *Writer, file_reader: *File.Reader, limit: std.io.Limit) FileError!usize {
2626 fn sendFile(w: *Writer, file_reader: *File.Reader, limit: Limit) FileError!usize {
25992627 if (File.Handle == void) return error.Unimplemented;
26002628 if (limit == .nothing) return 0;
26012629 const a: *Allocating = @fieldParentPtr("writer", w);
......@@ -2612,6 +2640,15 @@ pub const Allocating = struct {
26122640 return n;
26132641 }
26142642
2643 fn growingRebase(w: *Writer, preserve: usize, minimum_len: usize) Error!void {
2644 _ = preserve; // This implementation always preserves the entire buffer.
2645 const a: *Allocating = @fieldParentPtr("writer", w);
2646 const gpa = a.allocator;
2647 var list = a.toArrayList();
2648 defer setArrayList(a, list);
2649 list.ensureUnusedCapacity(gpa, minimum_len) catch return error.WriteFailed;
2650 }
2651
26152652 fn setArrayList(a: *Allocating, list: std.ArrayListUnmanaged(u8)) void {
26162653 a.writer.buffer = list.allocatedSlice();
26172654 a.writer.end = list.items.len;
......@@ -2645,7 +2682,7 @@ test "discarding sendFile" {
26452682 try file_reader.seekTo(0);
26462683
26472684 var w_buffer: [256]u8 = undefined;
2648 var discarding: std.io.Writer.Discarding = .init(&w_buffer);
2685 var discarding: Writer.Discarding = .init(&w_buffer);
26492686
26502687 _ = try file_reader.interface.streamRemaining(&discarding.writer);
26512688}
......@@ -2664,7 +2701,7 @@ test "allocating sendFile" {
26642701 var file_reader = file_writer.moveToReader();
26652702 try file_reader.seekTo(0);
26662703
2667 var allocating: std.io.Writer.Allocating = .init(testing.allocator);
2704 var allocating: Writer.Allocating = .init(testing.allocator);
26682705 defer allocating.deinit();
26692706
26702707 _ = try file_reader.interface.streamRemaining(&allocating.writer);
......@@ -2702,3 +2739,10 @@ test writeSliceEndian {
27022739 try writeSliceEndian(&w, u16, &array, .big);
27032740 try testing.expectEqualSlices(u8, &.{ 'x', 0x12, 0x34, 0x56, 0x78 }, &buffer);
27042741}
2742
2743test "writableSlice with fixed writer" {
2744 var buf: [2]u8 = undefined;
2745 var w: std.Io.Writer = .fixed(&buf);
2746 try w.writeByte(1);
2747 try std.testing.expectError(error.WriteFailed, w.writableSlice(2));
2748}
lib/std/compress/flate/Decompress.zig+1-3
......@@ -76,7 +76,7 @@ const indirect_vtable: Reader.VTable = .{
7676/// `input` buffer is asserted to be at least 10 bytes, or EOF before then.
7777///
7878/// If `buffer` is provided then asserted to have `flate.max_window_len`
79/// capacity, as well as `flate.history_len` unused capacity on every write.
79/// capacity.
8080pub fn init(input: *Reader, container: Container, buffer: []u8) Decompress {
8181 if (buffer.len != 0) assert(buffer.len >= flate.max_window_len);
8282 return .{
......@@ -239,8 +239,6 @@ fn decodeSymbol(self: *Decompress, decoder: anytype) !Symbol {
239239}
240240
241241fn streamDirect(r: *Reader, w: *Writer, limit: std.Io.Limit) Reader.StreamError!usize {
242 assert(w.buffer.len >= flate.max_window_len);
243 assert(w.unusedCapacityLen() >= flate.history_len);
244242 const d: *Decompress = @alignCast(@fieldParentPtr("reader", r));
245243 return streamFallible(d, w, limit);
246244}