commit 71f23402bcbb8554d5663ac5aa870abd69fda10f (tree)
parent 5b647b792c680a32c44823a050672537424c95c1
Author: Ryan Liptak <squeek502@hotmail.com>
Date: Mon, 18 May 2026 03:47:27 -0700
Writer: Fix splatBytePreserve losing data in certain cases
In cases where the now removed `preserve < w.end` branch was taken, some of the data in the buffer would never make it to logical sink. For example, in one of the newly added test cases where a buffer of 10 is filled with 9 bytes, and then `splatBytePreserve` is called with a preserve of 5 and a splat of 2, the 5 preserved bytes in the buffer would get memmoved to the front of the buffer in the `preserve < w.end` branch, clobbering the first 4 bytes (without any chance of them making it to the logical sink).
After this commit, `rebase` is called to allow the implementation to do the preservation while sending any relevant bytes to the logical sink in the process.
Addresses part of https://github.com/ziglang/zig/issues/24767
Supersedes https://github.com/ziglang/zig/pull/24875
Diffstat:
2 files changed, 112 insertions(+), 6 deletions(-)
diff --git a/lib/std/Io/Writer.zig b/lib/std/Io/Writer.zig
@@ -802,12 +802,14 @@ pub fn splatBytePreserve(w: *Writer, preserve: usize, byte: u8, n: usize) Error!
return;
}
}
- // All the next bytes received must be preserved.
- if (preserve < w.end) {
- @memmove(w.buffer[0..preserve], w.buffer[w.end - preserve ..][0..preserve]);
- w.end = preserve;
- }
- while (remaining > 0) remaining -= try w.splatByte(byte, remaining);
+ // Ensure the contract of `rebase` is upheld.
+ assert(w.end + remaining > w.buffer.len);
+ // Offset the amount preserved by the amount we have left to splat
+ // since the remaining splat is always going to be part of that
+ // preservation.
+ try w.vtable.rebase(w, preserve -| remaining, remaining);
+ @memset(w.buffer[w.end..][0..remaining], byte);
+ w.end += remaining;
}
/// Writes the same byte many times, allowing short writes.
@@ -2887,3 +2889,46 @@ test "writableSlice with fixed writer" {
try w.writeByte(1);
try std.testing.expectError(error.WriteFailed, w.writableSlice(2));
}
+
+test splatBytePreserve {
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 5, .splat_len = 5 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 9, .preserve = 5, .splat_len = 2 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 5, .splat_len = 6 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 6, .splat_len = 6 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 5, .splat_len = 10 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 6, .splat_len = 10 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 6, .splat_len = 11 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 6, .splat_len = 80 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 6, .splat_len = 85 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 10, .splat_len = 6 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 10, .splat_len = 11 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 10, .splat_len = 80 });
+ try testSplatBytePreserve(.{ .buf_len = 10, .fill_len = 5, .preserve = 10, .splat_len = 85 });
+}
+
+fn testSplatBytePreserve(options: struct { buf_len: u4, fill_len: u4, preserve: u4, splat_len: u8 }) !void {
+ assert(options.fill_len <= options.buf_len);
+ assert(options.preserve <= options.buf_len);
+
+ const fill_buf = "abcdefghijklmno";
+ const fill = fill_buf[0..options.fill_len];
+ var expected_out_buf: [256]u8 = @splat('X');
+ @memcpy(expected_out_buf[0..options.fill_len], fill);
+ const expected_out = expected_out_buf[0 .. options.fill_len + options.splat_len];
+ const expected_preserved = expected_out[expected_out.len -| options.preserve..];
+
+ var out_buf: [256]u8 = undefined;
+ var fw: Writer = .fixed(&out_buf);
+ var indirect_buffer: [16]u8 = undefined;
+ var twi: std.testing.WriterIndirect = .init(&fw, indirect_buffer[0..options.buf_len]);
+ const w = &twi.interface;
+
+ try w.writeAll(fill);
+ try w.splatBytePreserve(options.preserve, 'X', options.splat_len);
+
+ try std.testing.expectEqualStrings(expected_preserved, w.buffer[w.end -| options.preserve..w.end]);
+
+ try w.flush();
+
+ try std.testing.expectEqualStrings(expected_out, fw.buffer[0..fw.end]);
+}
diff --git a/lib/std/testing.zig b/lib/std/testing.zig
@@ -1334,6 +1334,67 @@ pub const ReaderIndirect = struct {
}
};
+/// A `Io.Writer` that writes its data to another `Io.Writer`, and only
+/// writes new data to its own buffer during `drain`.
+pub const WriterIndirect = struct {
+ out: *Io.Writer,
+ interface: Io.Writer,
+
+ pub fn init(out: *Io.Writer, buffer: []u8) WriterIndirect {
+ return .{
+ .out = out,
+ .interface = .{
+ .vtable = &.{
+ .drain = drain,
+ },
+ .buffer = buffer,
+ .end = 0,
+ },
+ };
+ }
+
+ fn drain(w: *Io.Writer, data: []const []const u8, splat: usize) std.Io.Writer.Error!usize {
+ const w_indirect: *WriterIndirect = @alignCast(@fieldParentPtr("interface", w));
+
+ // Write all data in the buffer to `out`
+ try w_indirect.out.writeAll(w.buffer[0..w.end]);
+ w.end = 0;
+
+ // Refill buffer using data
+ {
+ const end_before_fill = w.end;
+ for (data[0 .. data.len - 1]) |bytes| {
+ const dest = w.buffer[w.end..];
+ const len = @min(bytes.len, dest.len);
+ @memcpy(dest[0..len], bytes[0..len]);
+ w.end += len;
+ }
+ const pattern = data[data.len - 1];
+ switch (pattern.len) {
+ 0 => {},
+ 1 => {
+ const len = @min(w.buffer[w.end..].len, splat);
+ @memset(w.buffer[w.end..][0..len], pattern[0]);
+ w.end += len;
+ },
+ else => {
+ const dest = w.buffer[w.end..];
+ for (0..splat) |i| {
+ const start_i = i * pattern.len;
+ if (start_i >= dest.len) break;
+ const remaining = dest[start_i..];
+ const len = @min(pattern.len, remaining.len);
+ @memcpy(remaining[0..len], pattern[0..len]);
+ w.end += len;
+ }
+ },
+ }
+
+ return w.end - end_before_fill;
+ }
+ }
+};
+
test {
_ = &Smith;
}
