refactor cherry pick code to move state access out of helper

pkg/gui/controllers.go

@@ -7,7 +7,6 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/controllers"
 	"github.com/jesseduffield/lazygit/pkg/gui/controllers/helpers"
-	"github.com/jesseduffield/lazygit/pkg/gui/modes/cherrypicking"
 	"github.com/jesseduffield/lazygit/pkg/gui/services/custom_commands"
 	"github.com/jesseduffield/lazygit/pkg/gui/types"
 )
@@ -53,7 +52,6 @@ func (gui *Gui) resetControllers() {
 			helperCommon,
 			gui.git,
 			gui.State.Contexts,
-			func() *cherrypicking.CherryPicking { return gui.State.Modes.CherryPicking },
 			rebaseHelper,
 		),
 		Upstream:        helpers.NewUpstreamHelper(helperCommon, model, suggestionsHelper.GetRemoteBranchesSuggestionsFunc),

pkg/gui/controllers/helpers/cherry_pick_helper.go

@@ -1,8 +1,6 @@
 package helpers
 
 import (
-	"github.com/jesseduffield/generics/set"
-	"github.com/jesseduffield/generics/slices"
 	"github.com/jesseduffield/lazygit/pkg/commands"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 	"github.com/jesseduffield/lazygit/pkg/gui/context"
@@ -16,7 +14,6 @@ type CherryPickHelper struct {
 	git *commands.GitCommand
 
 	contexts *context.ContextTree
-	getData  func() *cherrypicking.CherryPicking
 
 	rebaseHelper *MergeAndRebaseHelper
 }
@@ -28,35 +25,32 @@ func NewCherryPickHelper(
 	c *types.HelperCommon,
 	git *commands.GitCommand,
 	contexts *context.ContextTree,
-	getData func() *cherrypicking.CherryPicking,
 	rebaseHelper *MergeAndRebaseHelper,
 ) *CherryPickHelper {
 	return &CherryPickHelper{
 		c:            c,
 		git:          git,
 		contexts:     contexts,
-		getData:      getData,
 		rebaseHelper: rebaseHelper,
 	}
 }
 
+func (self *CherryPickHelper) getData() *cherrypicking.CherryPicking {
+	return self.c.Modes().CherryPicking
+}
+
 func (self *CherryPickHelper) Copy(commit *models.Commit, commitsList []*models.Commit, context types.Context) error {
 	if err := self.resetIfNecessary(context); err != nil {
 		return err
 	}
 
 	// we will un-copy it if it's already copied
-	for index, cherryPickedCommit := range self.getData().CherryPickedCommits {
-		if commit.Sha == cherryPickedCommit.Sha {
-			self.getData().CherryPickedCommits = append(
-				self.getData().CherryPickedCommits[0:index],
-				self.getData().CherryPickedCommits[index+1:]...,
-			)
-			return self.rerender()
-		}
+	if self.getData().SelectedShaSet().Includes(commit.Sha) {
+		self.getData().Remove(commit, commitsList)
+	} else {
+		self.getData().Add(commit, commitsList)
 	}
 
-	self.add(commit, commitsList)
 	return self.rerender()
 }
 
@@ -65,7 +59,7 @@ func (self *CherryPickHelper) CopyRange(selectedIndex int, commitsList []*models
 		return err
 	}
 
-	commitSet := self.CherryPickedCommitShaSet()
+	commitSet := self.getData().SelectedShaSet()
 
 	// find the last commit that is copied that's above our position
 	// if there are none, startIndex = 0
@@ -78,7 +72,7 @@ func (self *CherryPickHelper) CopyRange(selectedIndex int, commitsList []*models
 
 	for index := startIndex; index <= selectedIndex; index++ {
 		commit := commitsList[index]
-		self.add(commit, commitsList)
+		self.getData().Add(commit, commitsList)
 	}
 
 	return self.rerender()
@@ -107,26 +101,6 @@ func (self *CherryPickHelper) Reset() error {
 	return self.rerender()
 }
 
-func (self *CherryPickHelper) CherryPickedCommitShaSet() *set.Set[string] {
-	shas := slices.Map(self.getData().CherryPickedCommits, func(commit *models.Commit) string {
-		return commit.Sha
-	})
-	return set.NewFromSlice(shas)
-}
-
-func (self *CherryPickHelper) add(selectedCommit *models.Commit, commitsList []*models.Commit) {
-	commitSet := self.CherryPickedCommitShaSet()
-	commitSet.Add(selectedCommit.Sha)
-
-	cherryPickedCommits := slices.Filter(commitsList, func(commit *models.Commit) bool {
-		return commitSet.Includes(commit.Sha)
-	})
-
-	self.getData().CherryPickedCommits = slices.Map(cherryPickedCommits, func(commit *models.Commit) *models.Commit {
-		return &models.Commit{Name: commit.Name, Sha: commit.Sha}
-	})
-}
-
 // you can only copy from one context at a time, because the order and position of commits matter
 func (self *CherryPickHelper) resetIfNecessary(context types.Context) error {
 	oldContextKey := types.ContextKey(self.getData().ContextKey)

pkg/gui/list_context_config.go

@@ -42,16 +42,6 @@ func (gui *Gui) remoteBranchesListContext() *context.RemoteBranchesContext {
 	)
 }
 
-func (gui *Gui) withDiffModeCheck(f func() error) func() error {
-	return func() error {
-		if gui.State.Modes.Diffing.Active() {
-			return gui.helpers.Diff.RenderDiff()
-		}
-
-		return f()
-	}
-}
-
 func (gui *Gui) tagsListContext() *context.TagsContext {
 	return context.NewTagsContext(
 		func(startIdx int, length int) [][]string {
@@ -78,7 +68,7 @@ func (gui *Gui) branchCommitsListContext() *context.LocalCommitsContext {
 				gui.Common,
 				gui.State.Model.Commits,
 				gui.State.ScreenMode != types.SCREEN_NORMAL,
-				gui.helpers.CherryPick.CherryPickedCommitShaSet(),
+				gui.c.Modes().CherryPicking.SelectedShaSet(),
 				gui.State.Modes.Diffing.Ref,
 				gui.c.UserConfig.Gui.TimeFormat,
 				gui.c.UserConfig.Git.ParseEmoji,
@@ -108,7 +98,7 @@ func (gui *Gui) subCommitsListContext() *context.SubCommitsContext {
 				gui.Common,
 				gui.State.Model.SubCommits,
 				gui.State.ScreenMode != types.SCREEN_NORMAL,
-				gui.helpers.CherryPick.CherryPickedCommitShaSet(),
+				gui.c.Modes().CherryPicking.SelectedShaSet(),
 				gui.State.Modes.Diffing.Ref,
 				gui.c.UserConfig.Gui.TimeFormat,
 				gui.c.UserConfig.Git.ParseEmoji,
@@ -149,7 +139,7 @@ func (gui *Gui) reflogCommitsListContext() *context.ReflogCommitsContext {
 			return presentation.GetReflogCommitListDisplayStrings(
 				gui.State.Model.FilteredReflogCommits,
 				gui.State.ScreenMode != types.SCREEN_NORMAL,
-				gui.helpers.CherryPick.CherryPickedCommitShaSet(),
+				gui.c.Modes().CherryPicking.SelectedShaSet(),
 				gui.State.Modes.Diffing.Ref,
 				gui.c.UserConfig.Gui.TimeFormat,
 				gui.c.UserConfig.Git.ParseEmoji,

pkg/gui/modes/cherrypicking/cherry_picking.go

@@ -1,6 +1,8 @@
 package cherrypicking
 
 import (
+	"github.com/jesseduffield/generics/set"
+	"github.com/jesseduffield/generics/slices"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
 )
 
@@ -18,6 +20,37 @@ func New() *CherryPicking {
 	}
 }
 
-func (m *CherryPicking) Active() bool {
-	return len(m.CherryPickedCommits) > 0
+func (self *CherryPicking) Active() bool {
+	return len(self.CherryPickedCommits) > 0
+}
+
+func (self *CherryPicking) SelectedShaSet() *set.Set[string] {
+	shas := slices.Map(self.CherryPickedCommits, func(commit *models.Commit) string {
+		return commit.Sha
+	})
+	return set.NewFromSlice(shas)
+}
+
+func (self *CherryPicking) Add(selectedCommit *models.Commit, commitsList []*models.Commit) {
+	commitSet := self.SelectedShaSet()
+	commitSet.Add(selectedCommit.Sha)
+
+	self.update(commitSet, commitsList)
+}
+
+func (self *CherryPicking) Remove(selectedCommit *models.Commit, commitsList []*models.Commit) {
+	commitSet := self.SelectedShaSet()
+	commitSet.Remove(selectedCommit.Sha)
+
+	self.update(commitSet, commitsList)
+}
+
+func (self *CherryPicking) update(selectedShaSet *set.Set[string], commitsList []*models.Commit) {
+	cherryPickedCommits := slices.Filter(commitsList, func(commit *models.Commit) bool {
+		return selectedShaSet.Includes(commit.Sha)
+	})
+
+	self.CherryPickedCommits = slices.Map(cherryPickedCommits, func(commit *models.Commit) *models.Commit {
+		return &models.Commit{Name: commit.Name, Sha: commit.Sha}
+	})
 }