index.FiveLines: fix context lines with 0, 1, 2 previous lines

Also add a test while we’re here.

fixes #119

internal/index/read.go

@@ -436,38 +436,44 @@ func (i *Index) matchesWithBuffer(t Trigram, buffers *bufferPair) ([]Match, erro
 
 // FiveLines returns five \n-separated lines surrounding pos. The first two
 // lines are context above the line containing pos (which is always element
-// [2]), the last two lines are context above that line.
+// [2]), the last two lines are context below that line.
 func FiveLines(b []byte, pos int) [5]string {
-	//fmt.Printf("FiveLines(%q, %d)", string(b), pos)
+	//fmt.Printf("FiveLines(%q, %d)\n", string(b), pos)
 	var five [5]string
 	prev := pos
 	start := 2 // no before lines
-	if prev > 0 {
-		// start of line of match
+
+	// move prev to the beginning of the line in case the match starts in the
+	// middle of the line
+	if idx := bytes.LastIndexByte(b[:prev], '\n'); idx != -1 {
+		prev = idx + 1
+	} else {
+		prev = 0
+	}
+
+	for start > 0 {
+		// fmt.Printf("  looking for newline in %q\n  (would extract from: %q)\n", string(b[:prev]), string(b[prev:]))
 		if idx := bytes.LastIndexByte(b[:prev], '\n'); idx != -1 {
-			prev = idx + 1
-		}
-		// first context line
-		if idx := bytes.LastIndexByte(b[:prev-1], '\n'); idx != -1 {
-			prev = idx + 1
-			start--
-			// second context line (maybe start of file)
-			if idx := bytes.LastIndexByte(b[:prev-1], '\n'); idx != -1 {
+			// fmt.Printf("  new line found at %d, could add one more context line\n", idx)
+			// idx points to the end of the line (\n),
+			// but we want to position prev at the start of the line.
+			if idx := bytes.LastIndexByte(b[:idx], '\n'); idx != -1 {
 				prev = idx + 1
 			} else {
 				prev = 0
 			}
 			start--
 		} else {
-			prev = 0
-			start--
+			break
 		}
+		// fmt.Println()
 	}
+	// fmt.Printf("  will extract from: %q", string(b[prev:]))
 
 	if prev == -1 {
 		return five // TODO: BUG
 	}
-	//fmt.Printf("start=%d, prev=%d, window = %q\n", start, prev, b[prev:])
+	// fmt.Printf("start=%d, prev=%d, window = %q\n", start, prev, b[prev:])
 	scanner := bufio.NewScanner(bytes.NewReader(b[prev:]))
 	for ; start < 5; start++ {
 		if scanner.Scan() {

internal/index/read_test.go

@@ -0,0 +1,151 @@
+package index
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestFiveLines(t *testing.T) {
+	golden := []byte(`first line
+second line
+third line
+foo bar
+baz 234
+qux 567 890<no newline>`)
+
+	binShGolden := []byte(`#!/bin/sh
+set -e
+
+if [ "$(uname -s)" = "Linux" ]; then
+	update-alternatives --install /usr/bin/pager pager /bin/more 50 \
+		--slave /usr/share/man/man1/pager.1.gz pager.1.gz \
+		/usr/share/man/man1/more.1.gz
+fi
+
+#DEBHELPER#
+`)
+
+	for _, tt := range []struct {
+		input []byte
+		query string
+		want  [5]string
+	}{
+
+		{
+			input: golden,
+			query: "foo bar",
+			want: [5]string{
+				"second line",
+				"third line",
+				"foo bar",
+				"baz 234",
+				"qux 567 890<no newline>",
+			},
+		},
+
+		{
+			input: golden,
+			query: "third line",
+			want: [5]string{
+				"first line",
+				"second line",
+				"third line",
+				"foo bar",
+				"baz 234",
+			},
+		},
+
+		{
+			input: golden,
+			query: "qux 567",
+			want: [5]string{
+				"foo bar",
+				"baz 234",
+				"qux 567 890<no newline>",
+				"",
+				"",
+			},
+		},
+
+		{
+			input: []byte("oneline"),
+			query: "one",
+			want: [5]string{
+				"",
+				"",
+				"oneline",
+				"",
+				"",
+			},
+		},
+
+		{
+			input: []byte("oneline\ntwoline"),
+			query: "one",
+			want: [5]string{
+				"",
+				"",
+				"oneline",
+				"twoline",
+				"",
+			},
+		},
+
+		{
+			input: []byte("oneline\ntwoline\nthreeline"),
+			query: "one",
+			want: [5]string{
+				"",
+				"",
+				"oneline",
+				"twoline",
+				"threeline",
+			},
+		},
+
+		{
+			input: []byte("oneline\ntwoline\nthreeline\nfourline\n"),
+			query: "one",
+			want: [5]string{
+				"",
+				"",
+				"oneline",
+				"twoline",
+				"threeline",
+			},
+		},
+
+		{
+			input: []byte(binShGolden),
+			query: "bin/sh",
+			want: [5]string{
+				"",
+				"",
+				"#!/bin/sh",
+				"set -e",
+				"",
+			},
+		},
+
+		{
+			input: []byte(binShGolden),
+			query: "set -e",
+			want: [5]string{
+				"",
+				"#!/bin/sh",
+				"set -e",
+				"",
+				`if [ "$(uname -s)" = "Linux" ]; then`,
+			},
+		},
+	} {
+		t.Run("", func(t *testing.T) {
+			got := FiveLines(tt.input, bytes.Index(tt.input, []byte(tt.query)))
+			if diff := cmp.Diff(tt.want, got); diff != "" {
+				t.Errorf("FiveLines(%s): unexpected diff (-want +got):\n%s", tt.query, diff)
+			}
+		})
+	}
+}