Move to next stageable line after adding a line to a custom patch

While it's true that the behavior is a little different from the staging panel,
where the staged lines are actually removed from the view and in many cases the
selection stays more or less in the same place, it is still very useful to move
to the next stageable thing in the custom patch building view too.

Author: stefanhaller

pkg/commands/patch/patch.go

@@ -115,29 +115,41 @@ func (self *Patch) HunkContainingLine(idx int) int {
 	return -1
 }
 
-// Returns the patch line index of the next change (i.e. addition or deletion).
-func (self *Patch) GetNextChangeIdx(idx int) int {
+// Returns the patch line index of the next change (i.e. addition or deletion)
+// that matches the same "included" state, given the includedLines. If you don't
+// care about included states, pass nil for includedLines and false for included.
+func (self *Patch) GetNextChangeIdxOfSameIncludedState(idx int, includedLines []int, included bool) (int, bool) {
 	idx = lo.Clamp(idx, 0, self.LineCount()-1)
 
 	lines := self.Lines()
 
+	isMatch := func(i int, line *PatchLine) bool {
+		sameIncludedState := lo.Contains(includedLines, i) == included
+		return line.IsChange() && sameIncludedState
+	}
+
 	for i, line := range lines[idx:] {
-		if line.isChange() {
-			return i + idx
+		if isMatch(i+idx, line) {
+			return i + idx, true
 		}
 	}
 
 	// there are no changes from the cursor onwards so we'll instead
 	// return the index of the last change
 	for i := len(lines) - 1; i >= 0; i-- {
 		line := lines[i]
-		if line.isChange() {
-			return i
+		if isMatch(i, line) {
+			return i, true
 		}
 	}
 
-	// should not be possible
-	return 0
+	return 0, false
+}
+
+// Returns the patch line index of the next change (i.e. addition or deletion).
+func (self *Patch) GetNextChangeIdx(idx int) int {
+	result, _ := self.GetNextChangeIdxOfSameIncludedState(idx, nil, false)
+	return result
 }
 
 // Returns the length of the patch in lines

pkg/gui/controllers/patch_building_controller.go

@@ -161,6 +161,8 @@ func (self *PatchBuildingController) toggleSelection() error {
 		state.SetLineSelectMode()
 	}
 
+	state.SelectNextStageableLineOfSameIncludedState(self.context().GetIncludedLineIndices(), firstSelectedChangeLineIsStaged)
+
 	return nil
 }
 

pkg/gui/patch_exploring/state.go

@@ -338,3 +338,11 @@ func wrapPatchLines(diff string, view *gocui.View) ([]int, []int) {
 		view.Wrap, view.Editable, strings.TrimSuffix(diff, "\n"), view.InnerWidth(), view.TabWidth)
 	return viewLineIndices, patchLineIndices
 }
+
+func (s *State) SelectNextStageableLineOfSameIncludedState(includedLines []int, included bool) {
+	_, lastLineIdx := s.SelectedPatchRange()
+	patchLineIdx, found := s.patch.GetNextChangeIdxOfSameIncludedState(lastLineIdx+1, includedLines, included)
+	if found {
+		s.SelectLine(s.viewLineIndices[patchLineIdx])
+	}
+}

pkg/integration/tests/patch_building/move_to_index_partial.go

@@ -48,7 +48,6 @@ var MoveToIndexPartial = NewIntegrationTest(NewIntegrationTestArgs{
 				Contains(`+third line2`),
 			).
 			PressPrimaryAction().
-			SelectNextItem().
 			PressPrimaryAction().
 			Tap(func() {
 				t.Views().Information().Content(Contains("Building patch"))

pkg/integration/tests/patch_building/specific_selection.go

@@ -66,18 +66,20 @@ var SpecificSelection = NewIntegrationTest(NewIntegrationTestArgs{
 				Contains(` 1f`),
 			).
 			PressPrimaryAction().
-			// unlike in the staging panel, we don't remove lines from the patch building panel
-			// upon 'adding' them. So the same lines will be selected
 			SelectedLines(
-				Contains(`@@ -1,6 +1,6 @@`),
-				Contains(`-1a`),
-				Contains(`+aa`),
-				Contains(` 1b`),
-				Contains(`-1c`),
-				Contains(`+cc`),
-				Contains(` 1d`),
-				Contains(` 1e`),
-				Contains(` 1f`),
+				Contains(`@@ -17,9 +17,9 @@`),
+				Contains(` 1q`),
+				Contains(` 1r`),
+				Contains(` 1s`),
+				Contains(`-1t`),
+				Contains(`-1u`),
+				Contains(`-1v`),
+				Contains(`+tt`),
+				Contains(`+uu`),
+				Contains(`+vv`),
+				Contains(` 1w`),
+				Contains(` 1x`),
+				Contains(` 1y`),
 			).
 			Tap(func() {
 				t.Views().Information().Content(Contains("Building patch"))
@@ -106,12 +108,21 @@ var SpecificSelection = NewIntegrationTest(NewIntegrationTestArgs{
 				Contains("+2a"),
 			).
 			PressPrimaryAction().
+			SelectedLines(
+				Contains("+2b"),
+			).
 			NavigateToLine(Contains("+2c")).
 			Press(keys.Universal.ToggleRangeSelect).
 			NavigateToLine(Contains("+2e")).
 			PressPrimaryAction().
+			SelectedLines(
+				Contains("+2f"),
+			).
 			NavigateToLine(Contains("+2g")).
 			PressPrimaryAction().
+			SelectedLines(
+				Contains("+2h"),
+			).
 			Tap(func() {
 				t.Views().Information().Content(Contains("Building patch"))