Merge pull request #5175 from daurnimator/multi-out-stream

std: add io.MultiOutStream

lib/std/io.zig

@@ -123,6 +123,9 @@ pub const cOutStream = @import("io/c_out_stream.zig").cOutStream;
 pub const CountingOutStream = @import("io/counting_out_stream.zig").CountingOutStream;
 pub const countingOutStream = @import("io/counting_out_stream.zig").countingOutStream;
 
+pub const MultiOutStream = @import("io/multi_out_stream.zig").MultiOutStream;
+pub const multiOutStream = @import("io/multi_out_stream.zig").multiOutStream;
+
 pub const BitInStream = @import("io/bit_in_stream.zig").BitInStream;
 pub const bitInStream = @import("io/bit_in_stream.zig").bitInStream;
 

lib/std/io/multi_out_stream.zig

@@ -0,0 +1,51 @@
+const std = @import("../std.zig");
+const io = std.io;
+const testing = std.testing;
+
+/// Takes a tuple of streams, and constructs a new stream that writes to all of them
+pub fn MultiOutStream(comptime OutStreams: type) type {
+    comptime var ErrSet = error{};
+    inline for (@typeInfo(OutStreams).Struct.fields) |field| {
+        const StreamType = field.field_type;
+        ErrSet = ErrSet || StreamType.Error;
+    }
+
+    return struct {
+        const Self = @This();
+
+        streams: OutStreams,
+
+        pub const Error = ErrSet;
+        pub const OutStream = io.OutStream(*Self, Error, write);
+        pub fn outStream(self: *Self) OutStream {
+            return .{ .context = self };
+        }
+
+        pub fn write(self: *Self, bytes: []const u8) Error!usize {
+            var batch = std.event.Batch(Error!void, self.streams.len, .auto_async).init();
+            comptime var i = 0;
+            inline while (i < self.streams.len) : (i += 1) {
+                const stream = self.streams[i];
+                // TODO: remove ptrCast: https://github.com/ziglang/zig/issues/5258
+                batch.add(@ptrCast(anyframe->Error!void, &async stream.writeAll(bytes)));
+            }
+            try batch.wait();
+            return bytes.len;
+        }
+    };
+}
+
+pub fn multiOutStream(streams: var) MultiOutStream(@TypeOf(streams)) {
+    return .{ .streams = streams };
+}
+
+test "MultiOutStream" {
+    var buf1: [255]u8 = undefined;
+    var fbs1 = io.fixedBufferStream(&buf1);
+    var buf2: [255]u8 = undefined;
+    var fbs2 = io.fixedBufferStream(&buf2);
+    var stream = multiOutStream(.{fbs1.outStream(), fbs2.outStream()});
+    try stream.outStream().print("HI", .{});
+    testing.expectEqualSlices(u8, "HI", fbs1.getWritten());
+    testing.expectEqualSlices(u8, "HI", fbs2.getWritten());
+}