separate commits from cherry pick state

5 files changed, 74 insertions, 51 deletions

diff --git a/pkg/commands/commit.go b/pkg/commands/commit.go
index 32daeaa23..a400da928 100644
--- a/pkg/commands/commit.go
+++ b/pkg/commands/commit.go
@@ -2,16 +2,14 @@ package commands
 
 // Commit : A git commit
 type Commit struct {
-	Sha           string
-	Name          string
-	Status        string // one of "unpushed", "pushed", "merged", "rebasing" or "selected"
-	DisplayString string
-	Action        string // one of "", "pick", "edit", "squash", "reword", "drop", "fixup"
-	Copied        bool   // to know if this commit is ready to be cherry-picked somewhere
-	Tags          []string
-	ExtraInfo     string // something like 'HEAD -> master, tag: v0.15.2'
-	Author        string
-	Date          string
+	Sha       string
+	Name      string
+	Status    string // one of "unpushed", "pushed", "merged", "rebasing" or "selected"
+	Action    string // one of "", "pick", "edit", "squash", "reword", "drop", "fixup"
+	Tags      []string
+	ExtraInfo string // something like 'HEAD -> master, tag: v0.15.2'
+	Author    string
+	Date      string
 }
 
 func (c *Commit) ShortSha() string {
diff --git a/pkg/commands/commit_list_builder.go b/pkg/commands/commit_list_builder.go
index 834d478e0..7f7a979b5 100644
--- a/pkg/commands/commit_list_builder.go
+++ b/pkg/commands/commit_list_builder.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"os/exec"
 	"path/filepath"
 	"regexp"
 	"strings"
@@ -70,13 +71,12 @@ func (c *CommitListBuilder) extractCommitFromLine(line string) *Commit {
 	}
 
 	return &Commit{
-		Sha:           sha,
-		Name:          message,
-		DisplayString: line,
-		Tags:          tags,
-		ExtraInfo:     extraInfo,
-		Date:          date,
-		Author:        author,
+		Sha:       sha,
+		Name:      message,
+		Tags:      tags,
+		ExtraInfo: extraInfo,
+		Date:      date,
+		Author:    author,
 	}
 }
 
@@ -100,15 +100,20 @@ func (c *CommitListBuilder) GetCommits(limit bool) ([]*Commit, error) {
 	}
 
 	unpushedCommits := c.getUnpushedCommits()
-	log := c.getLog(limit)
+	cmd := c.getLogCmd(limit)
 
-	// now we can split it up and turn it into commits
-	for _, line := range utils.SplitLines(log) {
+	err = RunLineOutputCmd(cmd, func(line string) (bool, error) {
 		commit := c.extractCommitFromLine(line)
 		_, unpushed := unpushedCommits[commit.ShortSha()]
 		commit.Status = map[bool]string{true: "unpushed", false: "pushed"}[unpushed]
 		commits = append(commits, commit)
+
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
 	}
+
 	if rebaseMode != "" {
 		currentCommit := commits[len(rebasingCommits)]
 		blue := color.New(color.FgYellow)
@@ -121,11 +126,6 @@ func (c *CommitListBuilder) GetCommits(limit bool) ([]*Commit, error) {
 		return nil, err
 	}
 
-	commits, err = c.setCommitCherryPickStatuses(commits)
-	if err != nil {
-		return nil, err
-	}
-
 	for _, commit := range commits {
 		for _, entry := range c.DiffEntries {
 			if entry.Sha == commit.Sha {
@@ -267,17 +267,6 @@ func (c *CommitListBuilder) setCommitMergedStatuses(commits []*Commit) ([]*Commi
 	return commits, nil
 }
 
-func (c *CommitListBuilder) setCommitCherryPickStatuses(commits []*Commit) ([]*Commit, error) {
-	for _, commit := range commits {
-		for _, cherryPickedCommit := range c.CherryPickedCommits {
-			if commit.Sha == cherryPickedCommit.Sha {
-				commit.Copied = true
-			}
-		}
-	}
-	return commits, nil
-}
-
 func (c *CommitListBuilder) getMergeBase() (string, error) {
 	currentBranch, _, err := c.GitCommand.CurrentBranchName()
 	if err != nil {
diff --git a/pkg/commands/os.go b/pkg/commands/os.go
index 7b96b99c4..517372662 100644
--- a/pkg/commands/os.go
+++ b/pkg/commands/os.go
@@ -1,6 +1,7 @@
 package commands
 
 import (
+	"bufio"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -425,3 +426,31 @@ func Kill(cmd *exec.Cmd) error {
 	}
 	return cmd.Process.Kill()
 }
+
+func RunLineOutputCmd(cmd *exec.Cmd, onLine func(line string) (bool, error)) error {
+	stdoutPipe, err := cmd.StdoutPipe()
+	if err != nil {
+		return err
+	}
+
+	scanner := bufio.NewScanner(stdoutPipe)
+	scanner.Split(bufio.ScanLines)
+	if err := cmd.Start(); err != nil {
+		return err
+	}
+
+	for scanner.Scan() {
+		line := scanner.Text()
+		stop, err := onLine(line)
+		if err != nil {
+			return err
+		}
+		if stop {
+			cmd.Process.Kill()
+			break
+		}
+	}
+
+	cmd.Wait()
+	return nil
+}
diff --git a/pkg/gui/commits_panel.go b/pkg/gui/commits_panel.go
index a31cd863a..7b620e571 100644
--- a/pkg/gui/commits_panel.go
+++ b/pkg/gui/commits_panel.go
@@ -351,14 +351,21 @@ func (gui *Gui) handleCopyCommit(g *gocui.Gui, v *gocui.View) error {
 	return gui.refreshCommits(gui.g)
 }
 
+func (gui *Gui) cherryPickedCommitShaMap() map[string]bool {
+	commitShaMap := map[string]bool{}
+	for _, commit := range gui.State.CherryPickedCommits {
+		commitShaMap[commit.Sha] = true
+	}
+	return commitShaMap
+}
+
 func (gui *Gui) addCommitToCherryPickedCommits(index int) {
-	// not super happy with modifying the state of the Commits array here
-	// but the alternative would be very tricky
-	gui.State.Commits[index].Copied = true
+	commitShaMap := gui.cherryPickedCommitShaMap()
+	commitShaMap[gui.State.Commits[index].Sha] = true
 
 	newCommits := []*commands.Commit{}
 	for _, commit := range gui.State.Commits {
-		if commit.Copied {
+		if commitShaMap[commit.Sha] {
 			// duplicating just the things we need to put in the rebase TODO list
 			newCommits = append(newCommits, &commands.Commit{Name: commit.Name, Sha: commit.Sha})
 		}
@@ -368,13 +375,13 @@ func (gui *Gui) addCommitToCherryPickedCommits(index int) {
 }
 
 func (gui *Gui) handleCopyCommitRange(g *gocui.Gui, v *gocui.View) error {
-	// whenever I add a commit, I need to make sure I retain its order
+	commitShaMap := gui.cherryPickedCommitShaMap()
 
 	// find the last commit that is copied that's above our position
 	// if there are none, startIndex = 0
 	startIndex := 0
 	for index, commit := range gui.State.Commits[0:gui.State.Panels.Commits.SelectedLine] {
-		if commit.Copied {
+		if commitShaMap[commit.Sha] {
 			startIndex = index
 		}
 	}
@@ -385,7 +392,7 @@ func (gui *Gui) handleCopyCommitRange(g *gocui.Gui, v *gocui.View) error {
 		gui.addCommitToCherryPickedCommits(index)
 	}
 
-	return gui.refreshCommits(gui.g)
+	return gui.renderBranchCommitsWithSelection()
 }
 
 // HandlePasteCommits begins a cherry-pick rebase with the commits the user has copied
@@ -548,7 +555,7 @@ func (gui *Gui) renderBranchCommitsWithSelection() error {
 	commitsView := gui.getCommitsView()
 
 	gui.refreshSelectedLine(&gui.State.Panels.Commits.SelectedLine, len(gui.State.Commits))
-	displayStrings := presentation.GetCommitListDisplayStrings(gui.State.Commits, gui.State.ScreenMode != SCREEN_NORMAL)
+	displayStrings := presentation.GetCommitListDisplayStrings(gui.State.Commits, gui.State.ScreenMode != SCREEN_NORMAL, gui.cherryPickedCommitShaMap())
 	gui.renderDisplayStrings(commitsView, displayStrings)
 	if gui.g.CurrentView() == commitsView && commitsView.Context == "branch-commits" {
 		if err := gui.handleCommitSelect(gui.g, commitsView); err != nil {
diff --git a/pkg/gui/presentation/commits.go b/pkg/gui/presentation/commits.go
index 00b8e7f34..dedf4a01b 100644
--- a/pkg/gui/presentation/commits.go
+++ b/pkg/gui/presentation/commits.go
@@ -9,10 +9,10 @@ import (
 	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
-func GetCommitListDisplayStrings(commits []*commands.Commit, fullDescription bool) [][]string {
+func GetCommitListDisplayStrings(commits []*commands.Commit, fullDescription bool, cherryPickedCommitShaMap map[string]bool) [][]string {
 	lines := make([][]string, len(commits))
 
-	var displayFunc func(*commands.Commit) []string
+	var displayFunc func(*commands.Commit, map[string]bool) []string
 	if fullDescription {
 		displayFunc = getFullDescriptionDisplayStringsForCommit
 	} else {
@@ -20,13 +20,13 @@ func GetCommitListDisplayStrings(commits []*commands.Commit, fullDescription boo
 	}
 
 	for i := range commits {
-		lines[i] = displayFunc(commits[i])
+		lines[i] = displayFunc(commits[i], cherryPickedCommitShaMap)
 	}
 
 	return lines
 }
 
-func getFullDescriptionDisplayStringsForCommit(c *commands.Commit) []string {
+func getFullDescriptionDisplayStringsForCommit(c *commands.Commit, cherryPickedCommitShaMap map[string]bool) []string {
 	red := color.New(color.FgRed)
 	yellow := color.New(color.FgYellow)
 	green := color.New(color.FgGreen)
@@ -58,7 +58,7 @@ func getFullDescriptionDisplayStringsForCommit(c *commands.Commit) []string {
 		shaColor = defaultColor
 	}
 
-	if c.Copied {
+	if cherryPickedCommitShaMap[c.Sha] {
 		shaColor = copied
 	}
 
@@ -77,7 +77,7 @@ func getFullDescriptionDisplayStringsForCommit(c *commands.Commit) []string {
 	return []string{shaColor.Sprint(c.ShortSha()), secondColumnString, yellow.Sprint(truncatedAuthor), tagString + defaultColor.Sprint(c.Name)}
 }
 
-func getDisplayStringsForCommit(c *commands.Commit) []string {
+func getDisplayStringsForCommit(c *commands.Commit, cherryPickedCommitShaMap map[string]bool) []string {
 	red := color.New(color.FgRed)
 	yellow := color.New(color.FgYellow)
 	green := color.New(color.FgGreen)
@@ -109,7 +109,7 @@ func getDisplayStringsForCommit(c *commands.Commit) []string {
 		shaColor = defaultColor
 	}
 
-	if c.Copied {
+	if cherryPickedCommitShaMap[c.Sha] {
 		shaColor = copied
 	}