Merge pull request #1282 from nwsharp/master

std.io: PeekStream and SliceStream

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

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