From 120dd1530ae329199928c3494ea6063d741fc54d Mon Sep 17 00:00:00 2001
From: Stefan Haller <stefan@haller-berlin.de>
Date: Sun, 2 Apr 2023 19:16:28 +0200
Subject: [PATCH] Make EditRebaseTodo more robust

It used to work on the assumption that rebasing commits in lazygit's model
correspond one-to-one to lines in the git-rebase-todo file, which isn't
necessarily true (e.g. when users use "git rebase --edit-todo" at the custom
command prompt and add a "break" between lines).
---
 pkg/commands/git_commands/rebase.go           | 29 +++++++++-------
 .../controllers/local_commits_controller.go   |  4 +--
 .../drop_todo_commit_with_update_ref.go       |  3 --
 pkg/utils/rebaseTodo.go                       | 34 +++++++++++++++++++
 4 files changed, 51 insertions(+), 19 deletions(-)
 create mode 100644 pkg/utils/rebaseTodo.go

diff --git a/pkg/commands/git_commands/rebase.go b/pkg/commands/git_commands/rebase.go
index 1207bdd56..27b5a6ba4 100644
--- a/pkg/commands/git_commands/rebase.go
+++ b/pkg/commands/git_commands/rebase.go
@@ -12,6 +12,7 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/app/daemon"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/commands/oscommands"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
 type RebaseCommands struct {
@@ -202,25 +203,27 @@ func (self *RebaseCommands) AmendTo(commit *models.Commit) error {
 	return self.SquashAllAboveFixupCommits(commit)
 }
 
-// EditRebaseTodo sets the action at a given index in the git-rebase-todo file
-func (self *RebaseCommands) EditRebaseTodo(index int, action todo.TodoCommand) error {
+// EditRebaseTodo sets the action for a given rebase commit in the git-rebase-todo file
+func (self *RebaseCommands) EditRebaseTodo(commit *models.Commit, action todo.TodoCommand) error {
 	fileName := filepath.Join(self.dotGitDir, "rebase-merge/git-rebase-todo")
-	bytes, err := os.ReadFile(fileName)
+	todos, err := utils.ReadRebaseTodoFile(fileName)
 	if err != nil {
 		return err
 	}
 
-	content := strings.Split(string(bytes), "\n")
-	commitCount := self.getTodoCommitCount(content)
+	for i := range todos {
+		t := &todos[i]
+		// Comparing just the sha is not enough; we need to compare both the
+		// action and the sha, as the sha could appear multiple times (e.g. in a
+		// pick and later in a merge)
+		if t.Command == commit.Action && t.Commit == commit.Sha {
+			t.Command = action
+			return utils.WriteRebaseTodoFile(fileName, todos)
+		}
+	}
 
-	// we have the most recent commit at the bottom whereas the todo file has
-	// it at the bottom, so we need to subtract our index from the commit count
-	contentIndex := commitCount - 1 - index
-	splitLine := strings.Split(content[contentIndex], " ")
-	content[contentIndex] = action.String() + " " + strings.Join(splitLine[1:], " ")
-	result := strings.Join(content, "\n")
-
-	return os.WriteFile(fileName, []byte(result), 0o644)
+	// Should never get here
+	return fmt.Errorf("Todo %s not found in git-rebase-todo", commit.Sha)
 }
 
 func (self *RebaseCommands) getTodoCommitCount(content []string) int {
diff --git a/pkg/gui/controllers/local_commits_controller.go b/pkg/gui/controllers/local_commits_controller.go
index 7afb7e0ed..72169a5f0 100644
--- a/pkg/gui/controllers/local_commits_controller.go
+++ b/pkg/gui/controllers/local_commits_controller.go
@@ -354,9 +354,7 @@ func (self *LocalCommitsController) handleMidRebaseCommand(action todo.TodoComma
 		false,
 	)
 
-	if err := self.git.Rebase.EditRebaseTodo(
-		self.context().GetSelectedLineIdx(), action,
-	); err != nil {
+	if err := self.git.Rebase.EditRebaseTodo(commit, action); err != nil {
 		return false, self.c.Error(err)
 	}
 
diff --git a/pkg/integration/tests/interactive_rebase/drop_todo_commit_with_update_ref.go b/pkg/integration/tests/interactive_rebase/drop_todo_commit_with_update_ref.go
index 11d693d29..c247f1743 100644
--- a/pkg/integration/tests/interactive_rebase/drop_todo_commit_with_update_ref.go
+++ b/pkg/integration/tests/interactive_rebase/drop_todo_commit_with_update_ref.go
@@ -60,10 +60,7 @@ var DropTodoCommitWithUpdateRef = NewIntegrationTest(NewIntegrationTestArgs{
 			IsFocused().
 			Lines(
 				Contains("(*) commit 06"),
-				/* EXPECTED:
 				Contains("commit 04"),
-				ACTUAL: */
-				Contains("commit 05"),
 				Contains("(*) commit 03"),
 				Contains("commit 02"),
 				Contains("commit 01"),
diff --git a/pkg/utils/rebaseTodo.go b/pkg/utils/rebaseTodo.go
new file mode 100644
index 000000000..3e071321b
--- /dev/null
+++ b/pkg/utils/rebaseTodo.go
@@ -0,0 +1,34 @@
+package utils
+
+import (
+	"os"
+
+	"github.com/fsmiamoto/git-todo-parser/todo"
+)
+
+func ReadRebaseTodoFile(fileName string) ([]todo.Todo, error) {
+	f, err := os.Open(fileName)
+	if err != nil {
+		return nil, err
+	}
+
+	todos, err := todo.Parse(f)
+	err2 := f.Close()
+	if err == nil {
+		err = err2
+	}
+	return todos, err
+}
+
+func WriteRebaseTodoFile(fileName string, todos []todo.Todo) error {
+	f, err := os.Create(fileName)
+	if err != nil {
+		return err
+	}
+	err = todo.Write(f, todos)
+	err2 := f.Close()
+	if err == nil {
+		err = err2
+	}
+	return err
+}