Keep same selection range when quick-starting an interactive rebase (#3247)

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.

6 files changed, 143 insertions, 13 deletions

diff --git a/pkg/gui/context/traits/list_cursor.go b/pkg/gui/context/traits/list_cursor.go
index 368485c05..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 {
@@ -128,7 +143,7 @@ func (self *ListCursor) AreMultipleItemsSelected() bool {
 
 func (self *ListCursor) GetSelectionRange() (int, int) {
 	if self.IsSelectingRange() {
-		return utils.MinMax(self.selectedIdx, self.rangeStartIdx)
+		return utils.SortRange(self.selectedIdx, self.rangeStartIdx)
 	}
 
 	return self.selectedIdx, self.selectedIdx
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.go b/pkg/integration/tests/interactive_rebase/quick_start_keep_selection.go
new file mode 100644
index 000000000..4589d010d
--- /dev/null
+++ b/pkg/integration/tests/interactive_rebase/quick_start_keep_selection.go
@@ -0,0 +1,51 @@
+package interactive_rebase
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var QuickStartKeepSelection = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Starts an interactive rebase and checks that the same commit stays selected",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	GitVersion:   AtLeast("2.38.0"),
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Git.MainBranches = []string{"master"}
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.
+			CreateNCommits(1).
+			NewBranch("branch1").
+			CreateNCommitsStartingAt(3, 2).
+			NewBranch("branch2").
+			CreateNCommitsStartingAt(3, 5)
+
+		shell.SetConfig("rebase.updateRefs", "true")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Commits().
+			Focus().
+			Lines(
+				Contains("CI commit 07").IsSelected(),
+				Contains("CI commit 06"),
+				Contains("CI commit 05"),
+				Contains("CI * commit 04"),
+				Contains("CI commit 03"),
+				Contains("CI commit 02"),
+				Contains("CI commit 01"),
+			).
+			NavigateToLine(Contains("commit 02")).
+			Press(keys.Commits.StartInteractiveRebase).
+			Lines(
+				Contains("pick").Contains("CI commit 07"),
+				Contains("pick").Contains("CI commit 06"),
+				Contains("pick").Contains("CI commit 05"),
+				Contains("update-ref").Contains("branch1"),
+				Contains("pick").Contains("CI * commit 04"),
+				Contains("pick").Contains("CI commit 03"),
+				Contains("CI commit 02").IsSelected(),
+				Contains("CI <-- YOU ARE HERE --- commit 01"),
+			)
+	},
+})
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
new file mode 100644
index 000000000..1b1039cf7
--- /dev/null
+++ b/pkg/integration/tests/interactive_rebase/quick_start_keep_selection_range.go
@@ -0,0 +1,56 @@
+package interactive_rebase
+
+import (
+	"github.com/jesseduffield/lazygit/pkg/config"
+	. "github.com/jesseduffield/lazygit/pkg/integration/components"
+)
+
+var QuickStartKeepSelectionRange = NewIntegrationTest(NewIntegrationTestArgs{
+	Description:  "Starts an interactive rebase and checks that the same commit range stays selected",
+	ExtraCmdArgs: []string{},
+	Skip:         false,
+	GitVersion:   AtLeast("2.38.0"),
+	SetupConfig: func(config *config.AppConfig) {
+		config.GetUserConfig().Git.MainBranches = []string{"master"}
+	},
+	SetupRepo: func(shell *Shell) {
+		shell.
+			CreateNCommits(1).
+			NewBranch("branch1").
+			CreateNCommitsStartingAt(2, 2).
+			NewBranch("branch2").
+			CreateNCommitsStartingAt(2, 4).
+			NewBranch("branch3").
+			CreateNCommitsStartingAt(2, 6)
+
+		shell.SetConfig("rebase.updateRefs", "true")
+	},
+	Run: func(t *TestDriver, keys config.KeybindingConfig) {
+		t.Views().Commits().
+			Focus().
+			NavigateToLine(Contains("commit 04")).
+			Press(keys.Universal.RangeSelectDown).
+			Press(keys.Universal.RangeSelectDown).
+			Lines(
+				Contains("CI commit 07"),
+				Contains("CI commit 06"),
+				Contains("CI * commit 05"),
+				Contains("CI commit 04").IsSelected(),
+				Contains("CI * commit 03").IsSelected(),
+				Contains("CI commit 02").IsSelected(),
+				Contains("CI commit 01"),
+			).
+			Press(keys.Commits.StartInteractiveRebase).
+			Lines(
+				Contains("CI commit 07"),
+				Contains("CI commit 06"),
+				Contains("update-ref").Contains("branch2"),
+				Contains("CI * commit 05"),
+				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"),
+			)
+	},
+})
diff --git a/pkg/integration/tests/test_list.go b/pkg/integration/tests/test_list.go
index 86d7f6682..2406b2f78 100644
--- a/pkg/integration/tests/test_list.go
+++ b/pkg/integration/tests/test_list.go
@@ -172,6 +172,8 @@ var tests = []*components.IntegrationTest{
 	interactive_rebase.OutsideRebaseRangeSelect,
 	interactive_rebase.PickRescheduled,
 	interactive_rebase.QuickStart,
+	interactive_rebase.QuickStartKeepSelection,
+	interactive_rebase.QuickStartKeepSelectionRange,
 	interactive_rebase.Rebase,
 	interactive_rebase.RewordCommitWithEditorAndFail,
 	interactive_rebase.RewordFirstCommit,
diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
index ae876568f..053b7ac2f 100644
--- a/pkg/utils/utils.go
+++ b/pkg/utils/utils.go
@@ -50,7 +50,7 @@ func Max(x, y int) int {
 	return y
 }
 
-func MinMax(x int, y int) (int, int) {
+func SortRange(x int, y int) (int, int) {
 	if x < y {
 		return x, y
 	}