fmt: Implement pp.WriteString method

This allows io.WriteString to make use of WriteString method
implemented by pp when writing a string to fmt.State.

Fixes #20786

Change-Id: Ice7a92bf303127ad87f05562217fa076f5c589ad
Reviewed-on: https://go-review.googlesource.com/61430
Reviewed-by: Rob Pike <r@golang.org>

2 files changed, 25 insertions, 0 deletions

diff --git a/src/fmt/fmt_test.go b/src/fmt/fmt_test.go
index e6239a51ba..0e09f16dbf 100644
--- a/src/fmt/fmt_test.go
+++ b/src/fmt/fmt_test.go
@@ -131,6 +131,19 @@ func (byteFormatter) Format(f State, _ rune) {
 
 var byteFormatterSlice = []byteFormatter{'h', 'e', 'l', 'l', 'o'}
 
+// Copy of io.stringWriter interface used by writeStringFormatter for type assertion.
+type stringWriter interface {
+	WriteString(s string) (n int, err error)
+}
+
+type writeStringFormatter string
+
+func (sf writeStringFormatter) Format(f State, c rune) {
+	if sw, ok := f.(stringWriter); ok {
+		sw.WriteString("***" + string(sf) + "***")
+	}
+}
+
 var fmtTests = []struct {
 	fmt string
 	val interface{}
@@ -977,6 +990,11 @@ var fmtTests = []struct {
 	// This next case seems wrong, but the docs say the Formatter wins here.
 	{"%#v", byteFormatterSlice, "[]fmt_test.byteFormatter{X, X, X, X, X}"},
 
+	// pp.WriteString
+	{"%s", writeStringFormatter(""), "******"},
+	{"%s", writeStringFormatter("xyz"), "***xyz***"},
+	{"%s", writeStringFormatter("⌘/⌘"), "***⌘/⌘***"},
+
 	// reflect.Value handled specially in Go 1.5, making it possible to
 	// see inside non-exported fields (which cannot be accessed with Interface()).
 	// Issue 8965.
diff --git a/src/fmt/print.go b/src/fmt/print.go
index 38ce928fc0..98c156a121 100644
--- a/src/fmt/print.go
+++ b/src/fmt/print.go
@@ -172,6 +172,13 @@ func (p *pp) Write(b []byte) (ret int, err error) {
 	return len(b), nil
 }
 
+// Implement WriteString so that we can call io.WriteString
+// on a pp (through state), for efficiency.
+func (p *pp) WriteString(s string) (ret int, err error) {
+	p.buf.WriteString(s)
+	return len(s), nil
+}
+
 // These routines end in 'f' and take a format string.
 
 // Fprintf formats according to a format specifier and writes to w.