Add a lineReader to provide two-line lookahead

Author: mrnugget

diff/diff_test.go

@@ -74,6 +74,7 @@ func TestParseHunksAndPrintHunks(t *testing.T) {
 		{filename: "oneline_hunk.diff"},
 		{filename: "empty.diff"},
 		{filename: "sample_hunk_lines_start_with_minuses.diff"},
+		{filename: "sample_hunk_lines_start_with_minuses_pluses.diff"},
 	}
 	for _, test := range tests {
 		diffData, err := ioutil.ReadFile(filepath.Join("testdata", test.filename))

diff/parse.go

@@ -23,14 +23,14 @@ func ParseMultiFileDiff(diff []byte) ([]*FileDiff, error) {
 // NewMultiFileDiffReader returns a new MultiFileDiffReader that reads
 // a multi-file unified diff from r.
 func NewMultiFileDiffReader(r io.Reader) *MultiFileDiffReader {
-	return &MultiFileDiffReader{reader: bufio.NewReader(r)}
+	return &MultiFileDiffReader{reader: newLineReader(r)}
 }
 
 // MultiFileDiffReader reads a multi-file unified diff.
 type MultiFileDiffReader struct {
 	line   int
 	offset int64
-	reader *bufio.Reader
+	reader *lineReader
 
 	// TODO(sqs): line and offset tracking in multi-file diffs is broken; add tests and fix
 
@@ -85,7 +85,7 @@ func (r *MultiFileDiffReader) ReadFile() (*FileDiff, error) {
 	// caused by the lack of any hunks, or a malformatted hunk, so we
 	// need to perform the check here.
 	hr := fr.HunksReader()
-	line, err := readLine(r.reader)
+	line, err := r.reader.readLine()
 	if err != nil && err != io.EOF {
 		return fd, err
 	}
@@ -141,14 +141,14 @@ func ParseFileDiff(diff []byte) (*FileDiff, error) {
 // NewFileDiffReader returns a new FileDiffReader that reads a file
 // unified diff.
 func NewFileDiffReader(r io.Reader) *FileDiffReader {
-	return &FileDiffReader{reader: bufio.NewReader(r)}
+	return &FileDiffReader{reader: &lineReader{reader: bufio.NewReader(r)}}
 }
 
 // FileDiffReader reads a unified file diff.
 type FileDiffReader struct {
 	line   int
 	offset int64
-	reader *bufio.Reader
+	reader *lineReader
 
 	// fileHeaderLine is the first file header line, set by:
 	//
@@ -266,7 +266,7 @@ func (r *FileDiffReader) readOneFileHeader(prefix []byte) (filename string, time
 
 	if r.fileHeaderLine == nil {
 		var err error
-		line, err = readLine(r.reader)
+		line, err = r.reader.readLine()
 		if err == io.EOF {
 			return "", nil, &ParseError{r.line, r.offset, ErrNoFileHeader}
 		} else if err != nil {
@@ -318,7 +318,7 @@ func (r *FileDiffReader) ReadExtendedHeaders() ([]string, error) {
 		var line []byte
 		if r.fileHeaderLine == nil {
 			var err error
-			line, err = readLine(r.reader)
+			line, err = r.reader.readLine()
 			if err == io.EOF {
 				return xheaders, &ParseError{r.line, r.offset, ErrExtendedHeadersEOF}
 			} else if err != nil {
@@ -447,15 +447,15 @@ func ParseHunks(diff []byte) ([]*Hunk, error) {
 // NewHunksReader returns a new HunksReader that reads unified diff hunks
 // from r.
 func NewHunksReader(r io.Reader) *HunksReader {
-	return &HunksReader{reader: bufio.NewReader(r)}
+	return &HunksReader{reader: &lineReader{reader: bufio.NewReader(r)}}
 }
 
 // A HunksReader reads hunks from a unified diff.
 type HunksReader struct {
 	line   int
 	offset int64
 	hunk   *Hunk
-	reader *bufio.Reader
+	reader *lineReader
 
 	nextHunkHeaderLine []byte
 }
@@ -474,7 +474,7 @@ func (r *HunksReader) ReadHunk() (*Hunk, error) {
 			line = r.nextHunkHeaderLine
 			r.nextHunkHeaderLine = nil
 		} else {
-			line, err = readLine(r.reader)
+			line, err = r.reader.readLine()
 			if err != nil {
 				if err == io.EOF && r.hunk != nil {
 					return r.hunk, nil
@@ -518,12 +518,15 @@ func (r *HunksReader) ReadHunk() (*Hunk, error) {
 			// If the line starts with `---` and the next one with `+++` we're
 			// looking at a non-extended file header and need to abort.
 			if bytes.HasPrefix(line, []byte("---")) {
-				ok, err := peekPrefix(r.reader, "+++")
+				ok, err := r.reader.nextLineStartsWith("+++")
 				if err != nil {
 					return r.hunk, err
 				}
 				if ok {
-					return r.hunk, &ParseError{r.line, r.offset, &ErrBadHunkLine{Line: line}}
+					ok2, _ := r.reader.nextNextLineStartsWith(string(hunkPrefix))
+					if ok2 {
+						return r.hunk, &ParseError{r.line, r.offset, &ErrBadHunkLine{Line: line}}
+					}
 				}
 			}
 
@@ -593,19 +596,6 @@ func linePrefix(c byte) bool {
 	return false
 }
 
-// peekPrefix peeks into the given reader to check whether the next
-// bytes match the given prefix.
-func peekPrefix(reader *bufio.Reader, prefix string) (bool, error) {
-	next, err := reader.Peek(len(prefix))
-	if err != nil {
-		if err == io.EOF {
-			return false, nil
-		}
-		return false, err
-	}
-	return bytes.HasPrefix(next, []byte(prefix)), nil
-}
-
 // normalizeHeader takes a header of the form:
 // "@@ -linestart[,chunksize] +linestart[,chunksize] @@ section"
 // and returns two strings, with the first in the form:

diff/reader_util.go

@@ -2,9 +2,77 @@ package diff
 
 import (
 	"bufio"
+	"bytes"
 	"io"
 )
 
+func newLineReader(r io.Reader) *lineReader {
+	return &lineReader{reader: bufio.NewReader(r)}
+}
+
+// lineReader is a wrapper around a bufio.Reader that caches the next line to
+// provide lookahead functionality for the next two lines.
+type lineReader struct {
+	reader *bufio.Reader
+
+	cachedNextLine    []byte
+	cachedNextLineErr error
+}
+
+// readLine returns the next unconsumed line and advances the internal cache of
+// the lineReader.
+func (l *lineReader) readLine() ([]byte, error) {
+	if l.cachedNextLine == nil && l.cachedNextLineErr == nil {
+		l.cachedNextLine, l.cachedNextLineErr = readLine(l.reader)
+	}
+
+	if l.cachedNextLineErr != nil {
+		return nil, l.cachedNextLineErr
+	}
+
+	next := l.cachedNextLine
+
+	l.cachedNextLine, l.cachedNextLineErr = readLine(l.reader)
+
+	return next, nil
+}
+
+// nextLineStartsWith looks at the line that would be returned by the next call
+// to readLine to check whether it has the given prefix.
+//
+// io.EOF and bufio.ErrBufferFull errors are ignored so that the function can
+// be used when at the end of the file.
+func (l *lineReader) nextLineStartsWith(prefix string) (bool, error) {
+	return l.lineHasPrefix(l.cachedNextLine, prefix, l.cachedNextLineErr)
+}
+
+// nextNextLineStartsWith checks the prefix of the line *after* the line that
+// would be returned by the next readLine.
+//
+// io.EOF and bufio.ErrBufferFull errors are ignored so that the function can
+// be used when at the end of the file.
+func (l *lineReader) nextNextLineStartsWith(prefix string) (bool, error) {
+	next, err := l.reader.Peek(len(prefix))
+	return l.lineHasPrefix(next, prefix, err)
+}
+
+// lineHasPrefix checks whether the given line has the given prefix with
+// bytes.HasPrefix.
+//
+// The readErr should be the error that was returned when the line was read.
+// lineHasPrefix checks the error to adjust its return value to, e.g., return
+// false and ignore the error when readErr is io.EOF.
+func (l *lineReader) lineHasPrefix(line []byte, prefix string, readErr error) (bool, error) {
+	if readErr != nil {
+		if readErr == io.EOF || readErr == bufio.ErrBufferFull {
+			return false, nil
+		}
+		return false, readErr
+	}
+
+	return bytes.HasPrefix(line, []byte(prefix)), nil
+}
+
 // readLine is a helper that mimics the functionality of calling bufio.Scanner.Scan() and
 // bufio.Scanner.Bytes(), but without the token size limitation. It will read and return
 // the next line in the Reader with the trailing newline stripped. It will return an

diff/testdata/sample_hunk_lines_start_with_minuses_pluses.diff

@@ -0,0 +1,8 @@
+@@ -1,5 +1,5 @@
+ select 1;
+--- this is my query
++++ this is my query
+ select 2;
+ select 3;
+--- this is the last line
++++ this is the last line