From 9c514e149c1d986c404460121c1a5fc01ab8356c Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Tue, 21 Jul 2015 12:24:33 -0700
Subject: [PATCH] io: add WriteString support to MultiWriter

Fixes #11805

Change-Id: I081e16b869dc706bd847ee645bb902bc671c123f
Reviewed-on: https://go-review.googlesource.com/12485
Reviewed-by: Josh Bleecher Snyder <josharian@gmail.com>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
---
 src/io/multi.go      | 24 ++++++++++++++++++++
 src/io/multi_test.go | 53 +++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/src/io/multi.go b/src/io/multi.go
index e26cc53e9e..16860aa361 100644
--- a/src/io/multi.go
+++ b/src/io/multi.go
@@ -52,6 +52,30 @@ func (t *multiWriter) Write(p []byte) (n int, err error) {
 	return len(p), nil
 }
 
+var _ stringWriter = (*multiWriter)(nil)
+
+func (t *multiWriter) WriteString(s string) (n int, err error) {
+	var p []byte // lazily initialized if/when needed
+	for _, w := range t.writers {
+		if sw, ok := w.(stringWriter); ok {
+			n, err = sw.WriteString(s)
+		} else {
+			if p == nil {
+				p = []byte(s)
+			}
+			n, err = w.Write(p)
+		}
+		if err != nil {
+			return
+		}
+		if n != len(s) {
+			err = ErrShortWrite
+			return
+		}
+	}
+	return len(s), nil
+}
+
 // MultiWriter creates a writer that duplicates its writes to all the
 // provided writers, similar to the Unix tee(1) command.
 func MultiWriter(writers ...Writer) Writer {
diff --git a/src/io/multi_test.go b/src/io/multi_test.go
index 56c6769a9e..e80592e0f5 100644
--- a/src/io/multi_test.go
+++ b/src/io/multi_test.go
@@ -62,8 +62,59 @@ func TestMultiReader(t *testing.T) {
 }
 
 func TestMultiWriter(t *testing.T) {
-	sha1 := sha1.New()
 	sink := new(bytes.Buffer)
+	// Hide bytes.Buffer's WriteString method:
+	testMultiWriter(t, struct {
+		Writer
+		fmt.Stringer
+	}{sink, sink})
+}
+
+func TestMultiWriter_String(t *testing.T) {
+	testMultiWriter(t, new(bytes.Buffer))
+}
+
+// test that a multiWriter.WriteString calls results in at most 1 allocation,
+// even if multiple targets don't support WriteString.
+func TestMultiWriter_WriteStringSingleAlloc(t *testing.T) {
+	var sink1, sink2 bytes.Buffer
+	type simpleWriter struct { // hide bytes.Buffer's WriteString
+		Writer
+	}
+	mw := MultiWriter(simpleWriter{&sink1}, simpleWriter{&sink2})
+	allocs := int(testing.AllocsPerRun(1000, func() {
+		WriteString(mw, "foo")
+	}))
+	if allocs != 1 {
+		t.Errorf("num allocations = %d; want 1", allocs)
+	}
+}
+
+type writeStringChecker struct{ called bool }
+
+func (c *writeStringChecker) WriteString(s string) (n int, err error) {
+	c.called = true
+	return len(s), nil
+}
+
+func (c *writeStringChecker) Write(p []byte) (n int, err error) {
+	return len(p), nil
+}
+
+func TestMultiWriter_StringCheckCall(t *testing.T) {
+	var c writeStringChecker
+	mw := MultiWriter(&c)
+	WriteString(mw, "foo")
+	if !c.called {
+		t.Error("did not see WriteString call to writeStringChecker")
+	}
+}
+
+func testMultiWriter(t *testing.T, sink interface {
+	Writer
+	fmt.Stringer
+}) {
+	sha1 := sha1.New()
 	mw := MultiWriter(sha1, sink)
 
 	sourceString := "My input text."
-- 
2.48.1