std.io.Writer: implement writableVectorIterator

This commit is contained in:
Andrew Kelley 2025-06-05 19:27:51 -07:00
parent 084e45fd86
commit 0027cc1b1a
2 changed files with 61 additions and 23 deletions

View File

@ -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;

View File

@ -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);
}