diff --git a/std/io.zig b/std/io.zig index 1c468f6f4..71a982239 100644 --- a/std/io.zig +++ b/std/io.zig @@ -331,6 +331,104 @@ pub fn BufferedInStreamCustom(comptime buffer_size: usize, comptime Error: type) }; } +/// Creates a stream which supports 'un-reading' data, so that it can be read again. +/// This makes look-ahead style parsing much easier. +pub fn PeekStream(comptime buffer_size: usize, comptime InStreamError: type) type { + return struct { + const Self = this; + pub const Error = InStreamError; + pub const Stream = InStream(Error); + + pub stream: Stream, + base: *Stream, + + // Right now the look-ahead space is statically allocated, but a version with dynamic allocation + // is not too difficult to derive from this. + buffer: [buffer_size]u8, + index: usize, + at_end: bool, + + pub fn init(base: *Stream) Self { + return Self{ + .base = base, + .buffer = undefined, + .index = 0, + .at_end = false, + .stream = Stream{ .readFn = readFn }, + }; + } + + pub fn putBackByte(self: *Self, byte: u8) void { + self.buffer[self.index] = byte; + self.index += 1; + } + + pub fn putBack(self: *Self, bytes: []const u8) void { + var pos = bytes.len; + while (pos != 0) { + pos -= 1; + self.putBackByte(bytes[pos]); + } + } + + fn readFn(in_stream: *Stream, dest: []u8) Error!usize { + const self = @fieldParentPtr(Self, "stream", in_stream); + + // copy over anything putBack()'d + var pos: usize = 0; + while (pos < dest.len and self.index != 0) { + dest[pos] = self.buffer[self.index - 1]; + self.index -= 1; + pos += 1; + } + + if (pos == dest.len or self.at_end) { + return pos; + } + + // ask the backing stream for more + const left = dest.len - pos; + const read = try self.base.read(dest[pos..]); + assert(read <= left); + + self.at_end = (read < left); + return pos + read; + } + + }; +} + +pub const SliceStream = struct { + const Self = this; + pub const Error = error { }; + pub const Stream = InStream(Error); + + pub stream: Stream, + + pos: usize, + slice: []const u8, + + pub fn init(slice: []const u8) Self { + return Self{ + .slice = slice, + .pos = 0, + .stream = Stream{ .readFn = readFn }, + }; + } + + fn readFn(in_stream: *Stream, dest: []u8) Error!usize { + const self = @fieldParentPtr(Self, "stream", in_stream); + const size = math.min(dest.len, self.slice.len - self.pos); + const end = self.pos + size; + + mem.copy(u8, dest[0..size], self.slice[self.pos..end]); + self.pos = end; + + return size; + } + +}; + pub fn BufferedOutStream(comptime Error: type) type { return BufferedOutStreamCustom(os.page_size, Error); } diff --git a/std/io_test.zig b/std/io_test.zig index 301a9a4cd..8d5c35c5f 100644 --- a/std/io_test.zig +++ b/std/io_test.zig @@ -60,3 +60,54 @@ test "BufferOutStream" { assert(mem.eql(u8, buffer.toSlice(), "x: 42\ny: 1234\n")); } + +test "SliceStream" { + const bytes = []const u8 { 1, 2, 3, 4, 5, 6, 7 }; + var ss = io.SliceStream.init(bytes); + + var dest: [4]u8 = undefined; + + var read = try ss.stream.read(dest[0..4]); + assert(read == 4); + assert(mem.eql(u8, dest[0..4], bytes[0..4])); + + read = try ss.stream.read(dest[0..4]); + assert(read == 3); + assert(mem.eql(u8, dest[0..3], bytes[4..7])); + + read = try ss.stream.read(dest[0..4]); + assert(read == 0); +} + +test "PeekStream" { + const bytes = []const u8 { 1, 2, 3, 4, 5, 6, 7, 8 }; + var ss = io.SliceStream.init(bytes); + var ps = io.PeekStream(2, io.SliceStream.Error).init(&ss.stream); + + var dest: [4]u8 = undefined; + + ps.putBackByte(9); + ps.putBackByte(10); + + var read = try ps.stream.read(dest[0..4]); + assert(read == 4); + assert(dest[0] == 10); + assert(dest[1] == 9); + assert(mem.eql(u8, dest[2..4], bytes[0..2])); + + read = try ps.stream.read(dest[0..4]); + assert(read == 4); + assert(mem.eql(u8, dest[0..4], bytes[2..6])); + + read = try ps.stream.read(dest[0..4]); + assert(read == 2); + assert(mem.eql(u8, dest[0..2], bytes[6..8])); + + ps.putBackByte(11); + ps.putBackByte(12); + + read = try ps.stream.read(dest[0..4]); + assert(read == 2); + assert(dest[0] == 12); + assert(dest[1] == 11); +}