Keep the same line selected after squashing fixup commits

This uses a bit of a heuristic that is hopefully correct most of the time.

4 files changed, 210 insertions, 18 deletions

diff --git a/pkg/gui/controllers/local_commits_controller.go b/pkg/gui/controllers/local_commits_controller.go
index e6b8bc801..ef6b5be80 100644
--- a/pkg/gui/controllers/local_commits_controller.go
+++ b/pkg/gui/controllers/local_commits_controller.go
@@ -2,6 +2,7 @@ package controllers
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/fsmiamoto/git-todo-parser/todo"
 	"github.com/go-errors/errors"
@@ -847,37 +848,89 @@ func (self *LocalCommitsController) squashFixupCommits() error {
 }
 
 func (self *LocalCommitsController) squashAllFixupsAboveSelectedCommit(commit *models.Commit) error {
-	return self.squashFixupsImpl(commit)
+	return self.squashFixupsImpl(commit, self.context().GetSelectedLineIdx())
 }
 
 func (self *LocalCommitsController) squashAllFixupsInCurrentBranch() error {
-	commit, err := self.findCommitForSquashFixupsInCurrentBranch()
+	commit, rebaseStartIdx, err := self.findCommitForSquashFixupsInCurrentBranch()
 	if err != nil {
 		return self.c.Error(err)
 	}
 
-	return self.squashFixupsImpl(commit)
+	return self.squashFixupsImpl(commit, rebaseStartIdx)
 }
 
-func (self *LocalCommitsController) squashFixupsImpl(commit *models.Commit) error {
-	return self.c.WithWaitingStatus(self.c.Tr.SquashingStatus, func(gocui.Task) error {
+func (self *LocalCommitsController) squashFixupsImpl(commit *models.Commit, rebaseStartIdx int) error {
+	selectionOffset := countSquashableCommitsAbove(self.c.Model().Commits, self.context().GetSelectedLineIdx(), rebaseStartIdx)
+	return self.c.WithWaitingStatusSync(self.c.Tr.SquashingStatus, func() error {
 		self.c.LogAction(self.c.Tr.Actions.SquashAllAboveFixupCommits)
 		err := self.c.Git().Rebase.SquashAllAboveFixupCommits(commit)
-		return self.c.Helpers().MergeAndRebase.CheckMergeOrRebase(err)
+		self.context().MoveSelectedLine(-selectionOffset)
+		return self.c.Helpers().MergeAndRebase.CheckMergeOrRebaseWithRefreshOptions(
+			err, types.RefreshOptions{Mode: types.SYNC})
 	})
 }
 
-func (self *LocalCommitsController) findCommitForSquashFixupsInCurrentBranch() (*models.Commit, error) {
+func (self *LocalCommitsController) findCommitForSquashFixupsInCurrentBranch() (*models.Commit, int, error) {
 	commits := self.c.Model().Commits
 	_, index, ok := lo.FindIndexOf(commits, func(c *models.Commit) bool {
 		return c.IsMerge() || c.Status == models.StatusMerged
 	})
 
 	if !ok || index == 0 {
-		return nil, errors.New(self.c.Tr.CannotSquashCommitsInCurrentBranch)
+		return nil, -1, errors.New(self.c.Tr.CannotSquashCommitsInCurrentBranch)
+	}
+
+	return commits[index-1], index - 1, nil
+}
+
+// Anticipate how many commits above the selectedIdx are going to get squashed
+// by the SquashAllAboveFixupCommits call, so that we can adjust the selection
+// afterwards. Let's hope we're matching git's behavior correctly here.
+func countSquashableCommitsAbove(commits []*models.Commit, selectedIdx int, rebaseStartIdx int) int {
+	result := 0
+
+	// For each commit _above_ the selection, ...
+	for i, commit := range commits[0:selectedIdx] {
+		// ... see if it is a fixup commit, and get the base subject it applies to
+		if baseSubject, isFixup := isFixupCommit(commit.Name); isFixup {
+			// Then, for each commit after the fixup, up to and including the
+			// rebase start commit, see if we find the base commit
+			for _, baseCommit := range commits[i+1 : rebaseStartIdx+1] {
+				if strings.HasPrefix(baseCommit.Name, baseSubject) {
+					result++
+				}
+			}
+		}
+	}
+	return result
+}
+
+// Check whether the given subject line is the subject of a fixup commit, and
+// returns (trimmedSubject, true) if so (where trimmedSubject is the subject
+// with all fixup prefixes removed), or (subject, false) if not.
+func isFixupCommit(subject string) (string, bool) {
+	prefixes := []string{"fixup! ", "squash! ", "amend! "}
+	trimPrefix := func(s string) (string, bool) {
+		for _, prefix := range prefixes {
+			if strings.HasPrefix(s, prefix) {
+				return strings.TrimPrefix(s, prefix), true
+			}
+		}
+		return s, false
+	}
+
+	if subject, wasTrimmed := trimPrefix(subject); wasTrimmed {
+		for {
+			// handle repeated prefixes like "fixup! amend! fixup! Subject"
+			if subject, wasTrimmed = trimPrefix(subject); !wasTrimmed {
+				break
+			}
+		}
+		return subject, true
 	}
 
-	return commits[index-1], nil
+	return subject, false
 }
 
 func (self *LocalCommitsController) createTag(commit *models.Commit) error {
@@ -1070,7 +1123,7 @@ func (self *LocalCommitsController) canFindCommitForQuickStart() *types.Disabled
 }
 
 func (self *LocalCommitsController) canFindCommitForSquashFixupsInCurrentBranch() *types.DisabledReason {
-	if _, err := self.findCommitForSquashFixupsInCurrentBranch(); err != nil {
+	if _, _, err := self.findCommitForSquashFixupsInCurrentBranch(); err != nil {
 		return &types.DisabledReason{Text: err.Error()}
 	}
 
diff --git a/pkg/gui/controllers/local_commits_controller_test.go b/pkg/gui/controllers/local_commits_controller_test.go
new file mode 100644
index 000000000..d425821d7
--- /dev/null
+++ b/pkg/gui/controllers/local_commits_controller_test.go
@@ -0,0 +1,141 @@
+package controllers
+
+import (
+	"testing"
+
+	"github.com/jesseduffield/lazygit/pkg/commands/models"
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_countSquashableCommitsAbove(t *testing.T) {
+	scenarios := []struct {
+		name           string
+		commits        []*models.Commit
+		selectedIdx    int
+		rebaseStartIdx int
+		expectedResult int
+	}{
+		{
+			name: "no squashable commits",
+			commits: []*models.Commit{
+				{Name: "abc"},
+				{Name: "def"},
+				{Name: "ghi"},
+			},
+			selectedIdx:    2,
+			rebaseStartIdx: 2,
+			expectedResult: 0,
+		},
+		{
+			name: "some squashable commits, including for the selected commit",
+			commits: []*models.Commit{
+				{Name: "fixup! def"},
+				{Name: "fixup! ghi"},
+				{Name: "abc"},
+				{Name: "def"},
+				{Name: "ghi"},
+			},
+			selectedIdx:    4,
+			rebaseStartIdx: 4,
+			expectedResult: 2,
+		},
+		{
+			name: "base commit is below rebase start",
+			commits: []*models.Commit{
+				{Name: "fixup! def"},
+				{Name: "abc"},
+				{Name: "def"},
+			},
+			selectedIdx:    1,
+			rebaseStartIdx: 1,
+			expectedResult: 0,
+		},
+		{
+			name: "base commit does not exist at all",
+			commits: []*models.Commit{
+				{Name: "fixup! xyz"},
+				{Name: "abc"},
+				{Name: "def"},
+			},
+			selectedIdx:    2,
+			rebaseStartIdx: 2,
+			expectedResult: 0,
+		},
+		{
+			name: "selected commit is in the middle of fixups",
+			commits: []*models.Commit{
+				{Name: "fixup! def"},
+				{Name: "abc"},
+				{Name: "fixup! ghi"},
+				{Name: "def"},
+				{Name: "ghi"},
+			},
+			selectedIdx:    1,
+			rebaseStartIdx: 4,
+			expectedResult: 1,
+		},
+		{
+			name: "selected commit is after rebase start",
+			commits: []*models.Commit{
+				{Name: "fixup! def"},
+				{Name: "abc"},
+				{Name: "def"},
+				{Name: "ghi"},
+			},
+			selectedIdx:    3,
+			rebaseStartIdx: 2,
+			expectedResult: 1,
+		},
+	}
+	for _, s := range scenarios {
+		t.Run(s.name, func(t *testing.T) {
+			assert.Equal(t, s.expectedResult, countSquashableCommitsAbove(s.commits, s.selectedIdx, s.rebaseStartIdx))
+		})
+	}
+}
+
+func Test_isFixupCommit(t *testing.T) {
+	scenarios := []struct {
+		subject                string
+		expectedTrimmedSubject string
+		expectedIsFixup        bool
+	}{
+		{
+			subject:                "Bla",
+			expectedTrimmedSubject: "Bla",
+			expectedIsFixup:        false,
+		},
+		{
+			subject:                "fixup Bla",
+			expectedTrimmedSubject: "fixup Bla",
+			expectedIsFixup:        false,
+		},
+		{
+			subject:                "fixup! Bla",
+			expectedTrimmedSubject: "Bla",
+			expectedIsFixup:        true,
+		},
+		{
+			subject:                "fixup! fixup! Bla",
+			expectedTrimmedSubject: "Bla",
+			expectedIsFixup:        true,
+		},
+		{
+			subject:                "amend! squash! Bla",
+			expectedTrimmedSubject: "Bla",
+			expectedIsFixup:        true,
+		},
+		{
+			subject:                "fixup!",
+			expectedTrimmedSubject: "fixup!",
+			expectedIsFixup:        false,
+		},
+	}
+	for _, s := range scenarios {
+		t.Run(s.subject, func(t *testing.T) {
+			trimmedSubject, isFixupCommit := isFixupCommit(s.subject)
+			assert.Equal(t, s.expectedTrimmedSubject, trimmedSubject)
+			assert.Equal(t, s.expectedIsFixup, isFixupCommit)
+		})
+	}
+}
diff --git a/pkg/integration/tests/interactive_rebase/squash_fixups_above.go b/pkg/integration/tests/interactive_rebase/squash_fixups_above.go
index b412f93ed..e87addce0 100644
--- a/pkg/integration/tests/interactive_rebase/squash_fixups_above.go
+++ b/pkg/integration/tests/interactive_rebase/squash_fixups_above.go
@@ -46,10 +46,9 @@ var SquashFixupsAbove = NewIntegrationTest(NewIntegrationTestArgs{
 			}).
 			Lines(
 				Contains("commit 03"),
-				Contains("commit 02"),
-				Contains("commit 01").IsSelected(), // wrong, we want the previous line
-			).
-			SelectPreviousItem()
+				Contains("commit 02").IsSelected(),
+				Contains("commit 01"),
+			)
 
 		t.Views().Main().
 			Content(Contains("fixup content"))
diff --git a/pkg/integration/tests/interactive_rebase/squash_fixups_in_current_branch.go b/pkg/integration/tests/interactive_rebase/squash_fixups_in_current_branch.go
index 75bfbf159..c6721d829 100644
--- a/pkg/integration/tests/interactive_rebase/squash_fixups_in_current_branch.go
+++ b/pkg/integration/tests/interactive_rebase/squash_fixups_in_current_branch.go
@@ -45,11 +45,10 @@ var SquashFixupsInCurrentBranch = NewIntegrationTest(NewIntegrationTestArgs{
 			}).
 			Lines(
 				Contains("commit 02"),
-				Contains("commit 01"),
-				Contains("fixup! master commit").IsSelected(), // wrong, we want the previous line
+				Contains("commit 01").IsSelected(),
+				Contains("fixup! master commit"),
 				Contains("master commit"),
-			).
-			NavigateToLine(Contains("commit 01"))
+			)
 
 		t.Views().Main().
 			Content(Contains("fixup content"))