authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-07-24 23:30:46-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-07-25 14:15:33-07:00
log7f1c04423e5cfd628c58824aee543b4e79525c5f
tree353c58790bd4781f591902060038729cd67fc645
parent3687eada6a4c7a2f9444b090f49094c9755cb364

std.Io: add "preserve" variants to Reader/Writer

2 files changed, 96 insertions(+), 23 deletions(-)

lib/std/Io/Reader.zig+26
......@@ -185,6 +185,32 @@ pub fn streamExact64(r: *Reader, w: *Writer, n: u64) StreamError!void {
185185 while (remaining != 0) remaining -= try r.stream(w, .limited64(remaining));
186186}
187187
188/// "Pump" exactly `n` bytes from the reader to the writer.
189///
190/// When draining `w`, ensures that at least `preserve_len` bytes remain
191/// buffered.
192///
193/// Asserts `Writer.buffer` capacity exceeds `preserve_len`.
194pub fn streamExactPreserve(r: *Reader, w: *Writer, preserve_len: usize, n: usize) StreamError!void {
195 if (w.end + n <= w.buffer.len) {
196 @branchHint(.likely);
197 return streamExact(r, w, n);
198 }
199 // If `n` is large, we can ignore `preserve_len` up to a point.
200 var remaining = n;
201 while (remaining > preserve_len) {
202 assert(remaining != 0);
203 remaining -= try r.stream(w, .limited(remaining - preserve_len));
204 if (w.end + remaining <= w.buffer.len) return streamExact(r, w, remaining);
205 }
206 // All the next bytes received must be preserved.
207 if (preserve_len < w.end) {
208 @memmove(w.buffer[0..preserve_len], w.buffer[w.end - preserve_len ..][0..preserve_len]);
209 w.end = preserve_len;
210 }
211 return streamExact(r, w, remaining);
212}
213
188214/// "Pump" data from the reader to the writer, handling `error.EndOfStream` as
189215/// a success case.
190216///
lib/std/Io/Writer.zig+70-23
......@@ -256,10 +256,10 @@ test "fixed buffer flush" {
256256 try testing.expectEqual(10, buffer[0]);
257257}
258258
259/// Calls `VTable.drain` but hides the last `preserve_length` bytes from the
259/// Calls `VTable.drain` but hides the last `preserve_len` bytes from the
260260/// implementation, keeping them buffered.
261pub fn drainPreserve(w: *Writer, preserve_length: usize) Error!void {
262 const temp_end = w.end -| preserve_length;
261pub fn drainPreserve(w: *Writer, preserve_len: usize) Error!void {
262 const temp_end = w.end -| preserve_len;
263263 const preserved = w.buffer[temp_end..w.end];
264264 w.end = temp_end;
265265 defer w.end += preserved.len;
......@@ -310,24 +310,38 @@ pub fn writableSliceGreedy(w: *Writer, minimum_length: usize) Error![]u8 {
310310}
311311
312312/// Asserts the provided buffer has total capacity enough for `minimum_length`
313/// and `preserve_length` combined.
313/// and `preserve_len` combined.
314314///
315315/// Does not `advance` the buffer end position.
316316///
317/// When draining the buffer, ensures that at least `preserve_length` bytes
317/// When draining the buffer, ensures that at least `preserve_len` bytes
318318/// remain buffered.
319319///
320/// If `preserve_length` is zero, this is equivalent to `writableSliceGreedy`.
321pub fn writableSliceGreedyPreserve(w: *Writer, preserve_length: usize, minimum_length: usize) Error![]u8 {
322 assert(w.buffer.len >= preserve_length + minimum_length);
320/// If `preserve_len` is zero, this is equivalent to `writableSliceGreedy`.
321pub fn writableSliceGreedyPreserve(w: *Writer, preserve_len: usize, minimum_length: usize) Error![]u8 {
322 assert(w.buffer.len >= preserve_len + minimum_length);
323323 while (w.buffer.len - w.end < minimum_length) {
324 try drainPreserve(w, preserve_length);
324 try drainPreserve(w, preserve_len);
325325 } else {
326326 @branchHint(.likely);
327327 return w.buffer[w.end..];
328328 }
329329}
330330
331/// Asserts the provided buffer has total capacity enough for `len`.
332///
333/// Advances the buffer end position by `len`.
334///
335/// When draining the buffer, ensures that at least `preserve_len` bytes
336/// remain buffered.
337///
338/// If `preserve_len` is zero, this is equivalent to `writableSlice`.
339pub fn writableSlicePreserve(w: *Writer, preserve_len: usize, len: usize) Error![]u8 {
340 const big_slice = try w.writableSliceGreedyPreserve(preserve_len, len);
341 advance(w, len);
342 return big_slice[0..len];
343}
344
331345pub const WritableVectorIterator = struct {
332346 first: []u8,
333347 middle: []const []u8 = &.{},
......@@ -523,16 +537,16 @@ pub fn write(w: *Writer, bytes: []const u8) Error!usize {
523537 return w.vtable.drain(w, &.{bytes}, 1);
524538}
525539
526/// Asserts `buffer` capacity exceeds `preserve_length`.
527pub fn writePreserve(w: *Writer, preserve_length: usize, bytes: []const u8) Error!usize {
528 assert(preserve_length <= w.buffer.len);
540/// Asserts `buffer` capacity exceeds `preserve_len`.
541pub fn writePreserve(w: *Writer, preserve_len: usize, bytes: []const u8) Error!usize {
542 assert(preserve_len <= w.buffer.len);
529543 if (w.end + bytes.len <= w.buffer.len) {
530544 @branchHint(.likely);
531545 @memcpy(w.buffer[w.end..][0..bytes.len], bytes);
532546 w.end += bytes.len;
533547 return bytes.len;
534548 }
535 const temp_end = w.end -| preserve_length;
549 const temp_end = w.end -| preserve_len;
536550 const preserved = w.buffer[temp_end..w.end];
537551 w.end = temp_end;
538552 defer w.end += preserved.len;
......@@ -552,13 +566,13 @@ pub fn writeAll(w: *Writer, bytes: []const u8) Error!void {
552566/// Calls `drain` as many times as necessary such that all of `bytes` are
553567/// transferred.
554568///
555/// When draining the buffer, ensures that at least `preserve_length` bytes
569/// When draining the buffer, ensures that at least `preserve_len` bytes
556570/// remain buffered.
557571///
558/// Asserts `buffer` capacity exceeds `preserve_length`.
559pub fn writeAllPreserve(w: *Writer, preserve_length: usize, bytes: []const u8) Error!void {
572/// Asserts `buffer` capacity exceeds `preserve_len`.
573pub fn writeAllPreserve(w: *Writer, preserve_len: usize, bytes: []const u8) Error!void {
560574 var index: usize = 0;
561 while (index < bytes.len) index += try w.writePreserve(preserve_length, bytes[index..]);
575 while (index < bytes.len) index += try w.writePreserve(preserve_len, bytes[index..]);
562576}
563577
564578/// Renders fmt string with args, calling `writer` with slices of bytes.
......@@ -761,11 +775,11 @@ pub fn writeByte(w: *Writer, byte: u8) Error!void {
761775 }
762776}
763777
764/// When draining the buffer, ensures that at least `preserve_length` bytes
778/// When draining the buffer, ensures that at least `preserve_len` bytes
765779/// remain buffered.
766pub fn writeBytePreserve(w: *Writer, preserve_length: usize, byte: u8) Error!void {
780pub fn writeBytePreserve(w: *Writer, preserve_len: usize, byte: u8) Error!void {
767781 while (w.buffer.len - w.end == 0) {
768 try drainPreserve(w, preserve_length);
782 try drainPreserve(w, preserve_len);
769783 } else {
770784 @branchHint(.likely);
771785 w.buffer[w.end] = byte;
......@@ -788,10 +802,42 @@ test splatByteAll {
788802 try testing.expectEqualStrings("7" ** 45, aw.writer.buffered());
789803}
790804
805pub fn splatBytePreserve(w: *Writer, preserve_len: usize, byte: u8, n: usize) Error!void {
806 const new_end = w.end + n;
807 if (new_end <= w.buffer.len) {
808 @memset(w.buffer[w.end..][0..n], byte);
809 w.end = new_end;
810 return;
811 }
812 // If `n` is large, we can ignore `preserve_len` up to a point.
813 var remaining = n;
814 while (remaining > preserve_len) {
815 assert(remaining != 0);
816 remaining -= try splatByte(w, byte, remaining - preserve_len);
817 if (w.end + remaining <= w.buffer.len) {
818 @memset(w.buffer[w.end..][0..remaining], byte);
819 w.end += remaining;
820 return;
821 }
822 }
823 // All the next bytes received must be preserved.
824 if (preserve_len < w.end) {
825 @memmove(w.buffer[0..preserve_len], w.buffer[w.end - preserve_len ..][0..preserve_len]);
826 w.end = preserve_len;
827 }
828 while (remaining > 0) remaining -= try w.splatByte(byte, remaining);
829}
830
791831/// Writes the same byte many times, allowing short writes.
792832///
793833/// Does maximum of one underlying `VTable.drain`.
794834pub fn splatByte(w: *Writer, byte: u8, n: usize) Error!usize {
835 if (w.end + n <= w.buffer.len) {
836 @branchHint(.likely);
837 @memset(w.buffer[w.end..][0..n], byte);
838 w.end += n;
839 return n;
840 }
795841 return writeSplat(w, &.{&.{byte}}, n);
796842}
797843
......@@ -801,9 +847,10 @@ pub fn splatBytesAll(w: *Writer, bytes: []const u8, splat: usize) Error!void {
801847 var remaining_bytes: usize = bytes.len * splat;
802848 remaining_bytes -= try w.splatBytes(bytes, splat);
803849 while (remaining_bytes > 0) {
804 const leftover = remaining_bytes % bytes.len;
805 const buffers: [2][]const u8 = .{ bytes[bytes.len - leftover ..], bytes };
806 remaining_bytes -= try w.writeSplat(&buffers, splat);
850 const leftover_splat = remaining_bytes / bytes.len;
851 const leftover_bytes = remaining_bytes % bytes.len;
852 const buffers: [2][]const u8 = .{ bytes[bytes.len - leftover_bytes ..], bytes };
853 remaining_bytes -= try w.writeSplat(&buffers, leftover_splat);
807854 }
808855}
809856