From 68c870934385d58b404d2491bf9dfae8de319d2e Mon Sep 17 00:00:00 2001
From: Leigh McCulloch <leigh@mcchouse.com>
Date: Sat, 8 May 2021 04:21:53 -0700
Subject: [PATCH] io: add MultiWriter (#10021)

---
 vlib/io/multi_writer.v      | 33 +++++++++++++++++++
 vlib/io/multi_writer_test.v | 66 +++++++++++++++++++++++++++++++++++++
 2 files changed, 99 insertions(+)
 create mode 100644 vlib/io/multi_writer.v
 create mode 100644 vlib/io/multi_writer_test.v

diff --git a/vlib/io/multi_writer.v b/vlib/io/multi_writer.v
new file mode 100644
index 0000000000..75d9525729
--- /dev/null
+++ b/vlib/io/multi_writer.v
@@ -0,0 +1,33 @@
+module io
+
+// new_multi_writer returns a Writer that writes to all writers. The write
+// function of the returned Writer writes to all writers of the MultiWriter,
+// returns the length of bytes written, and if any writer fails to write the
+// full length an error is returned and writing to other writers stops, and if
+// any writer returns an error the error is returned immediately and writing to
+// other writers stops.
+pub fn new_multi_writer(writers []Writer) Writer {
+	return &MultiWriter{
+		writers: writers
+	}
+}
+
+// MultiWriter writes to all its writers.
+pub struct MultiWriter {
+pub:
+	writers []Writer
+}
+
+// write writes to all writers of the MultiWriter. Returns the length of bytes
+// written. If any writer fails to write the full length an error is returned
+// and writing to other writers stops. If any writer returns an error the error
+// is returned immediately and writing to other writers stops.
+pub fn (m MultiWriter) write(buf []byte) ?int {
+	for w in m.writers {
+		n := w.write(buf) ?
+		if n != buf.len {
+			return error('io: incomplete write to writer of MultiWriter')
+		}
+	}
+	return buf.len
+}
diff --git a/vlib/io/multi_writer_test.v b/vlib/io/multi_writer_test.v
new file mode 100644
index 0000000000..6d46236bdc
--- /dev/null
+++ b/vlib/io/multi_writer_test.v
@@ -0,0 +1,66 @@
+module io
+
+fn test_multi_writer_write_successful() {
+	w0 := TestWriter{}
+	w1 := TestWriter{}
+	mw := new_multi_writer([w0, w1])
+	n := mw.write('0123456789'.bytes()) or {
+		assert false
+		return
+	}
+	assert n == 10
+	assert w0.bytes == '0123456789'.bytes()
+	assert w1.bytes == '0123456789'.bytes()
+}
+
+fn test_multi_writer_write_incomplete() {
+	w0 := TestWriter{}
+	w1 := TestIncompleteWriter{}
+	mw := new_multi_writer([w0, w1])
+	n := mw.write('0123456789'.bytes()) or {
+		assert w0.bytes == '0123456789'.bytes()
+		assert w1.bytes == '012345678'.bytes()
+		return
+	}
+	assert false
+}
+
+fn test_multi_writer_write_error() {
+	w0 := TestWriter{}
+	w1 := TestErrorWriter{}
+	w2 := TestWriter{}
+	mw := new_multi_writer([w0, w1, w2])
+	n := mw.write('0123456789'.bytes()) or {
+		assert w0.bytes == '0123456789'.bytes()
+		assert w2.bytes == []
+		return
+	}
+	assert false
+}
+
+struct TestWriter {
+pub mut:
+	bytes []byte
+}
+
+fn (mut w TestWriter) write(buf []byte) ?int {
+	w.bytes << buf
+	return buf.len
+}
+
+struct TestIncompleteWriter {
+pub mut:
+	bytes []byte
+}
+
+fn (mut w TestIncompleteWriter) write(buf []byte) ?int {
+	b := buf[..buf.len - 1]
+	w.bytes << b
+	return b.len
+}
+
+struct TestErrorWriter {}
+
+fn (mut w TestErrorWriter) write(buf []byte) ?int {
+	return error('error writer errored')
+}