1const Limited = @This();
  2
  3const std = @import("../../std.zig");
  4const Reader = std.Io.Reader;
  5const Writer = std.Io.Writer;
  6const Limit = std.Io.Limit;
  7
  8unlimited: *Reader,
  9remaining: Limit,
 10interface: Reader,
 11
 12pub fn init(reader: *Reader, limit: Limit, buffer: []u8) Limited {
 13    return .{
 14        .unlimited = reader,
 15        .remaining = limit,
 16        .interface = .{
 17            .vtable = &.{
 18                .stream = stream,
 19                .discard = discard,
 20            },
 21            .buffer = buffer,
 22            .seek = 0,
 23            .end = 0,
 24        },
 25    };
 26}
 27
 28fn stream(r: *Reader, w: *Writer, limit: Limit) Reader.StreamError!usize {
 29    const l: *Limited = @fieldParentPtr("interface", r);
 30    if (l.remaining == .nothing) return error.EndOfStream;
 31    const combined_limit = limit.min(l.remaining);
 32    const n = try l.unlimited.stream(w, combined_limit);
 33    l.remaining = l.remaining.subtract(n).?;
 34    return n;
 35}
 36
 37test stream {
 38    var orig_buf: [10]u8 = undefined;
 39    @memcpy(&orig_buf, "test bytes");
 40    var fixed: std.Io.Reader = .fixed(&orig_buf);
 41
 42    var limit_buf: [1]u8 = undefined;
 43    var limited: std.Io.Reader.Limited = .init(&fixed, @enumFromInt(4), &limit_buf);
 44
 45    var result_buf: [10]u8 = undefined;
 46    var fixed_writer: std.Io.Writer = .fixed(&result_buf);
 47    const streamed = try limited.interface.stream(&fixed_writer, @enumFromInt(7));
 48
 49    try std.testing.expect(streamed == 4);
 50    try std.testing.expectEqualStrings("test", result_buf[0..streamed]);
 51}
 52
 53fn discard(r: *Reader, limit: Limit) Reader.Error!usize {
 54    const l: *Limited = @fieldParentPtr("interface", r);
 55    if (l.remaining == .nothing) return error.EndOfStream;
 56    const combined_limit = limit.min(l.remaining);
 57    const n = try l.unlimited.discard(combined_limit);
 58    l.remaining = l.remaining.subtract(n).?;
 59    return n;
 60}
 61
 62test "end of stream, read, hit limit exactly" {
 63    var f: Reader = .fixed("i'm dying");
 64    var l = f.limited(.limited(4), &.{});
 65    const r = &l.interface;
 66
 67    var buf: [2]u8 = undefined;
 68    try r.readSliceAll(&buf);
 69    try r.readSliceAll(&buf);
 70    try std.testing.expectError(error.EndOfStream, l.interface.readSliceAll(&buf));
 71}
 72
 73test "end of stream, read, hit limit after partial read" {
 74    var f: Reader = .fixed("i'm dying");
 75    var l = f.limited(.limited(5), &.{});
 76    const r = &l.interface;
 77
 78    var buf: [2]u8 = undefined;
 79    try r.readSliceAll(&buf);
 80    try r.readSliceAll(&buf);
 81    try std.testing.expectError(error.EndOfStream, l.interface.readSliceAll(&buf));
 82}
 83
 84test "end of stream, discard, hit limit exactly" {
 85    var f: Reader = .fixed("i'm dying");
 86    var l = f.limited(.limited(4), &.{});
 87    const r = &l.interface;
 88
 89    try r.discardAll(2);
 90    try r.discardAll(2);
 91    try std.testing.expectError(error.EndOfStream, l.interface.discardAll(2));
 92}
 93
 94test "end of stream, discard, hit limit after partial read" {
 95    var f: Reader = .fixed("i'm dying");
 96    var l = f.limited(.limited(5), &.{});
 97    const r = &l.interface;
 98
 99    try r.discardAll(2);
100    try r.discardAll(2);
101    try std.testing.expectError(error.EndOfStream, l.interface.discardAll(2));
102}