authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-06-05 19:27:51-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-07-01 16:35:29-07:00
log0027cc1b1a2fb7233ccb2f820b58365d9ebf1ea4
treec20c254fced5294298a57ddaf961ceaa8a0e3756
parent084e45fd865e73f862d85da2a634f8f2cb74580d

std.io.Writer: implement writableVectorIterator

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

lib/std/io/Reader.zig+18-23
...@@ -279,14 +279,6 @@ pub fn readVec(r: *Reader, data: []const []u8) Error!usize {...@@ -279,14 +279,6 @@ pub fn readVec(r: *Reader, data: []const []u8) Error!usize {
279 return readVec(r, data, .unlimited);279 return readVec(r, data, .unlimited);
280}280}
281281
282const VectorWrapped = struct {
283 writer: Writer,
284 first: []u8,
285 middle: []const []u8,
286 last: []u8,
287 var unique_address: u8 = undefined;
288};
289
290/// Equivalent to `readVec` but reads at most `limit` bytes.282/// Equivalent to `readVec` but reads at most `limit` bytes.
291///283///
292/// This ultimately will lower to a call to `stream`, but it must ensure284/// 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 {...@@ -313,29 +305,32 @@ pub fn readVecLimit(r: *Reader, data: []const []u8, limit: Limit) Error!usize {
313 r.seek = 0;305 r.seek = 0;
314 r.end = 0;306 r.end = 0;
315 const first = buf[copy_len..];307 const first = buf[copy_len..];
316 var wrapped: VectorWrapped = .{308 const middle = data[i + 1 ..];
317 .first = first,309 var wrapper: Writer.VectorWrapper = .{
318 .middle = data[i + 1 ..],310 .it = .{
319 .last = r.buffer,311 .first = first,
312 .middle = middle,
313 .last = r.buffer,
314 },
320 .writer = .{315 .writer = .{
321 .context = &VectorWrapped.unique_address,316 .context = &Writer.VectorWrapper.unique_address,
322 .buffer = if (first.len >= r.buffer.len) first else r.buffer,317 .buffer = if (first.len >= r.buffer.len) first else r.buffer,
323 .vtable = &.{ .drain = Writer.fixedDrain },318 .vtable = &.{ .drain = Writer.fixedDrain },
324 },319 },
325 };320 };
326 var n = r.vtable.stream(r, &wrapped.writer, .limited(remaining)) catch |err| switch (err) {321 var n = r.vtable.stream(r, &wrapper.writer, .limited(remaining)) catch |err| switch (err) {
327 error.WriteFailed => {322 error.WriteFailed => {
328 if (wrapped.writer.buffer.ptr == first.ptr) {323 if (wrapper.writer.buffer.ptr == first.ptr) {
329 remaining -= wrapped.writer.end;324 remaining -= wrapper.writer.end;
330 } else {325 } else {
331 r.end = wrapped.writer.end;326 r.end = wrapper.writer.end;
332 }327 }
333 break;328 break;
334 },329 },
335 else => |e| return e,330 else => |e| return e,
336 };331 };
337 assert(n == wrapped.writer.end);332 assert(n == wrapper.writer.end);
338 if (wrapped.writer.buffer.ptr != first.ptr) {333 if (wrapper.writer.buffer.ptr != first.ptr) {
339 r.end = n;334 r.end = n;
340 break;335 break;
341 }336 }
...@@ -345,13 +340,13 @@ pub fn readVecLimit(r: *Reader, data: []const []u8, limit: Limit) Error!usize {...@@ -345,13 +340,13 @@ pub fn readVecLimit(r: *Reader, data: []const []u8, limit: Limit) Error!usize {
345 }340 }
346 remaining -= first.len;341 remaining -= first.len;
347 n -= first.len;342 n -= first.len;
348 for (wrapped.middle) |middle| {343 for (middle) |mid| {
349 if (n < middle.len) {344 if (n < mid.len) {
350 remaining -= n;345 remaining -= n;
351 break;346 break;
352 }347 }
353 remaining -= middle.len;348 remaining -= mid.len;
354 n -= middle.len;349 n -= mid.len;
355 }350 }
356 r.end = n;351 r.end = n;
357 break;352 break;
lib/std/io/Writer.zig+43
...@@ -256,6 +256,49 @@ pub fn writableSliceGreedy(w: *Writer, minimum_length: usize) Error![]u8 {...@@ -256,6 +256,49 @@ pub fn writableSliceGreedy(w: *Writer, minimum_length: usize) Error![]u8 {
256 }256 }
257}257}
258258
259pub const WritableVectorIterator = struct {
260 first: []u8,
261 middle: []const []u8 = &.{},
262 last: []u8 = &.{},
263 index: usize = 0,
264
265 pub fn next(it: *WritableVectorIterator) ?[]u8 {
266 while (true) {
267 const i = it.index;
268 it.index += 1;
269 if (i == 0) {
270 if (it.first.len == 0) continue;
271 return it.first;
272 }
273 const middle_index = i - 1;
274 if (middle_index < it.middle.len) {
275 const middle = it.middle[middle_index];
276 if (middle.len == 0) continue;
277 return middle;
278 }
279 if (middle_index == it.middle.len) {
280 if (it.last.len == 0) continue;
281 return it.last;
282 }
283 return null;
284 }
285 }
286};
287
288pub const VectorWrapper = struct {
289 writer: Writer,
290 it: WritableVectorIterator,
291 pub var unique_address: u8 = undefined;
292};
293
294pub fn writableVectorIterator(w: *Writer) Error!WritableVectorIterator {
295 if (w.context == &VectorWrapper.unique_address) {
296 const wrapper: *VectorWrapper = @fieldParentPtr("writer", w);
297 return wrapper.it;
298 }
299 return .{ .first = try writableSliceGreedy(w, 1) };
300}
301
259pub fn ensureUnusedCapacity(w: *Writer, n: usize) Error!void {302pub fn ensureUnusedCapacity(w: *Writer, n: usize) Error!void {
260 _ = try writableSliceGreedy(w, n);303 _ = try writableSliceGreedy(w, n);
261}304}