| ... | ... | @@ -0,0 +1,54 @@ |
| 1 | const std = @import("../std.zig"); |
| 2 | const io = std.io; |
| 3 | const testing = std.testing; |
| 4 | |
| 5 | /// Takes a tuple of streams, and constructs a new stream that writes to all of them |
| 6 | pub fn MultiOutStream(comptime OutStreams: type) type { |
| 7 | comptime var ErrSet = error{}; |
| 8 | inline for (@typeInfo(OutStreams).Struct.fields) |field| { |
| 9 | const StreamType = field.field_type; |
| 10 | ErrSet = ErrSet || StreamType.Error; |
| 11 | } |
| 12 | |
| 13 | return struct { |
| 14 | const Self = @This(); |
| 15 | |
| 16 | streams: OutStreams, |
| 17 | |
| 18 | pub const Error = ErrSet; |
| 19 | pub const OutStream = io.OutStream(*Self, Error, write); |
| 20 | pub fn outStream(self: *Self) OutStream { |
| 21 | return .{ .context = self }; |
| 22 | } |
| 23 | |
| 24 | pub fn write(self: *Self, bytes: []const u8) Error!usize { |
| 25 | if (comptime self.streams.len == 0) return bytes.len; |
| 26 | |
| 27 | // only first stream is allowed to do a partial write |
| 28 | // all subsequent streams do a `.writeAll` |
| 29 | const bytes_to_write = try self.streams[0].write(bytes); |
| 30 | const slice_to_write = bytes[0..bytes_to_write]; |
| 31 | comptime var i = 1; |
| 32 | inline while (i < self.streams.len) : (i += 1) { |
| 33 | const stream = self.streams[i]; |
| 34 | try stream.writeAll(slice_to_write); |
| 35 | } |
| 36 | return bytes_to_write; |
| 37 | } |
| 38 | }; |
| 39 | } |
| 40 | |
| 41 | pub fn multiOutStream(streams: var) MultiOutStream(@TypeOf(streams)) { |
| 42 | return .{ .streams = streams }; |
| 43 | } |
| 44 | |
| 45 | test "MultiOutStream" { |
| 46 | var buf1: [255]u8 = undefined; |
| 47 | var fbs1 = io.fixedBufferStream(&buf1); |
| 48 | var buf2: [255]u8 = undefined; |
| 49 | var fbs2 = io.fixedBufferStream(&buf2); |
| 50 | var stream = multiOutStream(.{fbs1.outStream(), fbs2.outStream()}); |
| 51 | try stream.outStream().print("HI", .{}); |
| 52 | testing.expectEqualSlices(u8, "HI", fbs1.getWritten()); |
| 53 | testing.expectEqualSlices(u8, "HI", fbs2.getWritten()); |
| 54 | } |