Stream printed diffs instead of buffering them three times

diff/diff_test.go

@@ -5,6 +5,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"io"
 	"io/ioutil"
+	"math"
 	"path/filepath"
 	"reflect"
 	"strings"
@@ -952,6 +953,51 @@ func TestParseMultiFileDiffAndPrintMultiFileDiff(t *testing.T) {
 		if !bytes.Equal(printed, diffData) {
 			t.Errorf("%s: printed multi-file diff != original multi-file diff\n\n# PrintMultiFileDiff output - Original:\n%s", test.filename, cmp.Diff(diffData, printed))
 		}
+
+		var written bytes.Buffer
+		if err := WriteMultiFileDiff(&written, diffs); err != nil {
+			t.Errorf("%s: WriteMultiFileDiff: %s", test.filename, err)
+		}
+		if !bytes.Equal(written.Bytes(), diffData) {
+			t.Errorf("%s: written multi-file diff != original multi-file diff\n\n# WriteMultiFileDiff output - Original:\n%s", test.filename, cmp.Diff(diffData, written.Bytes()))
+		}
+	}
+}
+
+// failingWriter fails the write that goes past the first n bytes and accepts
+// all other writes, so a caller that ignores that one error goes unnoticed.
+type failingWriter struct{ n int }
+
+func (w *failingWriter) Write(p []byte) (int, error) {
+	if len(p) > w.n {
+		n := w.n
+		w.n = math.MaxInt
+		return n, io.ErrShortWrite
+	}
+	w.n -= len(p)
+	return len(p), nil
+}
+
+func TestWriteMultiFileDiffReturnsWriteError(t *testing.T) {
+	// Covers extended headers, "Only in" lines, timestamps, sections and
+	// missing newlines, so that every write in the printer can fail.
+	diffData, err := ioutil.ReadFile(filepath.Join("testdata", "sample_contains_added_deleted_files.diff"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	diffData = append(diffData, "--- a\t2009-10-11 15:12:20.000000000 -0700\n+++ b\n@@ -1,1 +1,1 @@ section\n-a\n\\ No newline at end of file\n+b\n\\ No newline at end of file\n"...)
+	diffs, err := ParseMultiFileDiff(diffData)
+	if err != nil {
+		t.Fatal(err)
+	}
+	printed, err := PrintMultiFileDiff(diffs)
+	if err != nil || !bytes.Equal(printed, diffData) {
+		t.Fatalf("PrintMultiFileDiff did not round-trip: %v", err)
+	}
+	for n := 0; n < len(printed); n++ {
+		if err := WriteMultiFileDiff(&failingWriter{n: n}, diffs); err != io.ErrShortWrite {
+			t.Fatalf("failing after %d bytes: got err %v, want %v", n, err, io.ErrShortWrite)
+		}
 	}
 }
 

diff/print.go

@@ -5,137 +5,176 @@ import (
 	"fmt"
 	"io"
 	"path/filepath"
+	"strconv"
 	"time"
 )
 
 // PrintMultiFileDiff prints a multi-file diff in unified diff format.
+//
+// To avoid holding the whole diff in memory, use WriteMultiFileDiff.
 func PrintMultiFileDiff(ds []*FileDiff) ([]byte, error) {
-	var buf bytes.Buffer
+	return printExact(func(w io.Writer) error { return WriteMultiFileDiff(w, ds) })
+}
+
+// PrintFileDiff prints a FileDiff in unified diff format.
+//
+// To avoid holding the whole diff in memory, use WriteFileDiff.
+func PrintFileDiff(d *FileDiff) ([]byte, error) {
+	return printExact(func(w io.Writer) error { return WriteFileDiff(w, d) })
+}
+
+// PrintHunks prints diff hunks in unified diff format.
+func PrintHunks(hunks []*Hunk) ([]byte, error) {
+	return printExact(func(w io.Writer) error { return (&printer{w: w}).hunks(hunks) })
+}
+
+// printExact calls write twice: first to measure the output, then to fill a
+// buffer of exactly that size. This keeps the peak memory of printing a large
+// diff at the size of the output, instead of the up to 2x that a growing buffer
+// needs while it copies itself into a larger allocation.
+func printExact(write func(io.Writer) error) ([]byte, error) {
+	var n countingWriter
+	if err := write(&n); err != nil {
+		return nil, err
+	}
+	if n == 0 {
+		return nil, nil
+	}
+	buf := bytes.NewBuffer(make([]byte, 0, n))
+	if err := write(buf); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// countingWriter counts the bytes written to it and discards them.
+type countingWriter int
+
+func (c *countingWriter) Write(p []byte) (int, error) {
+	*c += countingWriter(len(p))
+	return len(p), nil
+}
+
+// WriteMultiFileDiff writes a multi-file diff in unified diff format to w.
+//
+// Hunk bodies are written to w directly, without intermediate copies, so many
+// small writes go to w. Wrap w in a bufio.Writer if those writes are costly.
+func WriteMultiFileDiff(w io.Writer, ds []*FileDiff) error {
+	p := printer{w: w}
 	for _, d := range ds {
-		diff, err := PrintFileDiff(d)
-		if err != nil {
-			return nil, err
-		}
-		if _, err := buf.Write(diff); err != nil {
-			return nil, err
+		if err := p.fileDiff(d); err != nil {
+			return err
 		}
 	}
-	return buf.Bytes(), nil
+	return nil
 }
 
-// PrintFileDiff prints a FileDiff in unified diff format.
+// WriteFileDiff writes a FileDiff in unified diff format to w.
+//
+// Hunk bodies are written to w directly, without intermediate copies, so many
+// small writes go to w. Wrap w in a bufio.Writer if those writes are costly.
 //
 // TODO(sqs): handle escaping whitespace/etc. chars in filenames
-func PrintFileDiff(d *FileDiff) ([]byte, error) {
-	var buf bytes.Buffer
+func WriteFileDiff(w io.Writer, d *FileDiff) error {
+	return (&printer{w: w}).fileDiff(d)
+}
+
+// printer writes diffs to w. Header lines are assembled in line, which is
+// reused across lines so that printing does not allocate per line.
+type printer struct {
+	w    io.Writer
+	line []byte
+}
 
+func (p *printer) fileDiff(d *FileDiff) error {
 	for _, xheader := range d.Extended {
-		if _, err := fmt.Fprintln(&buf, xheader); err != nil {
-			return nil, err
+		p.line = append(append(p.line[:0], xheader...), '\n')
+		if err := p.writeLine(); err != nil {
+			return err
 		}
 	}
 
 	// FileDiff is added/deleted file
 	// No further hunks printing needed
 	if d.NewName == "" {
-		_, err := fmt.Fprintf(&buf, onlyInMessage, filepath.Dir(d.OrigName), filepath.Base(d.OrigName))
-		if err != nil {
-			return nil, err
-		}
-		return buf.Bytes(), nil
+		_, err := fmt.Fprintf(p.w, onlyInMessage, filepath.Dir(d.OrigName), filepath.Base(d.OrigName))
+		return err
 	}
 
 	if d.Hunks == nil {
-		return buf.Bytes(), nil
+		return nil
 	}
 
-	if err := printFileHeader(&buf, "--- ", d.OrigName, d.OrigTime); err != nil {
-		return nil, err
-	}
-	if err := printFileHeader(&buf, "+++ ", d.NewName, d.NewTime); err != nil {
-		return nil, err
-	}
-
-	ph, err := PrintHunks(d.Hunks)
-	if err != nil {
-		return nil, err
+	if err := p.fileHeader("--- ", d.OrigName, d.OrigTime); err != nil {
+		return err
 	}
-
-	if _, err := buf.Write(ph); err != nil {
-		return nil, err
+	if err := p.fileHeader("+++ ", d.NewName, d.NewTime); err != nil {
+		return err
 	}
-	return buf.Bytes(), nil
+	return p.hunks(d.Hunks)
 }
 
-func printFileHeader(w io.Writer, prefix string, filename string, timestamp *time.Time) error {
-	if _, err := fmt.Fprint(w, prefix, filename); err != nil {
-		return err
-	}
+func (p *printer) fileHeader(prefix string, filename string, timestamp *time.Time) error {
+	p.line = append(append(p.line[:0], prefix...), filename...)
 	if timestamp != nil {
-		if _, err := fmt.Fprint(w, "\t", timestamp.Format(diffTimeFormatLayout)); err != nil {
-			return err
-		}
+		p.line = timestamp.AppendFormat(append(p.line, '\t'), diffTimeFormatLayout)
 	}
-	if _, err := fmt.Fprintln(w); err != nil {
-		return err
-	}
-	return nil
+	p.line = append(p.line, '\n')
+	return p.writeLine()
 }
 
-// PrintHunks prints diff hunks in unified diff format.
-func PrintHunks(hunks []*Hunk) ([]byte, error) {
-	var buf bytes.Buffer
+func (p *printer) hunks(hunks []*Hunk) error {
 	for _, hunk := range hunks {
-		_, err := fmt.Fprintf(&buf,
-			"@@ -%d,%d +%d,%d @@", hunk.OrigStartLine, hunk.OrigLines, hunk.NewStartLine, hunk.NewLines,
-		)
-		if err != nil {
-			return nil, err
-		}
+		// Equivalent to fmt.Sprintf(hunkHeader, ...) plus the section.
+		p.line = append(p.line[:0], "@@ -"...)
+		p.line = strconv.AppendInt(p.line, int64(hunk.OrigStartLine), 10)
+		p.line = append(p.line, ',')
+		p.line = strconv.AppendInt(p.line, int64(hunk.OrigLines), 10)
+		p.line = append(p.line, " +"...)
+		p.line = strconv.AppendInt(p.line, int64(hunk.NewStartLine), 10)
+		p.line = append(p.line, ',')
+		p.line = strconv.AppendInt(p.line, int64(hunk.NewLines), 10)
+		p.line = append(p.line, " @@"...)
 		if hunk.Section != "" {
-			_, err := fmt.Fprint(&buf, " ", hunk.Section)
-			if err != nil {
-				return nil, err
-			}
+			p.line = append(append(p.line, ' '), hunk.Section...)
 		}
-		if _, err := fmt.Fprintln(&buf); err != nil {
-			return nil, err
+		p.line = append(p.line, '\n')
+		if err := p.writeLine(); err != nil {
+			return err
 		}
 
 		if hunk.OrigNoNewlineAt == 0 {
-			if _, err := buf.Write(hunk.Body); err != nil {
-				return nil, err
+			if _, err := p.w.Write(hunk.Body); err != nil {
+				return err
 			}
 		} else {
-			if _, err := buf.Write(hunk.Body[:hunk.OrigNoNewlineAt]); err != nil {
-				return nil, err
+			if _, err := p.w.Write(hunk.Body[:hunk.OrigNoNewlineAt]); err != nil {
+				return err
 			}
-			if err := printNoNewlineMessage(&buf); err != nil {
-				return nil, err
+			if err := p.noNewlineMessage(); err != nil {
+				return err
 			}
-			if _, err := buf.Write(hunk.Body[hunk.OrigNoNewlineAt:]); err != nil {
-				return nil, err
+			if _, err := p.w.Write(hunk.Body[hunk.OrigNoNewlineAt:]); err != nil {
+				return err
 			}
 		}
 
 		if !bytes.HasSuffix(hunk.Body, []byte{'\n'}) {
-			if _, err := fmt.Fprintln(&buf); err != nil {
-				return nil, err
-			}
-			if err := printNoNewlineMessage(&buf); err != nil {
-				return nil, err
+			p.line = append(append(append(p.line[:0], '\n'), noNewlineMessage...), '\n')
+			if err := p.writeLine(); err != nil {
+				return err
 			}
 		}
 	}
-	return buf.Bytes(), nil
+	return nil
 }
 
-func printNoNewlineMessage(w io.Writer) error {
-	if _, err := w.Write([]byte(noNewlineMessage)); err != nil {
-		return err
-	}
-	if _, err := fmt.Fprintln(w); err != nil {
-		return err
-	}
-	return nil
+func (p *printer) noNewlineMessage() error {
+	p.line = append(append(p.line[:0], noNewlineMessage...), '\n')
+	return p.writeLine()
+}
+
+func (p *printer) writeLine() error {
+	_, err := p.w.Write(p.line)
+	return err
 }