allow opening files on the selected line in the staging panel

3 files changed, 116 insertions, 42 deletions

diff --git a/pkg/commands/patch_modifier.go b/pkg/commands/patch_modifier.go
index 092814f90..f4a782015 100644
--- a/pkg/commands/patch_modifier.go
+++ b/pkg/commands/patch_modifier.go
@@ -14,23 +14,42 @@ var hunkHeaderRegexp = regexp.MustCompile(`(?m)^@@ -(\d+)[^\+]+\+(\d+)[^@]+@@(.*
 var patchHeaderRegexp = regexp.MustCompile(`(?ms)(^diff.*?)^@@`)
 
 type PatchHunk struct {
-	header       string
 	FirstLineIdx int
-	LastLineIdx  int
+	oldStart     int
+	newStart     int
+	heading      string
 	bodyLines    []string
 }
 
-func newHunk(header string, body string, firstLineIdx int) *PatchHunk {
-	bodyLines := strings.SplitAfter(header+body, "\n")[1:] // dropping the header line
+func (hunk *PatchHunk) LastLineIdx() int {
+	return hunk.FirstLineIdx + len(hunk.bodyLines)
+}
+
+func newHunk(lines []string, firstLineIdx int) *PatchHunk {
+	header := lines[0]
+	bodyLines := lines[1:]
+
+	oldStart, newStart, heading := headerInfo(header)
 
 	return &PatchHunk{
-		header:       header,
+		oldStart:     oldStart,
+		newStart:     newStart,
+		heading:      heading,
 		FirstLineIdx: firstLineIdx,
-		LastLineIdx:  firstLineIdx + len(bodyLines),
 		bodyLines:    bodyLines,
 	}
 }
 
+func headerInfo(header string) (int, int, string) {
+	match := hunkHeaderRegexp.FindStringSubmatch(header)
+
+	oldStart := mustConvertToInt(match[1])
+	newStart := mustConvertToInt(match[2])
+	heading := match[3]
+
+	return oldStart, newStart, heading
+}
+
 func (hunk *PatchHunk) updatedLines(lineIndices []int, reverse bool) []string {
 	skippedNewlineMessageIndex := -1
 	newLines := []string{}
@@ -94,38 +113,21 @@ func (hunk *PatchHunk) formatWithChanges(lineIndices []int, reverse bool, startO
 }
 
 func (hunk *PatchHunk) updatedHeader(newBodyLines []string, startOffset int, reverse bool) (int, string, bool) {
-	changeCount := 0
-	oldLength := 0
-	newLength := 0
-	for _, line := range newBodyLines {
-		switch line[:1] {
-		case "+":
-			newLength++
-			changeCount++
-		case "-":
-			oldLength++
-			changeCount++
-		case " ":
-			oldLength++
-			newLength++
-		}
-	}
+	changeCount := nLinesWithPrefix(newBodyLines, []string{"+", "-"})
+	oldLength := nLinesWithPrefix(newBodyLines, []string{" ", "-"})
+	newLength := nLinesWithPrefix(newBodyLines, []string{"+", " "})
 
 	if changeCount == 0 {
 		// if nothing has changed we just return nothing
 		return startOffset, "", false
 	}
 
-	// get oldstart, newstart, and heading from header
-	match := hunkHeaderRegexp.FindStringSubmatch(hunk.header)
-
 	var oldStart int
 	if reverse {
-		oldStart = mustConvertToInt(match[2])
+		oldStart = hunk.newStart
 	} else {
-		oldStart = mustConvertToInt(match[1])
+		oldStart = hunk.oldStart
 	}
-	heading := match[3]
 
 	var newStartOffset int
 	// if the hunk went from zero to positive length, we need to increment the starting point by one
@@ -141,7 +143,7 @@ func (hunk *PatchHunk) updatedHeader(newBodyLines []string, startOffset int, rev
 	newStart := oldStart + startOffset + newStartOffset
 
 	newStartOffset = startOffset + newLength - oldLength
-	formattedHeader := hunk.formatHeader(oldStart, oldLength, newStart, newLength, heading)
+	formattedHeader := hunk.formatHeader(oldStart, oldLength, newStart, newLength, hunk.heading)
 	return newStartOffset, formattedHeader, true
 }
 
@@ -162,19 +164,33 @@ func GetHeaderFromDiff(diff string) string {
 }
 
 func GetHunksFromDiff(diff string) []*PatchHunk {
-	headers := hunkHeaderRegexp.FindAllString(diff, -1)
-	bodies := hunkHeaderRegexp.Split(diff, -1)[1:] // discarding top bit
+	hunks := []*PatchHunk{}
+	firstLineIdx := -1
+	var hunkLines []string
+	pastDiffHeader := false
+
+	for lineIdx, line := range strings.SplitAfter(diff, "\n") {
+		isHunkHeader := strings.HasPrefix(line, "@@ -")
+
+		if isHunkHeader {
+			if pastDiffHeader { // we need to persist the current hunk
+				hunks = append(hunks, newHunk(hunkLines, firstLineIdx))
+			}
+			pastDiffHeader = true
+			firstLineIdx = lineIdx
+			hunkLines = []string{line}
+			continue
+		}
 
-	headerFirstLineIndices := []int{}
-	for lineIdx, line := range strings.Split(diff, "\n") {
-		if strings.HasPrefix(line, "@@ -") {
-			headerFirstLineIndices = append(headerFirstLineIndices, lineIdx)
+		if !pastDiffHeader { // skip through the stuff that precedes the first hunk
+			continue
 		}
+
+		hunkLines = append(hunkLines, line)
 	}
 
-	hunks := make([]*PatchHunk, len(headers))
-	for index, header := range headers {
-		hunks[index] = newHunk(header, bodies[index], headerFirstLineIndices[index])
+	if pastDiffHeader {
+		hunks = append(hunks, newHunk(hunkLines, firstLineIdx))
 	}
 
 	return hunks
@@ -203,7 +219,7 @@ outer:
 	for _, hunk := range d.hunks {
 		// if there is any line in our lineIndices array that the hunk contains, we append it
 		for _, lineIdx := range lineIndices {
-			if lineIdx >= hunk.FirstLineIdx && lineIdx <= hunk.LastLineIdx {
+			if lineIdx >= hunk.FirstLineIdx && lineIdx <= hunk.LastLineIdx() {
 				hunksInRange = append(hunksInRange, hunk)
 				continue outer
 			}
@@ -251,7 +267,7 @@ func (d *PatchModifier) OriginalPatchLength() int {
 		return 0
 	}
 
-	return d.hunks[len(d.hunks)-1].LastLineIdx
+	return d.hunks[len(d.hunks)-1].LastLineIdx()
 }
 
 func ModifiedPatchForRange(log *logrus.Entry, filename string, diffText string, firstLineIdx int, lastLineIdx int, reverse bool, keepOriginalHeader bool) string {
@@ -263,3 +279,24 @@ func ModifiedPatchForLines(log *logrus.Entry, filename string, diffText string,
 	p := NewPatchModifier(log, filename, diffText)
 	return p.ModifiedPatchForLines(includedLineIndices, reverse, keepOriginalHeader)
 }
+
+// I want to know, given a hunk, what line a given index is on
+func (hunk *PatchHunk) LineNumberOfLine(idx int) int {
+	lines := hunk.bodyLines[0 : idx-hunk.FirstLineIdx-1]
+
+	offset := nLinesWithPrefix(lines, []string{"+", " "})
+
+	return hunk.newStart + offset
+}
+
+func nLinesWithPrefix(lines []string, chars []string) int {
+	result := 0
+	for _, line := range lines {
+		for _, char := range chars {
+			if line[:1] == char {
+				result++
+			}
+		}
+	}
+	return result
+}
diff --git a/pkg/commands/patch_modifier_test.go b/pkg/commands/patch_modifier_test.go
index 0e44af812..66b675923 100644
--- a/pkg/commands/patch_modifier_test.go
+++ b/pkg/commands/patch_modifier_test.go
@@ -2,6 +2,7 @@ package commands
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -88,6 +89,15 @@ index e69de29..c6568ea 100644
 \ No newline at end of file
 `
 
+const exampleHunk = `@@ -1,5 +1,5 @@
+ apple
+-grape
++orange
+...
+...
+...
+`
+
 // TestModifyPatchForRange is a function.
 func TestModifyPatchForRange(t *testing.T) {
 	type scenario struct {
@@ -509,3 +519,30 @@ func TestModifyPatchForRange(t *testing.T) {
 		})
 	}
 }
+
+func TestLineNumberOfLine(t *testing.T) {
+	type scenario struct {
+		testName string
+		hunk     *PatchHunk
+		idx      int
+		expected int
+	}
+
+	scenarios := []scenario{
+		{
+			testName: "nothing selected",
+			hunk:     newHunk(strings.SplitAfter(exampleHunk, "\n"), 10),
+			idx:      15,
+			expected: 3,
+		},
+	}
+
+	for _, s := range scenarios {
+		t.Run(s.testName, func(t *testing.T) {
+			result := s.hunk.LineNumberOfLine(s.idx)
+			if !assert.Equal(t, s.expected, result) {
+				fmt.Println(result)
+			}
+		})
+	}
+}
diff --git a/pkg/commands/patch_parser.go b/pkg/commands/patch_parser.go
index d35aebb24..b54c57c1f 100644
--- a/pkg/commands/patch_parser.go
+++ b/pkg/commands/patch_parser.go
@@ -63,7 +63,7 @@ func (p *PatchParser) GetHunkContainingLine(lineIndex int, offset int) *PatchHun
 	}
 
 	for index, hunk := range p.PatchHunks {
-		if lineIndex >= hunk.FirstLineIdx && lineIndex <= hunk.LastLineIdx {
+		if lineIndex >= hunk.FirstLineIdx && lineIndex <= hunk.LastLineIdx() {
 			resultIndex := index + offset
 			if resultIndex < 0 {
 				resultIndex = 0
@@ -75,7 +75,7 @@ func (p *PatchParser) GetHunkContainingLine(lineIndex int, offset int) *PatchHun
 	}
 
 	// if your cursor is past the last hunk, select the last hunk
-	if lineIndex > p.PatchHunks[len(p.PatchHunks)-1].LastLineIdx {
+	if lineIndex > p.PatchHunks[len(p.PatchHunks)-1].LastLineIdx() {
 		return p.PatchHunks[len(p.PatchHunks)-1]
 	}