diff --git a/lib/std/Io/Writer.zig b/lib/std/Io/Writer.zig index 5bb813913ae9b622e1ca9f30bd33ce87c1857a7d..813142a7a35e48c1de0eaba49bea2a5f47ec1139 100644 --- 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 index 57bddc49bb9efc0a7a5ea5082a77f2d3c0314e48..d262ca0c43a52c5364f405cc623b380c252d3915 100644 --- 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; }