std: add LimitedReader: reader that returns EOF early

2 files changed, 53 insertions, 0 deletions

diff --git a/lib/std/io.zig b/lib/std/io.zig
index fb5f00805c..fa0d5d8698 100644
--- a/lib/std/io.zig
+++ b/lib/std/io.zig
@@ -125,6 +125,9 @@ pub const fixedBufferStream = @import("io/fixed_buffer_stream.zig").fixedBufferS
 pub const CWriter = @import("io/c_writer.zig").CWriter;
 pub const cWriter = @import("io/c_writer.zig").cWriter;
 
+pub const EarlyEOFReader = @import("io/early_eof_reader.zig").EarlyEOFReader;
+pub const earlyEOFReader = @import("io/early_eof_reader.zig").earlyEOFReader;
+
 pub const CountingWriter = @import("io/counting_writer.zig").CountingWriter;
 pub const countingWriter = @import("io/counting_writer.zig").countingWriter;
 pub const CountingReader = @import("io/counting_reader.zig").CountingReader;
diff --git a/lib/std/io/early_eof_reader.zig b/lib/std/io/early_eof_reader.zig
new file mode 100644
index 0000000000..5ddb340455
--- /dev/null
+++ b/lib/std/io/early_eof_reader.zig
@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2015-2020 Zig Contributors
+// This file is part of [zig](https://ziglang.org/), which is MIT licensed.
+// The MIT license requires this copyright notice to be included in all copies
+// and substantial portions of the software.
+const std = @import("../std.zig");
+const io = std.io;
+const assert = std.debug.assert;
+const testing = std.testing;
+
+pub fn EarlyEOFReader(comptime ReaderType: type) type {
+    return struct {
+        inner_reader: ReaderType,
+        bytes_left: u64,
+
+        pub const Error = ReaderType.Error;
+        pub const Reader = io.Reader(*Self, Error, read);
+
+        const Self = @This();
+
+        pub fn read(self: *Self, dest: []u8) Error!usize {
+            const max_read = std.math.min(self.bytes_left, dest.len);
+            const n = try self.inner_reader.read(dest[0..max_read]);
+            self.bytes_left -= n;
+            return n;
+        }
+
+        pub fn reader(self: *Self) Reader {
+            return .{ .context = self };
+        }
+    };
+}
+
+/// Returns an initialised `EarlyEOFReader`
+/// `bytes_left` is a `u64` to be able to take 64 bit file offsets
+pub fn earlyEOFReader(inner_reader: anytype, bytes_left: u64) EarlyEOFReader(@TypeOf(inner_reader)) {
+    return .{ .inner_reader = inner_reader, .bytes_left = bytes_left };
+}
+
+test "io.EarlyEOFReader" {
+    const data = "hello world";
+    var fbs = std.io.fixedBufferStream(data);
+    var early_stream = earlyEOFReader(fbs.reader(), 3);
+
+    var buf: [5]u8 = undefined;
+    testing.expectEqual(@as(usize, 3), try early_stream.reader().read(&buf));
+    testing.expectEqualSlices(u8, data[0..3], buf[0..3]);
+    testing.expectEqual(@as(usize, 0), try early_stream.reader().read(&buf));
+    testing.expectError(error.EndOfStream, early_stream.reader().skipBytes(10, .{}));
+}