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

pkg/commands/patch/patch.go

@@ -115,15 +115,22 @@ 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
 		}
 	}
 
@@ -131,13 +138,18 @@ func (self *Patch) GetNextChangeIdx(idx int) int {
 	// 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

@@ -155,7 +155,7 @@ func (self *PatchBuildingController) toggleSelection() error {
 		state.SetLineSelectMode()
 	}
 
-	state.SelectNextStageableLine()
+	state.SelectNextStageableLineOfSameIncludedState(self.context().GetIncludedLineIndices(), currentLineIsStaged)
 
 	return nil
 }

pkg/gui/patch_exploring/state.go

@@ -325,8 +325,10 @@ func wrapPatchLines(diff string, view *gocui.View) ([]int, []int) {
 	return viewLineIndices, patchLineIndices
 }
 
-func (s *State) SelectNextStageableLine() {
+func (s *State) SelectNextStageableLineOfSameIncludedState(includedLines []int, included bool) {
 	_, lastLineIdx := s.SelectedPatchRange()
-	patchLineIdx := s.patch.GetNextChangeIdx(lastLineIdx + 1)
+	patchLineIdx, found := s.patch.GetNextChangeIdxOfSameIncludedState(lastLineIdx+1, includedLines, included)
+	if found {
 		s.SelectLine(s.viewLineIndices[patchLineIdx])
+	}
 }