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

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,7 +161,7 @@ func (self *PatchBuildingController) toggleSelection() error {
 state.SetLineSelectMode()
 }
 
-state.SelectNextStageableLine()
+state.SelectNextStageableLineOfSameIncludedState(self.context().GetIncludedLineIndices(), firstSelectedChangeLineIsStaged)
 
 return nil
 }

pkg/gui/patch_exploring/state.go

@@ -339,8 +339,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)
-s.SelectLine(s.viewLineIndices[patchLineIdx])
+patchLineIdx, found := s.patch.GetNextChangeIdxOfSameIncludedState(lastLineIdx+1, includedLines, included)
+if found {
+s.SelectLine(s.viewLineIndices[patchLineIdx])
+}
 }