refactor cherry pick code to move state access out of helper

1 files changed, 35 insertions, 2 deletions

diff --git a/pkg/gui/modes/cherrypicking/cherry_picking.go b/pkg/gui/modes/cherrypicking/cherry_picking.go
index bd5c6437a..1fd34d473 100644
--- a/pkg/gui/modes/cherrypicking/cherry_picking.go
+++ b/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}
+	})
 }