diff --git a/diff/diff_test.go b/diff/diff_test.go index fc42e0b..673f076 100644 --- a/diff/diff_test.go +++ b/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 --git a/diff/print.go b/diff/print.go index 012651a..c0cebcb 100644 --- a/diff/print.go +++ b/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 }