Fix selection after staging an added line

In some cases, staging an added line could result in the previous deleted line
to become selected, rather than the next added line.

pkg/commands/patch/patch.go

@@ -51,6 +51,16 @@ func (self *Patch) Lines() []*PatchLine {
 	return lines
 }
 
+// Returns the old-file starting line number of the hunk containing the given
+// patch line index. Returns 0 if the line is not inside any hunk.
+func (self *Patch) HunkOldStartForLine(idx int) int {
+	hunkIdx := self.HunkContainingLine(idx)
+	if hunkIdx == -1 {
+		return 0
+	}
+	return self.hunks[hunkIdx].oldStart
+}
+
 // Returns the patch line index of the first line in the given hunk
 func (self *Patch) HunkStartIdx(hunkIndex int) int {
 	hunkIndex = lo.Clamp(hunkIndex, 0, len(self.hunks)-1)

pkg/commands/patch/patch_line.go

@@ -22,6 +22,14 @@ func (self *PatchLine) IsChange() bool {
 	return self.Kind == ADDITION || self.Kind == DELETION
 }
 
+func (self *PatchLine) IsAddition() bool {
+	return self.Kind == ADDITION
+}
+
+func (self *PatchLine) IsDeletion() bool {
+	return self.Kind == DELETION
+}
+
 // Returns the number of lines in the given slice that have one of the given kinds
 func nLinesWithKind(lines []*PatchLine, kinds []PatchLineKind) int {
 	return lo.CountBy(lines, func(line *PatchLine) bool {

pkg/gui/patch_exploring/state.go

@@ -89,7 +89,24 @@ func NewState(diff string, selectedLineIdx int, view *gocui.View, oldState *Stat
 		if oldState.selectMode != RANGE {
 			selectMode = oldState.selectMode
 		}
-		selectedLineIdx = viewLineIndices[patch.GetNextChangeIdx(oldState.patchLineIndices[oldState.selectedLineIdx])]
+		oldPatchLineIdx := oldState.patchLineIndices[oldState.selectedLineIdx]
+		newPatchLineIdx := patch.GetNextChangeIdx(oldPatchLineIdx)
+		// When staging an addition from a consecutive changes block, the unselected deletions get
+		// reordered to appear before the remaining additions in the new diff. This can cause the
+		// cursor to land on a deletion at the same patch line index where the staged addition used
+		// to be. In that case, skip forward past any deletions, then call GetNextChangeIdx from the
+		// first non-deletion position, which correctly lands on the next meaningful change.
+		newLines := patch.Lines()
+		if newPatchLineIdx == oldPatchLineIdx &&
+			oldState.patch.Lines()[oldPatchLineIdx].IsAddition() &&
+			newLines[newPatchLineIdx].IsDeletion() &&
+			patch.HunkOldStartForLine(newPatchLineIdx) == oldState.patch.HunkOldStartForLine(oldPatchLineIdx) {
+			for newPatchLineIdx < len(newLines) && newLines[newPatchLineIdx].IsDeletion() {
+				newPatchLineIdx++
+			}
+			newPatchLineIdx = patch.GetNextChangeIdx(newPatchLineIdx)
+		}
+		selectedLineIdx = viewLineIndices[newPatchLineIdx]
 	} else {
 		selectedLineIdx = viewLineIndices[patch.GetNextChangeIdx(0)]
 	}

pkg/integration/tests/staging/select_next_line_after_staging_in_two_hunk_diff.go

@@ -54,15 +54,8 @@ var SelectNextLineAfterStagingInTwoHunkDiff = NewIntegrationTest(NewIntegrationT
 			PressPrimaryAction().
 			SelectedLine(Contains("+1b")).
 			PressPrimaryAction().
-			/* EXPECTED:
 			SelectedLine(Contains("+2b")).
-			ACTUAL: */
-			SelectedLine(Contains("-1")).
-			NavigateToLine(Contains("+2b")).
 			PressPrimaryAction().
-			/* EXPECTED:
 			SelectedLine(Contains("-3"))
-			ACTUAL: */
-			SelectedLine(Contains("-1"))
 	},
 })