From 0027cc1b1a2fb7233ccb2f820b58365d9ebf1ea4 Mon Sep 17 00:00:00 2001
From: Andrew Kelley <andrew@ziglang.org>
Date: Thu, 5 Jun 2025 19:27:51 -0700
Subject: [PATCH] std.io.Writer: implement writableVectorIterator

---
 lib/std/io/Reader.zig | 41 ++++++++++++++++++-----------------------
 lib/std/io/Writer.zig | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 23 deletions(-)

diff --git a/lib/std/io/Reader.zig b/lib/std/io/Reader.zig
index d5d8fe11d5..5453e3a8f5 100644
--- a/lib/std/io/Reader.zig
+++ b/lib/std/io/Reader.zig
@@ -279,14 +279,6 @@ pub fn readVec(r: *Reader, data: []const []u8) Error!usize {
     return readVec(r, data, .unlimited);
 }
 
-const VectorWrapped = struct {
-    writer: Writer,
-    first: []u8,
-    middle: []const []u8,
-    last: []u8,
-    var unique_address: u8 = undefined;
-};
-
 /// Equivalent to `readVec` but reads at most `limit` bytes.
 ///
 /// This ultimately will lower to a call to `stream`, but it must ensure
@@ -313,29 +305,32 @@ pub fn readVecLimit(r: *Reader, data: []const []u8, limit: Limit) Error!usize {
         r.seek = 0;
         r.end = 0;
         const first = buf[copy_len..];
-        var wrapped: VectorWrapped = .{
-            .first = first,
-            .middle = data[i + 1 ..],
-            .last = r.buffer,
+        const middle = data[i + 1 ..];
+        var wrapper: Writer.VectorWrapper = .{
+            .it = .{
+                .first = first,
+                .middle = middle,
+                .last = r.buffer,
+            },
             .writer = .{
-                .context = &VectorWrapped.unique_address,
+                .context = &Writer.VectorWrapper.unique_address,
                 .buffer = if (first.len >= r.buffer.len) first else r.buffer,
                 .vtable = &.{ .drain = Writer.fixedDrain },
             },
         };
-        var n = r.vtable.stream(r, &wrapped.writer, .limited(remaining)) catch |err| switch (err) {
+        var n = r.vtable.stream(r, &wrapper.writer, .limited(remaining)) catch |err| switch (err) {
             error.WriteFailed => {
-                if (wrapped.writer.buffer.ptr == first.ptr) {
-                    remaining -= wrapped.writer.end;
+                if (wrapper.writer.buffer.ptr == first.ptr) {
+                    remaining -= wrapper.writer.end;
                 } else {
-                    r.end = wrapped.writer.end;
+                    r.end = wrapper.writer.end;
                 }
                 break;
             },
             else => |e| return e,
         };
-        assert(n == wrapped.writer.end);
-        if (wrapped.writer.buffer.ptr != first.ptr) {
+        assert(n == wrapper.writer.end);
+        if (wrapper.writer.buffer.ptr != first.ptr) {
             r.end = n;
             break;
         }
@@ -345,13 +340,13 @@ pub fn readVecLimit(r: *Reader, data: []const []u8, limit: Limit) Error!usize {
         }
         remaining -= first.len;
         n -= first.len;
-        for (wrapped.middle) |middle| {
-            if (n < middle.len) {
+        for (middle) |mid| {
+            if (n < mid.len) {
                 remaining -= n;
                 break;
             }
-            remaining -= middle.len;
-            n -= middle.len;
+            remaining -= mid.len;
+            n -= mid.len;
         }
         r.end = n;
         break;
diff --git a/lib/std/io/Writer.zig b/lib/std/io/Writer.zig
index 53238082c8..509356de0c 100644
--- a/lib/std/io/Writer.zig
+++ b/lib/std/io/Writer.zig
@@ -256,6 +256,49 @@ pub fn writableSliceGreedy(w: *Writer, minimum_length: usize) Error![]u8 {
     }
 }
 
+pub const WritableVectorIterator = struct {
+    first: []u8,
+    middle: []const []u8 = &.{},
+    last: []u8 = &.{},
+    index: usize = 0,
+
+    pub fn next(it: *WritableVectorIterator) ?[]u8 {
+        while (true) {
+            const i = it.index;
+            it.index += 1;
+            if (i == 0) {
+                if (it.first.len == 0) continue;
+                return it.first;
+            }
+            const middle_index = i - 1;
+            if (middle_index < it.middle.len) {
+                const middle = it.middle[middle_index];
+                if (middle.len == 0) continue;
+                return middle;
+            }
+            if (middle_index == it.middle.len) {
+                if (it.last.len == 0) continue;
+                return it.last;
+            }
+            return null;
+        }
+    }
+};
+
+pub const VectorWrapper = struct {
+    writer: Writer,
+    it: WritableVectorIterator,
+    pub var unique_address: u8 = undefined;
+};
+
+pub fn writableVectorIterator(w: *Writer) Error!WritableVectorIterator {
+    if (w.context == &VectorWrapper.unique_address) {
+        const wrapper: *VectorWrapper = @fieldParentPtr("writer", w);
+        return wrapper.it;
+    }
+    return .{ .first = try writableSliceGreedy(w, 1) };
+}
+
 pub fn ensureUnusedCapacity(w: *Writer, n: usize) Error!void {
     _ = try writableSliceGreedy(w, n);
 }