fixed some #397

2 files changed, 24 insertions, 23 deletions

diff --git a/pkg/gui/commits_panel.go b/pkg/gui/commits_panel.go
index 0d3dd6e84..cd7aa9530 100644
--- a/pkg/gui/commits_panel.go
+++ b/pkg/gui/commits_panel.go
@@ -54,7 +54,7 @@ func (gui *Gui) handleCommitSelect(g *gocui.Gui, v *gocui.View) error {
 
 func (gui *Gui) refreshCommits(g *gocui.Gui) error {
 	g.Update(func(*gocui.Gui) error {
-		builder, err := git.NewCommitListBuilder(gui.Log, gui.GitCommand, gui.OSCommand, gui.Tr, gui.State.CherryPickedCommits)
+		builder, err := git.NewCommitListBuilder(gui.Log, gui.GitCommand, gui.OSCommand, gui.Tr, gui.State.CherryPickedCommits, gui.State.DiffEntries)
 		if err != nil {
 			return err
 		}
@@ -62,13 +62,6 @@ func (gui *Gui) refreshCommits(g *gocui.Gui) error {
 		if err != nil {
 			return err
 		}
-
-		for _, commit := range commits {
-			if _, has := gui.State.DiffEntries[commit.Sha]; has {
-				commit.Selected = true
-			}
-		}
-
 		gui.State.Commits = commits
 
 		gui.refreshSelectedLine(&gui.State.Panels.Commits.SelectedLine, len(gui.State.Commits))
@@ -472,25 +465,20 @@ func (gui *Gui) handleToggleDiffCommit(g *gocui.Gui, v *gocui.View) error {
 	}
 
 	// if already selected commit delete
-	if _, has := gui.State.DiffEntries[commit.Sha]; has {
-		delete(gui.State.DiffEntries, commit.Sha)
-		gui.setDiffMode()
+	if idx, has := gui.hasCommit(gui.State.DiffEntries, commit.Sha); has {
+		gui.State.DiffEntries = gui.unchooseCommit(gui.State.DiffEntries, idx)
 	} else {
 		if len(gui.State.DiffEntries) == selectLimit {
-			return gui.createErrorPanel(gui.g, "you have already selected two commit, please deselect one of two")
+			gui.State.DiffEntries = gui.unchooseCommit(gui.State.DiffEntries, 0)
 		}
-		gui.State.DiffEntries[commit.Sha] = commit
-		gui.setDiffMode()
+		gui.State.DiffEntries = append(gui.State.DiffEntries, commit)
 	}
 
-	// if selected tow commits, display diff between
-	if len(gui.State.DiffEntries) == selectLimit {
-		var entries []string
-		for _, entry := range gui.State.DiffEntries {
-			entries = append(entries, entry.Sha)
-		}
+	gui.setDiffMode()
 
-		commitText, err := gui.GitCommand.DiffCommits(entries[0], entries[1])
+	// if selected two commits, display diff between
+	if len(gui.State.DiffEntries) == selectLimit {
+		commitText, err := gui.GitCommand.DiffCommits(gui.State.DiffEntries[0].Sha, gui.State.DiffEntries[1].Sha)
 
 		if err != nil {
 			return gui.createErrorPanel(gui.g, err.Error())
@@ -514,3 +502,16 @@ func (gui *Gui) setDiffMode() {
 
 	gui.refreshCommits(gui.g)
 }
+
+func (gui *Gui) hasCommit(commits []*commands.Commit, target string) (int, bool) {
+	for idx, commit := range commits {
+		if commit.Sha == target {
+			return idx, true
+		}
+	}
+	return -1, false
+}
+
+func (gui *Gui) unchooseCommit(commits []*commands.Commit, i int) []*commands.Commit {
+	return append(commits[:i], commits[i+1:]...)
+}
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index f6cdd093d..99250c9ac 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -136,7 +136,7 @@ type guiState struct {
 	Commits             []*commands.Commit
 	StashEntries        []*commands.StashEntry
 	CommitFiles         []*commands.CommitFile
-	DiffEntries         map[string]*commands.Commit
+	DiffEntries         []*commands.Commit
 	MenuItemCount       int // can't store the actual list because it's of interface{} type
 	PreviousView        string
 	Platform            commands.Platform
@@ -156,7 +156,7 @@ func NewGui(log *logrus.Entry, gitCommand *commands.GitCommand, oSCommand *comma
 		Commits:             make([]*commands.Commit, 0),
 		CherryPickedCommits: make([]*commands.Commit, 0),
 		StashEntries:        make([]*commands.StashEntry, 0),
-		DiffEntries:         make(map[string]*commands.Commit),
+		DiffEntries:         make([]*commands.Commit, 0),
 		Platform:            *oSCommand.Platform,
 		Panels: &panelStates{
 			Files:       &filePanelState{SelectedLine: -1},