From 7926081ec35761028c87cdb82204f177311866be Mon Sep 17 00:00:00 2001 From: Andrew Kelley Date: Tue, 24 Jun 2025 02:39:48 -0700 Subject: [PATCH] fix std.fs.File.Reader --- lib/std/fs/File.zig | 103 ++++++++++++++++++++++++++++++++---------- lib/std/io/Reader.zig | 1 - lib/std/io/Writer.zig | 32 ++++++++++++- 3 files changed, 109 insertions(+), 27 deletions(-) diff --git a/lib/std/fs/File.zig b/lib/std/fs/File.zig index 2ab77a6bf5..665e7de02a 100644 --- a/lib/std/fs/File.zig +++ b/lib/std/fs/File.zig @@ -969,6 +969,14 @@ pub const Reader = struct { }; } + pub fn initMode(file: File, buffer: []u8, init_mode: Reader.Mode) Reader { + return .{ + .file = file, + .interface = initInterface(buffer), + .mode = init_mode, + }; + } + pub fn getSize(r: *Reader) GetEndPosError!u64 { return r.size orelse { if (r.size_err) |err| return err; @@ -1038,7 +1046,7 @@ pub const Reader = struct { fn stream(io_reader: *std.io.Reader, w: *std.io.Writer, limit: std.io.Limit) std.io.Reader.StreamError!usize { const r: *Reader = @fieldParentPtr("interface", io_reader); switch (r.mode) { - .positional, .streaming => return w.writeFile(r, limit, &.{}, 0) catch |write_err| switch (write_err) { + .positional, .streaming => return w.sendFile(r, limit) catch |write_err| switch (write_err) { error.ReadFailed => return error.ReadFailed, error.WriteFailed => return error.WriteFailed, error.Unimplemented => { @@ -1056,7 +1064,7 @@ pub const Reader = struct { return n; } var iovecs_buffer: [max_buffers_len]posix.iovec = undefined; - const dest = w.writableVectorPosix(&iovecs_buffer, limit); + const dest = try w.writableVectorPosix(&iovecs_buffer, limit); assert(dest[0].len > 0); const n = posix.preadv(r.file.handle, dest, r.pos) catch |err| switch (err) { error.Unseekable => { @@ -1089,9 +1097,10 @@ pub const Reader = struct { return n; } var iovecs_buffer: [max_buffers_len]posix.iovec = undefined; - const dest = w.writableVectorPosix(&iovecs_buffer, limit); + const dest = try w.writableVectorPosix(&iovecs_buffer, limit); assert(dest[0].len > 0); - const n = posix.pread(r.file.handle, dest) catch |err| { + // TODO also add buffer at the end + const n = posix.readv(r.file.handle, dest) catch |err| { r.err = err; return error.ReadFailed; }; @@ -1296,33 +1305,70 @@ pub const Writer = struct { pub fn drain(io_writer: *std.io.Writer, data: []const []const u8, splat: usize) std.io.Writer.Error!usize { const w: *Writer = @fieldParentPtr("interface", io_writer); const handle = w.file.handle; - if (true) @panic("update to check for buffered data"); + const buffered = io_writer.buffered(); var splat_buffer: [256]u8 = undefined; if (is_windows) { - if (data.len == 1 and splat == 0) return 0; - return windows.WriteFile(handle, data[0], null); + var i: usize = 0; + while (i < buffered.len) { + const n = windows.WriteFile(handle, buffered[i..], null) catch |err| { + w.err = err; + w.pos += i; + _ = io_writer.consume(i); + return error.WriteFailed; + }; + i += n; + if (data.len > 0 and buffered.len - i < n) { + w.pos += i; + return io_writer.consume(i); + } + } + if (i != 0 or data.len == 0 or (data.len == 1 and splat == 0)) { + w.pos += i; + return io_writer.consume(i); + } + const n = windows.WriteFile(handle, data[0], null) catch |err| { + w.err = err; + return 0; + }; + w.pos += n; + return n; + } + if (data.len == 0) { + var i: usize = 0; + while (i < buffered.len) { + i += std.posix.write(handle, buffered) catch |err| { + w.err = err; + w.pos += i; + _ = io_writer.consume(i); + return error.WriteFailed; + }; + } + w.pos += i; + return io_writer.consumeAll(); } var iovecs: [max_buffers_len]std.posix.iovec_const = undefined; - var len: usize = @min(iovecs.len, data.len); - for (iovecs[0..len], data[0..len]) |*v, d| v.* = .{ - .base = if (d.len == 0) "" else d.ptr, // OS sadly checks ptr addr before length. - .len = d.len, - }; + var len: usize = 0; + if (buffered.len > 0) { + iovecs[len] = .{ .base = buffered.ptr, .len = buffered.len }; + len += 1; + } + for (data) |d| { + if (d.len == 0) continue; + if (iovecs.len - len == 0) break; + iovecs[len] = .{ .base = d.ptr, .len = d.len }; + len += 1; + } switch (splat) { - 0 => return std.posix.writev(handle, iovecs[0 .. len - 1]) catch |err| { - w.err = err; - return error.WriteFailed; + 0 => if (data[data.len - 1].len != 0) { + len -= 1; }, - 1 => return std.posix.writev(handle, iovecs[0..len]) catch |err| { - w.err = err; - return error.WriteFailed; - }, - else => { - const pattern = data[data.len - 1]; - if (pattern.len == 1) { + 1 => {}, + else => switch (data[data.len - 1].len) { + 0 => {}, + 1 => { const memset_len = @min(splat_buffer.len, splat); const buf = splat_buffer[0..memset_len]; - @memset(buf, pattern[0]); + @memset(buf, data[data.len - 1][0]); iovecs[len - 1] = .{ .base = buf.ptr, .len = buf.len }; var remaining_splat = splat - buf.len; while (remaining_splat > splat_buffer.len and len < iovecs.len) { @@ -1338,13 +1384,20 @@ pub const Writer = struct { w.err = err; return error.WriteFailed; }; - } + }, + else => for (0..splat - 1) |_| { + if (iovecs.len - len == 0) break; + iovecs[len] = .{ .base = data[data.len - 1].ptr, .len = data[data.len - 1].len }; + len += 1; + }, }, } - return std.posix.writev(handle, iovecs[0..len]) catch |err| { + const n = std.posix.writev(handle, iovecs[0..len]) catch |err| { w.err = err; return error.WriteFailed; }; + w.pos += n; + return io_writer.consume(n); } pub fn sendFile( diff --git a/lib/std/io/Reader.zig b/lib/std/io/Reader.zig index 88d9ebab98..5c999fff78 100644 --- a/lib/std/io/Reader.zig +++ b/lib/std/io/Reader.zig @@ -329,7 +329,6 @@ pub fn readVecLimit(r: *Reader, data: []const []u8, limit: Limit) Error!usize { }, else => |e| return e, }; - assert(n == wrapper.writer.end); if (wrapper.writer.buffer.ptr != first.ptr) { r.end = n; break; diff --git a/lib/std/io/Writer.zig b/lib/std/io/Writer.zig index 769d649329..08f8f8e89c 100644 --- a/lib/std/io/Writer.zig +++ b/lib/std/io/Writer.zig @@ -326,13 +326,29 @@ pub const VectorWrapper = struct { }; pub fn writableVectorIterator(w: *Writer) Error!WritableVectorIterator { - if (w.context == &VectorWrapper.unique_address) { + if (@as(*u8, @ptrCast(w.context)) == &VectorWrapper.unique_address) { const wrapper: *VectorWrapper = @fieldParentPtr("writer", w); return wrapper.it; } return .{ .first = try writableSliceGreedy(w, 1) }; } +pub fn writableVectorPosix(w: *Writer, buffer: []std.posix.iovec, limit: Limit) Error![]std.posix.iovec { + var it = try writableVectorIterator(w); + var i: usize = 0; + var remaining = limit; + while (it.next()) |full_buffer| { + if (!remaining.nonzero()) break; + if (buffer.len - i == 0) break; + const buf = remaining.slice(full_buffer); + if (buf.len == 0) continue; + buffer[i] = .{ .base = buf.ptr, .len = buf.len }; + i += 1; + remaining = remaining.subtract(buf.len).?; + } + return buffer[0..i]; +} + pub fn ensureUnusedCapacity(w: *Writer, n: usize) Error!void { _ = try writableSliceGreedy(w, n); } @@ -353,6 +369,13 @@ pub fn advance(w: *Writer, n: usize) void { w.count += n; } +/// After calling `writableVector`, this function tracks how many bytes were +/// written to it. +pub fn advanceVector(w: *Writer, n: usize) usize { + w.count += n; + return consume(w, n); +} + /// The `data` parameter is mutable because this function needs to mutate the /// fields in order to handle partial writes from `VTable.writeSplat`. pub fn writeVecAll(w: *Writer, data: [][]const u8) Error!void { @@ -1830,6 +1853,13 @@ pub fn consume(w: *Writer, n: usize) usize { return n - w.end; } +/// Shortcut for setting `end` to zero and returning zero. Equivalent to +/// calling `consume` with `end`. +pub fn consumeAll(w: *Writer) usize { + w.end = 0; + return 0; +} + /// For use when the `Writer` implementation can cannot offer a more efficient /// implementation than a basic read/write loop on the file. pub fn unimplementedSendFile(w: *Writer, file_reader: *File.Reader, limit: Limit) FileError!usize {