std: let PeekStream have static/dynamic variants

This completes a TODO in the existing implementation

2 files changed, 32 insertions, 12 deletions

diff --git a/lib/std/io.zig b/lib/std/io.zig
index 018c863bbd..2ebf8c3b87 100644
--- a/lib/std/io.zig
+++ b/lib/std/io.zig
@@ -196,7 +196,7 @@ test "io.BufferedInStream" {
 
 /// 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_type: usize, comptime InStreamError: type) type {
+pub fn PeekStream(comptime buffer_type: std.fifo.LinearFifoBufferType, comptime InStreamError: type) type {
     return struct {
         const Self = @This();
         pub const Error = InStreamError;
@@ -205,18 +205,38 @@ pub fn PeekStream(comptime buffer_type: usize, comptime InStreamError: type) typ
         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.
-        const FifoType = std.fifo.LinearFifo(u8, .{ .Static = buffer_size });
+        const FifoType = std.fifo.LinearFifo(u8, buffer_type);
         fifo: FifoType,
 
-        pub fn init(base: *Stream) Self {
-            return .{
-                .base = base,
-                .fifo = FifoType.init(),
-                .stream = Stream{ .readFn = readFn },
-            };
-        }
+        pub usingnamespace switch (buffer_type) {
+            .Static => struct {
+                pub fn init(base: *Stream) Self {
+                    return .{
+                        .base = base,
+                        .fifo = FifoType.init(),
+                        .stream = Stream{ .readFn = readFn },
+                    };
+                }
+            },
+            .Slice => struct {
+                pub fn init(base: *Stream, buf: []u8) Self {
+                    return .{
+                        .base = base,
+                        .fifo = FifoType.init(buf),
+                        .stream = Stream{ .readFn = readFn },
+                    };
+                }
+            },
+            .Dynamic => struct {
+                pub fn init(base: *Stream, allocator: *mem.Allocator) Self {
+                    return .{
+                        .base = base,
+                        .fifo = FifoType.init(allocator),
+                        .stream = Stream{ .readFn = readFn },
+                    };
+                }
+            },
+        };
 
         pub fn putBackByte(self: *Self, byte: u8) !void {
             try self.putBack(@ptrCast([*]const u8, &byte)[0..1]);
diff --git a/lib/std/io/test.zig b/lib/std/io/test.zig
index b2f8307310..6c194b0f3d 100644
--- a/lib/std/io/test.zig
+++ b/lib/std/io/test.zig
@@ -93,7 +93,7 @@ test "SliceInStream" {
 test "PeekStream" {
     const bytes = [_]u8{ 1, 2, 3, 4, 5, 6, 7, 8 };
     var ss = io.SliceInStream.init(&bytes);
-    var ps = io.PeekStream(2, io.SliceInStream.Error).init(&ss.stream);
+    var ps = io.PeekStream(.{ .Static = 2 }, io.SliceInStream.Error).init(&ss.stream);
 
     var dest: [4]u8 = undefined;