Keep same selection range when quick-starting an interactive rebase

This is useful if you want to move a range of commits, so you select them, and
then realize it's better to do it in an interactive rebase. Pressing 'i'
preserves the range now.

3 files changed, 35 insertions, 15 deletions

diff --git a/pkg/gui/context/traits/list_cursor.go b/pkg/gui/context/traits/list_cursor.go
index 4e74019ca..e42f1f5e5 100644
--- a/pkg/gui/context/traits/list_cursor.go
+++ b/pkg/gui/context/traits/list_cursor.go
@@ -64,6 +64,21 @@ func (self *ListCursor) SetSelection(value int) {
 	self.CancelRangeSelect()
 }
 
+func (self *ListCursor) SetSelectionRangeAndMode(selectedIdx, rangeStartIdx int, mode RangeSelectMode) {
+	self.selectedIdx = self.clampValue(selectedIdx)
+	self.rangeStartIdx = self.clampValue(rangeStartIdx)
+	self.rangeSelectMode = mode
+}
+
+// Returns the selectedIdx, the rangeStartIdx, and the mode of the current selection.
+func (self *ListCursor) GetSelectionRangeAndMode() (int, int, RangeSelectMode) {
+	if self.IsSelectingRange() {
+		return self.selectedIdx, self.rangeStartIdx, self.rangeSelectMode
+	} else {
+		return self.selectedIdx, self.selectedIdx, self.rangeSelectMode
+	}
+}
+
 func (self *ListCursor) clampValue(value int) int {
 	clampedValue := -1
 	if self.list.Len() > 0 {
diff --git a/pkg/gui/controllers/local_commits_controller.go b/pkg/gui/controllers/local_commits_controller.go
index 2062907f0..1feb6e488 100644
--- a/pkg/gui/controllers/local_commits_controller.go
+++ b/pkg/gui/controllers/local_commits_controller.go
@@ -116,7 +116,7 @@ func (self *LocalCommitsController) GetKeybindings(opts types.KeybindingsOpts) [
 			// we're calling it 'quick-start interactive rebase' to differentiate it from
 			// when you manually select the base commit.
 			Key:               opts.GetKey(opts.Config.Commits.StartInteractiveRebase),
-			Handler:           self.withItem(self.quickStartInteractiveRebase),
+			Handler:           self.quickStartInteractiveRebase,
 			GetDisabledReason: self.require(self.notMidRebase(self.c.Tr.AlreadyRebasing), self.canFindCommitForQuickStart),
 			Description:       self.c.Tr.QuickStartInteractiveRebase,
 			Tooltip: utils.ResolvePlaceholderString(self.c.Tr.QuickStartInteractiveRebaseTooltip, map[string]string{
@@ -429,36 +429,42 @@ func (self *LocalCommitsController) edit(selectedCommits []*models.Commit) error
 
 	selectedCommit := selectedCommits[0]
 
-	return self.startInteractiveRebaseWithEdit(selectedCommit, selectedCommit)
+	return self.startInteractiveRebaseWithEdit(selectedCommit)
 }
 
-func (self *LocalCommitsController) quickStartInteractiveRebase(selectedCommit *models.Commit) error {
+func (self *LocalCommitsController) quickStartInteractiveRebase() error {
 	commitToEdit, err := self.findCommitForQuickStartInteractiveRebase()
 	if err != nil {
 		return self.c.Error(err)
 	}
 
-	return self.startInteractiveRebaseWithEdit(commitToEdit, selectedCommit)
+	return self.startInteractiveRebaseWithEdit(commitToEdit)
 }
 
 func (self *LocalCommitsController) startInteractiveRebaseWithEdit(
 	commitToEdit *models.Commit,
-	selectedCommit *models.Commit,
 ) error {
 	return self.c.WithWaitingStatus(self.c.Tr.RebasingStatus, func(gocui.Task) error {
 		self.c.LogAction(self.c.Tr.Actions.EditCommit)
+		selectedIdx, rangeStartIdx, rangeSelectMode := self.context().GetSelectionRangeAndMode()
+		commits := self.c.Model().Commits
+		selectedSha := commits[selectedIdx].Sha
+		rangeStartSha := commits[rangeStartIdx].Sha
 		err := self.c.Git().Rebase.EditRebase(commitToEdit.Sha)
 		return self.c.Helpers().MergeAndRebase.CheckMergeOrRebaseWithRefreshOptions(
 			err,
 			types.RefreshOptions{Mode: types.BLOCK_UI, Then: func() {
-				// We need to select the same commit again because after starting a rebase,
+				// We need to select the same commit range again because after starting a rebase,
 				// new lines can be added for update-ref commands in the TODO file, due to
-				// stacked branches. So the commit may be in a different position in the list.
-				_, index, ok := lo.FindIndexOf(self.c.Model().Commits, func(c *models.Commit) bool {
-					return c.Sha == selectedCommit.Sha
+				// stacked branches. So the selected commits may be in different positions in the list.
+				_, newSelectedIdx, ok1 := lo.FindIndexOf(self.c.Model().Commits, func(c *models.Commit) bool {
+					return c.Sha == selectedSha
+				})
+				_, newRangeStartIdx, ok2 := lo.FindIndexOf(self.c.Model().Commits, func(c *models.Commit) bool {
+					return c.Sha == rangeStartSha
 				})
-				if ok {
-					self.context().SetSelection(index)
+				if ok1 && ok2 {
+					self.context().SetSelectionRangeAndMode(newSelectedIdx, newRangeStartIdx, rangeSelectMode)
 				}
 			}})
 	})
diff --git a/pkg/integration/tests/interactive_rebase/quick_start_keep_selection_range.go b/pkg/integration/tests/interactive_rebase/quick_start_keep_selection_range.go
index 1109e7914..1b1039cf7 100644
--- a/pkg/integration/tests/interactive_rebase/quick_start_keep_selection_range.go
+++ b/pkg/integration/tests/interactive_rebase/quick_start_keep_selection_range.go
@@ -46,10 +46,9 @@ var QuickStartKeepSelectionRange = NewIntegrationTest(NewIntegrationTestArgs{
 				Contains("CI commit 06"),
 				Contains("update-ref").Contains("branch2"),
 				Contains("CI * commit 05"),
-				// Only 01 remains selected, but we want 04 through 01 to stay selected:
-				Contains("CI commit 04"),
-				Contains("update-ref").Contains("branch1"),
-				Contains("CI * commit 03"),
+				Contains("CI commit 04").IsSelected(),
+				Contains("update-ref").Contains("branch1").IsSelected(),
+				Contains("CI * commit 03").IsSelected(),
 				Contains("CI commit 02").IsSelected(),
 				Contains("CI <-- YOU ARE HERE --- commit 01"),
 			)