vastly improve the logic for undo and redo

3 files changed, 120 insertions, 83 deletions

diff --git a/pkg/gui/undoing.go b/pkg/gui/undoing.go
index f050eb5c1..c10914e6d 100644
--- a/pkg/gui/undoing.go
+++ b/pkg/gui/undoing.go
@@ -1,10 +1,9 @@
 package gui
 
 import (
-	"regexp"
-
 	"github.com/jesseduffield/gocui"
 	"github.com/jesseduffield/lazygit/pkg/commands"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
 // Quick summary of how this all works:
@@ -18,108 +17,134 @@ import (
 // two user actions, meaning we end up undoing reflog entry C. Redoing works in a similar way.
 
 const (
-	USER_ACTION = iota
-	UNDO
-	REDO
+	CHECKOUT = iota
+	COMMIT
+	REBASE
+	CURRENT_REBASE
 )
 
 type reflogAction struct {
-	regexStr string
-	action   func(match []string, commitSha string, waitingStatus string, envVars []string, isUndo bool) error
-	kind     int
+	kind int // one of CHECKOUT, REBASE, and COMMIT
+	from string
+	to   string
 }
 
-func (gui *Gui) reflogActions() []reflogAction {
-	return []reflogAction{
-		{
-			regexStr: `^checkout: moving from ([\S]+) to ([\S]+)`,
-			kind:     USER_ACTION,
-			action: func(match []string, commitSha string, waitingStatus string, envVars []string, isUndo bool) error {
-				branchName := match[2]
-				if isUndo {
-					branchName = match[1]
-				}
-				return gui.handleCheckoutRef(branchName, handleCheckoutRefOptions{
-					WaitingStatus: waitingStatus,
-					EnvVars:       envVars,
-				},
-				)
-			},
-		},
-		{
-			regexStr: `^commit|^rebase -i \(start\)|^reset: moving to|^pull`,
-			kind:     USER_ACTION,
-			action: func(match []string, commitSha string, waitingStatus string, envVars []string, isUndo bool) error {
-				return gui.handleHardResetWithAutoStash(commitSha, handleHardResetWithAutoStashOptions{EnvVars: envVars, WaitingStatus: waitingStatus})
-			},
-		},
-		{
-			regexStr: `^\[lazygit undo\]`,
-			kind:     UNDO,
-		},
-		{
-			regexStr: `^\[lazygit redo\]`,
-			kind:     REDO,
-		},
+// Here we're going through the reflog and maintaining a counter that represents how many
+// undos/redos/user actions we've seen. when we hit a user action we call the callback specifying
+// what the counter is up to and the nature of the action.
+// We can't take you from a non-interactive rebase state into an interactive rebase state, so if we hit
+// a 'finish' or an 'abort' entry, we ignore everything else until we find the corresponding 'start' entry.
+// If we find ourselves already in an interactive rebase and we've hit the start entry,
+// we can't really do an undo because there's no way to redo back into the rebase.
+// instead we just ask the user if they want to abort the rebase instead.
+func (gui *Gui) parseReflogForActions(onUserAction func(counter int, action reflogAction) (bool, error)) error {
+	counter := 0
+	reflogCommits := gui.State.ReflogCommits
+	rebaseFinishCommitSha := ""
+	var action *reflogAction
+	for reflogCommitIdx, reflogCommit := range reflogCommits {
+		action = nil
+
+		prevCommitSha := ""
+		if len(reflogCommits)-1 >= reflogCommitIdx+1 {
+			prevCommitSha = reflogCommits[reflogCommitIdx+1].Sha
+		}
+
+		if rebaseFinishCommitSha == "" {
+			if ok, _ := utils.FindStringSubmatch(reflogCommit.Name, `^\[lazygit undo\]`); ok {
+				counter++
+			} else if ok, _ := utils.FindStringSubmatch(reflogCommit.Name, `^\[lazygit redo\]`); ok {
+				counter--
+			} else if ok, _ := utils.FindStringSubmatch(reflogCommit.Name, `^rebase -i \(abort\)|^rebase -i \(finish\)`); ok {
+				rebaseFinishCommitSha = reflogCommit.Sha
+			} else if ok, match := utils.FindStringSubmatch(reflogCommit.Name, `^checkout: moving from ([\S]+) to ([\S]+)`); ok {
+				action = &reflogAction{kind: CHECKOUT, from: match[1], to: match[2]}
+			} else if ok, _ := utils.FindStringSubmatch(reflogCommit.Name, `^commit|^reset: moving to|^pull`); ok {
+				action = &reflogAction{kind: COMMIT, from: prevCommitSha, to: reflogCommit.Sha}
+			} else if ok, _ := utils.FindStringSubmatch(reflogCommit.Name, `^rebase -i \(start\)`); ok {
+				// if we're here then we must be currently inside an interactive rebase
+				action = &reflogAction{kind: CURRENT_REBASE}
+			}
+		} else if ok, _ := utils.FindStringSubmatch(reflogCommit.Name, `^rebase -i \(start\)`); ok {
+			action = &reflogAction{kind: REBASE, from: prevCommitSha, to: rebaseFinishCommitSha}
+		}
+
+		if action != nil {
+			ok, err := onUserAction(counter, *action)
+			if ok {
+				return err
+			}
+			counter--
+			if action.kind == REBASE {
+				rebaseFinishCommitSha = ""
+			}
+		}
 	}
+	return nil
 }
 
 func (gui *Gui) reflogUndo(g *gocui.Gui, v *gocui.View) error {
-	return gui.iterateUserActions(func(match []string, reflogCommits []*commands.Commit, reflogIdx int, action reflogAction, counter int) (bool, error) {
-		if counter == -1 {
-			prevCommitSha := ""
-			if len(reflogCommits)-1 >= reflogIdx+1 {
-				prevCommitSha = reflogCommits[reflogIdx+1].Sha
-			}
-			return true, action.action(match, prevCommitSha, gui.Tr.SLocalize("UndoingStatus"), []string{"GIT_REFLOG_ACTION=[lazygit undo]"}, true)
-		} else {
+	undoEnvVars := []string{"GIT_REFLOG_ACTION=[lazygit undo]"}
+	undoingStatus := gui.Tr.SLocalize("UndoingStatus")
+
+	return gui.parseReflogForActions(func(counter int, action reflogAction) (bool, error) {
+		if counter != 0 {
 			return false, nil
 		}
+
+		switch action.kind {
+		case COMMIT, REBASE:
+			return true, gui.handleHardResetWithAutoStash(action.from, handleHardResetWithAutoStashOptions{
+				EnvVars:       undoEnvVars,
+				WaitingStatus: undoingStatus,
+			})
+		case CURRENT_REBASE:
+			return true, gui.createConfirmationPanel(g, v, true, gui.Tr.SLocalize("AbortRebase"), gui.Tr.SLocalize("UndoOutOfRebaseWarning"), func(g *gocui.Gui, v *gocui.View) error {
+				return gui.genericMergeCommand("abort")
+			}, nil)
+		case CHECKOUT:
+			return true, gui.handleCheckoutRef(action.from, handleCheckoutRefOptions{
+				EnvVars:       undoEnvVars,
+				WaitingStatus: undoingStatus,
+			})
+		}
+
+		gui.Log.Error("didn't match on the user action when trying to undo")
+		return true, nil
 	})
 }
 
 func (gui *Gui) reflogRedo(g *gocui.Gui, v *gocui.View) error {
-	return gui.iterateUserActions(func(match []string, reflogCommits []*commands.Commit, reflogIdx int, action reflogAction, counter int) (bool, error) {
+	redoEnvVars := []string{"GIT_REFLOG_ACTION=[lazygit redo]"}
+	redoingStatus := gui.Tr.SLocalize("RedoingStatus")
+
+	return gui.parseReflogForActions(func(counter int, action reflogAction) (bool, error) {
+		// if we're redoing and the counter is zero, we just return
 		if counter == 0 {
-			return true, action.action(match, reflogCommits[reflogIdx].Sha, gui.Tr.SLocalize("RedoingStatus"), []string{"GIT_REFLOG_ACTION=[lazygit redo]"}, false)
-		} else if counter < 0 {
 			return true, nil
-		} else {
+		} else if counter > 1 {
 			return false, nil
 		}
-	})
-}
 
-func (gui *Gui) iterateUserActions(onUserAction func(match []string, reflogCommits []*commands.Commit, reflogIdx int, action reflogAction, counter int) (bool, error)) error {
-	reflogCommits := gui.State.ReflogCommits
-
-	counter := 0
-	for i, reflogCommit := range reflogCommits {
-		for _, action := range gui.reflogActions() {
-			re := regexp.MustCompile(action.regexStr)
-			match := re.FindStringSubmatch(reflogCommit.Name)
-			if len(match) == 0 {
-				continue
-			}
-
-			switch action.kind {
-			case UNDO:
-				counter++
-			case REDO:
-				counter--
-			case USER_ACTION:
-				counter--
-				shouldReturn, err := onUserAction(match, reflogCommits, i, action, counter)
-				if err != nil {
-					return err
-				}
-				if shouldReturn {
-					return nil
-				}
-			}
+		switch action.kind {
+		case COMMIT, REBASE:
+			return true, gui.handleHardResetWithAutoStash(action.to, handleHardResetWithAutoStashOptions{
+				EnvVars:       redoEnvVars,
+				WaitingStatus: redoingStatus,
+			})
+		case CURRENT_REBASE:
+			// no idea if this is even possible but you certainly can't redo into the end of a rebase if you're still in the rebase
+			return true, nil
+		case CHECKOUT:
+			return true, gui.handleCheckoutRef(action.to, handleCheckoutRefOptions{
+				EnvVars:       redoEnvVars,
+				WaitingStatus: redoingStatus,
+			})
 		}
-	}
-	return nil
+
+		gui.Log.Error("didn't match on the user action when trying to redo")
+		return true, nil
+	})
 }
 
 type handleHardResetWithAutoStashOptions struct {
diff --git a/pkg/i18n/english.go b/pkg/i18n/english.go
index 0d624534d..57aae8021 100644
--- a/pkg/i18n/english.go
+++ b/pkg/i18n/english.go
@@ -1056,6 +1056,12 @@ func addEnglish(i18nObject *i18n.Bundle) error {
 		}, &i18n.Message{
 			ID:    "prevTab",
 			Other: "previous tab",
+		}, &i18n.Message{
+			ID:    "AbortRebase",
+			Other: "Abort rebase",
+		}, &i18n.Message{
+			ID:    "UndoOutOfRebaseWarning",
+			Other: "If you undo at this point, you won't be able to re-enter this rebase by pressing redo. Abort rebase?",
 		},
 	)
 }
diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
index f42ee130c..4afee73b7 100644
--- a/pkg/utils/utils.go
+++ b/pkg/utils/utils.go
@@ -316,3 +316,9 @@ func TruncateWithEllipsis(str string, limit int) string {
 	remainingLength := limit - len(ellipsis)
 	return str[0:remainingLength] + "..."
 }
+
+func FindStringSubmatch(str string, regexpStr string) (bool, []string) {
+	re := regexp.MustCompile(regexpStr)
+	match := re.FindStringSubmatch(str)
+	return len(match) > 0, match
+}