textutil: PrefixLineWriter should make single-line writes.

I realized that the whole point of buffering in PrefixLineWriter
is to ensure each underlying Write call contains a full line.
That ensures that if we have multiple PrefixLineWriters writing
out to os.Stdout, we won't be interleaving separate lines.

Note that if we didn't need this property, we could implement
PrefixLineWriter without any buffering or Flush.

So I've changed PrefixLineWriter to ensure single-line-Write.
Also changed to add '\n' to unterminated data when flushed.

Change-Id: I682fac1066bb9989cb232008a7c67959cb2f730a
diff --git a/textutil/writer.go b/textutil/writer.go
index 72b5386..842b48a 100644
--- a/textutil/writer.go
+++ b/textutil/writer.go
@@ -40,10 +40,11 @@
 	return w.w.Write(data)
 }
 
-// PrefixLineWriter returns a WriteFlusher that wraps w.  Any occurrence of EOL
+// PrefixLineWriter returns a WriteFlusher that wraps w.  Each occurrence of EOL
 // (\f, \n, \r, \v, LineSeparator or ParagraphSeparator) causes the preceeding
-// line to be written to w, with the given prefix.  Data without EOL is buffered
-// until the next EOL or Flush call.
+// line to be written to w, with the given prefix, in a single Write call.  Data
+// without EOL is buffered until the next EOL or Flush call.  Flush appends \n to
+// buffered data that doesn't end in EOL.
 //
 // A single Write call on the returned WriteFlusher may result in zero or more
 // Write calls on the underlying w.
@@ -51,18 +52,20 @@
 // If w implements WriteFlusher, each Flush call on the returned WriteFlusher
 // results in exactly one Flush call on the underlying w.
 func PrefixLineWriter(w io.Writer, prefix string) WriteFlusher {
-	return &prefixLineWriter{w, []byte(prefix), nil}
+	return &prefixLineWriter{w, []byte(prefix), len(prefix)}
 }
 
 type prefixLineWriter struct {
-	w      io.Writer
-	prefix []byte
-	buf    []byte
+	w         io.Writer
+	buf       []byte
+	prefixLen int
 }
 
 const eolRunesAsString = "\f\n\r\v" + string(LineSeparator) + string(ParagraphSeparator)
 
 func (w *prefixLineWriter) Write(data []byte) (int, error) {
+	// Write requires that the return arg is in the range [0, len(data)], and we
+	// must return len(data) on success.
 	totalLen := len(data)
 	for len(data) > 0 {
 		index := bytes.IndexAny(data, eolRunesAsString)
@@ -72,17 +75,13 @@
 			w.buf = append(w.buf, data...)
 			return totalLen, nil
 		}
-		// Saw EOL: write prefix, buffer, and data including EOL.
-		if _, err := w.w.Write(w.prefix); err != nil {
-			return totalLen - len(data), err
-		}
-		if _, err := w.w.Write(w.buf); err != nil {
-			return totalLen - len(data), err
-		}
-		w.buf = w.buf[:0]
+		// Saw EOL: single Write of buffer + data including EOL.
 		_, eolSize := utf8.DecodeRune(data[index:])
-		n, err := w.w.Write(data[:index+eolSize])
-		data = data[n:]
+		dataEnd := index + eolSize
+		w.buf = append(w.buf, data[:dataEnd]...)
+		data = data[dataEnd:]
+		_, err := w.w.Write(w.buf)
+		w.buf = w.buf[:w.prefixLen] // reset buf to prepare for the next line.
 		if err != nil {
 			return totalLen - len(data), err
 		}
@@ -98,14 +97,13 @@
 			}
 		}
 	}()
-	if len(w.buf) > 0 {
-		if _, err := w.w.Write(w.prefix); err != nil {
+	if len(w.buf) > w.prefixLen {
+		w.buf = append(w.buf, '\n') // add EOL to unterminated line.
+		_, err := w.w.Write(w.buf)
+		w.buf = w.buf[:w.prefixLen] // reset buf to prepare for the next line.
+		if err != nil {
 			return err
 		}
-		if _, err := w.w.Write(w.buf); err != nil {
-			return err
-		}
-		w.buf = w.buf[:0]
 	}
 	return nil
 }
diff --git a/textutil/writer_test.go b/textutil/writer_test.go
index 8aba8ab..dbb884d 100644
--- a/textutil/writer_test.go
+++ b/textutil/writer_test.go
@@ -8,6 +8,7 @@
 	"bytes"
 	"errors"
 	"fmt"
+	"reflect"
 	"strings"
 	"testing"
 )
@@ -75,53 +76,60 @@
 	tests := []struct {
 		Prefix string
 		Writes []string
-		Want   string
+		Wants  []string
 	}{
-		{"", nil, ""},
-		{"", []string{""}, ""},
-		{"", []string{"a"}, "a"},
-		{"", []string{"a", ""}, "a"},
-		{"", []string{"", "a"}, "a"},
-		{"", []string{"a", "b"}, "ab"},
-		{"", []string{"ab"}, "ab"},
-		{"", []string{"\n"}, "\n"},
-		{"", []string{"\n", ""}, "\n"},
-		{"", []string{"", "\n"}, "\n"},
-		{"", []string{"a", "\n"}, "a\n"},
-		{"", []string{"a\n"}, "a\n"},
-		{"", []string{"\n", "a"}, "\na"},
-		{"", []string{"\na"}, "\na"},
-		{"", []string{"a\nb\nc"}, "a\nb\nc"},
-		{"PRE", nil, ""},
-		{"PRE", []string{""}, ""},
-		{"PRE", []string{"a"}, "PREa"},
-		{"PRE", []string{"a", ""}, "PREa"},
-		{"PRE", []string{"", "a"}, "PREa"},
-		{"PRE", []string{"a", "b"}, "PREab"},
-		{"PRE", []string{"ab"}, "PREab"},
-		{"PRE", []string{"\n"}, "PRE\n"},
-		{"PRE", []string{"\n", ""}, "PRE\n"},
-		{"PRE", []string{"", "\n"}, "PRE\n"},
-		{"PRE", []string{"a", "\n"}, "PREa\n"},
-		{"PRE", []string{"a\n"}, "PREa\n"},
-		{"PRE", []string{"\n", "a"}, "PRE\nPREa"},
-		{"PRE", []string{"\na"}, "PRE\nPREa"},
-		{"PRE", []string{"a", "\n", "b", "\n", "c"}, "PREa\nPREb\nPREc"},
-		{"PRE", []string{"a\nb\nc"}, "PREa\nPREb\nPREc"},
-		{"PRE", []string{"a\nb\nc\n"}, "PREa\nPREb\nPREc\n"},
+		{"", nil, nil},
+		{"", []string{""}, nil},
+		{"", []string{"a"}, []string{"a."}},
+		{"", []string{"a", ""}, []string{"a."}},
+		{"", []string{"", "a"}, []string{"a."}},
+		{"", []string{"a", "b"}, []string{"ab."}},
+		{"", []string{"ab"}, []string{"ab."}},
+		{"", []string{"\n"}, []string{"\n"}},
+		{"", []string{"\n", ""}, []string{"\n"}},
+		{"", []string{"", "\n"}, []string{"\n"}},
+		{"", []string{"a", "\n"}, []string{"a\n"}},
+		{"", []string{"a\n"}, []string{"a\n"}},
+		{"", []string{"\n", "a"}, []string{"\n", "a."}},
+		{"", []string{"\na"}, []string{"\n", "a."}},
+		{"", []string{"a\nb\nc"}, []string{"a\n", "b\n", "c."}},
+		{"", []string{"a\nb\nc\n"}, []string{"a\n", "b\n", "c\n"}},
+		{"PRE", nil, nil},
+		{"PRE", []string{""}, nil},
+		{"PRE", []string{"a"}, []string{"PREa."}},
+		{"PRE", []string{"a", ""}, []string{"PREa."}},
+		{"PRE", []string{"", "a"}, []string{"PREa."}},
+		{"PRE", []string{"a", "b"}, []string{"PREab."}},
+		{"PRE", []string{"ab"}, []string{"PREab."}},
+		{"PRE", []string{"\n"}, []string{"PRE\n"}},
+		{"PRE", []string{"\n", ""}, []string{"PRE\n"}},
+		{"PRE", []string{"", "\n"}, []string{"PRE\n"}},
+		{"PRE", []string{"a", "\n"}, []string{"PREa\n"}},
+		{"PRE", []string{"a\n"}, []string{"PREa\n"}},
+		{"PRE", []string{"\n", "a"}, []string{"PRE\n", "PREa."}},
+		{"PRE", []string{"\na"}, []string{"PRE\n", "PREa."}},
+		{"PRE", []string{"a", "\n", "b", "\n", "c"}, []string{"PREa\n", "PREb\n", "PREc."}},
+		{"PRE", []string{"a\nb\nc"}, []string{"PREa\n", "PREb\n", "PREc."}},
+		{"PRE", []string{"a\nb\nc\n"}, []string{"PREa\n", "PREb\n", "PREc\n"}},
 	}
 	for _, test := range tests {
 		for _, eol := range eolRunesAsString {
-			// Replace \n in Want and Writes with the test eol rune.
-			want := strings.Replace(test.Want, "\n", string(eol), -1)
-			var writes []string
-			for _, write := range test.Writes {
-				writes = append(writes, strings.Replace(write, "\n", string(eol), -1))
+			// Replace '\n' in Writes and Wants with the test eol rune, and replace '.'
+			// in Wants with '\n'.
+			var writes, wants []string
+			for _, x := range test.Writes {
+				x = strings.Replace(x, "\n", string(eol), -1)
+				writes = append(writes, x)
+			}
+			for _, x := range test.Wants {
+				x = strings.Replace(x, "\n", string(eol), -1)
+				x = strings.Replace(x, ".", "\n", -1)
+				wants = append(wants, x)
 			}
 			// Run the actual tests.
-			var buf bytes.Buffer
-			w := PrefixLineWriter(&buf, test.Prefix)
-			name := fmt.Sprintf("(%q, %q)", want, writes)
+			capture := &captureWriter{}
+			w := PrefixLineWriter(capture, test.Prefix)
+			name := fmt.Sprintf("(%q, %q)", wants, writes)
 			for _, write := range writes {
 				name := name + fmt.Sprintf("(%q)", write)
 				n, err := w.Write([]byte(write))
@@ -135,13 +143,22 @@
 			if err := w.Flush(); err != nil {
 				t.Errorf("%s Flush got error: %v", name, err)
 			}
-			if got, want := buf.String(), want; got != want {
+			if got, want := capture.Writes, wants; !reflect.DeepEqual(got, want) {
 				t.Errorf("%s got %q, want %q", name, got, want)
 			}
 		}
 	}
 }
 
+type captureWriter struct {
+	Writes []string
+}
+
+func (w *captureWriter) Write(p []byte) (int, error) {
+	w.Writes = append(w.Writes, string(p))
+	return len(p), nil
+}
+
 var (
 	err1 = errors.New("error 1")
 	err2 = errors.New("error 2")
@@ -243,8 +260,8 @@
 func TestPrefixLineWriter_EOLWriteErrorFlushError(t *testing.T) {
 	fake := &fakeWriteFlusher{writeErr: err1, flushErr: err2}
 	w := PrefixLineWriter(fake, "prefix")
-	if n, err := w.Write([]byte("ab\n")); n != 0 || err != err1 {
-		t.Errorf("Write got (%v,%v), want (0,%v)", n, err, err1)
+	if n, err := w.Write([]byte("ab\n")); n != 3 || err != err1 {
+		t.Errorf("Write got (%v,%v), want (3,%v)", n, err, err1)
 	}
 	if err := w.Flush(); err != err2 {
 		t.Errorf("Flush got error %v, want %v", err, err2)