initial support for staging individual lines

12 files changed, 339 insertions, 9 deletions

diff --git a/pkg/commands/git.go b/pkg/commands/git.go
index adc26b8c8..262a57c2f 100644
--- a/pkg/commands/git.go
+++ b/pkg/commands/git.go
@@ -3,6 +3,7 @@ package commands
 import (
 	"errors"
 	"fmt"
+	"io/ioutil"
 	"os"
 	"os/exec"
 	"strings"
@@ -571,9 +572,10 @@ func (c *GitCommand) CheckRemoteBranchExists(branch *Branch) bool {
 }
 
 // Diff returns the diff of a file
-func (c *GitCommand) Diff(file *File) string {
+func (c *GitCommand) Diff(file *File, plain bool) string {
 	cachedArg := ""
 	trackedArg := "--"
+	colorArg := "--color"
 	fileName := c.OSCommand.Quote(file.Name)
 	if file.HasStagedChanges && !file.HasUnstagedChanges {
 		cachedArg = "--cached"
@@ -581,9 +583,36 @@ func (c *GitCommand) Diff(file *File) string {
 	if !file.Tracked && !file.HasStagedChanges {
 		trackedArg = "--no-index /dev/null"
 	}
-	command := fmt.Sprintf("git diff --color %s %s %s", cachedArg, trackedArg, fileName)
+	if plain {
+		colorArg = ""
+	}
+
+	command := fmt.Sprintf("git diff %s %s %s %s", colorArg, cachedArg, trackedArg, fileName)
 
 	// for now we assume an error means the file was deleted
 	s, _ := c.OSCommand.RunCommandWithOutput(command)
 	return s
 }
+
+func (c *GitCommand) ApplyPatch(patch string) (string, error) {
+
+	content := []byte(patch)
+	tmpfile, err := ioutil.TempFile("", "patch")
+	if err != nil {
+		c.Log.Error(err)
+		return "", errors.New("Could not create patch file") // TODO: i18n
+	}
+
+	defer os.Remove(tmpfile.Name()) // clean up
+
+	if _, err := tmpfile.Write(content); err != nil {
+		c.Log.Error(err)
+		return "", err
+	}
+	if err := tmpfile.Close(); err != nil {
+		c.Log.Error(err)
+		return "", err
+	}
+
+	return c.OSCommand.RunCommandWithOutput(fmt.Sprintf("git apply --cached %s", tmpfile.Name()))
+}
diff --git a/pkg/commands/git_test.go b/pkg/commands/git_test.go
index aec928ea5..0e1390535 100644
--- a/pkg/commands/git_test.go
+++ b/pkg/commands/git_test.go
@@ -1822,7 +1822,7 @@ func TestGitCommandDiff(t *testing.T) {
 		t.Run(s.testName, func(t *testing.T) {
 			gitCmd := newDummyGitCommand()
 			gitCmd.OSCommand.command = s.command
-			gitCmd.Diff(s.file)
+			gitCmd.Diff(s.file, false)
 		})
 	}
 }
diff --git a/pkg/git/patch_modifier.go b/pkg/git/patch_modifier.go
index 0e31af952..05afcf5ba 100644
--- a/pkg/git/patch_modifier.go
+++ b/pkg/git/patch_modifier.go
@@ -61,9 +61,13 @@ func (p *PatchModifier) getHunkStart(patchLines []string, lineNumber int) (int,
 func (p *PatchModifier) getModifiedHunk(patchLines []string, hunkStart int, lineNumber int) ([]string, error) {
 	lineChanges := 0
 	// strip the hunk down to just the line we want to stage
-	newHunk := []string{}
-	for offsetIndex, line := range patchLines[hunkStart:] {
-		index := offsetIndex + hunkStart
+	newHunk := []string{patchLines[hunkStart]}
+	for offsetIndex, line := range patchLines[hunkStart+1:] {
+		index := offsetIndex + hunkStart + 1
+		if strings.HasPrefix(line, "@@") {
+			newHunk = append(newHunk, "\n")
+			break
+		}
 		if index != lineNumber {
 			// we include other removals but treat them like context
 			if strings.HasPrefix(line, "-") {
@@ -98,7 +102,12 @@ func (p *PatchModifier) getModifiedHunk(patchLines []string, hunkStart int, line
 func (p *PatchModifier) updatedHeader(currentHeader string, lineChanges int) (string, error) {
 	// current counter is the number after the second comma
 	re := regexp.MustCompile(`^[^,]+,[^,]+,(\d+)`)
-	prevLengthString := re.FindStringSubmatch(currentHeader)[1]
+	matches := re.FindStringSubmatch(currentHeader)
+	if len(matches) < 2 {
+		re = regexp.MustCompile(`^[^,]+,[^+]+\+(\d+)`)
+		matches = re.FindStringSubmatch(currentHeader)
+	}
+	prevLengthString := matches[1]
 
 	prevLength, err := strconv.Atoi(prevLengthString)
 	if err != nil {
diff --git a/pkg/git/patch_parser.go b/pkg/git/patch_parser.go
new file mode 100644
index 000000000..8fa4d355b
--- /dev/null
+++ b/pkg/git/patch_parser.go
@@ -0,0 +1,35 @@
+package git
+
+import (
+	"strings"
+
+	"github.com/sirupsen/logrus"
+)
+
+type PatchParser struct {
+	Log *logrus.Entry
+}
+
+// NewPatchParser builds a new branch list builder
+func NewPatchParser(log *logrus.Entry) (*PatchParser, error) {
+	return &PatchParser{
+		Log: log,
+	}, nil
+}
+
+func (p *PatchParser) ParsePatch(patch string) ([]int, []int, error) {
+	lines := strings.Split(patch, "\n")
+	hunkStarts := []int{}
+	stageableLines := []int{}
+	headerLength := 4
+	for offsetIndex, line := range lines[headerLength:] {
+		index := offsetIndex + headerLength
+		if strings.HasPrefix(line, "@@") {
+			hunkStarts = append(hunkStarts, index)
+		}
+		if strings.HasPrefix(line, "-") || strings.HasPrefix(line, "+") {
+			stageableLines = append(stageableLines, index)
+		}
+	}
+	return hunkStarts, stageableLines, nil
+}
diff --git a/pkg/gui/files_panel.go b/pkg/gui/files_panel.go
index a00bd2843..3404903ef 100644
--- a/pkg/gui/files_panel.go
+++ b/pkg/gui/files_panel.go
@@ -45,6 +45,25 @@ func (gui *Gui) stageSelectedFile(g *gocui.Gui) error {
 	return gui.GitCommand.StageFile(file.Name)
 }
 
+func (gui *Gui) handleSwitchToStagingPanel(g *gocui.Gui, v *gocui.View) error {
+	stagingView, err := g.View("staging")
+	if err != nil {
+		return err
+	}
+	file, err := gui.getSelectedFile(g)
+	if err != nil {
+		if err != gui.Errors.ErrNoFiles {
+			return err
+		}
+		return nil
+	}
+	if !file.Tracked || !file.HasUnstagedChanges {
+		return gui.createErrorPanel(g, gui.Tr.SLocalize("FileStagingRequirements"))
+	}
+	gui.switchFocus(g, v, stagingView)
+	return gui.refreshStagingPanel()
+}
+
 func (gui *Gui) handleFilePress(g *gocui.Gui, v *gocui.View) error {
 	file, err := gui.getSelectedFile(g)
 	if err != nil {
@@ -188,12 +207,11 @@ func (gui *Gui) handleFileSelect(g *gocui.Gui, v *gocui.View) error {
 	if err := gui.renderfilesOptions(g, file); err != nil {
 		return err
 	}
-	var content string
 	if file.HasMergeConflicts {
 		return gui.refreshMergePanel(g)
 	}
 
-	content = gui.GitCommand.Diff(file)
+	content := gui.GitCommand.Diff(file, false)
 	return gui.renderString(g, "main", content)
 }
 
diff --git a/pkg/gui/gui.go b/pkg/gui/gui.go
index ee8fb1165..67415c5ce 100644
--- a/pkg/gui/gui.go
+++ b/pkg/gui/gui.go
@@ -72,6 +72,13 @@ type Gui struct {
 	statusManager *statusManager
 }
 
+type stagingState struct {
+	StageableLines   []int
+	HunkStarts       []int
+	CurrentLineIndex int
+	Diff             string
+}
+
 type guiState struct {
 	Files             []*commands.File
 	Branches          []*commands.Branch
@@ -85,6 +92,7 @@ type guiState struct {
 	EditHistory       *stack.Stack
 	Platform          commands.Platform
 	Updating          bool
+	StagingState      *stagingState
 }
 
 // NewGui builds a new gui handler
@@ -208,6 +216,20 @@ func (gui *Gui) layout(g *gocui.Gui) error {
 		v.FgColor = gocui.ColorWhite
 	}
 
+	v, err = g.SetView("staging", leftSideWidth+panelSpacing, 0, width-1, optionsTop, gocui.LEFT)
+	if err != nil {
+		if err != gocui.ErrUnknownView {
+			return err
+		}
+		v.Title = gui.Tr.SLocalize("StagingTitle")
+		v.Wrap = true
+		v.Highlight = true
+		v.FgColor = gocui.ColorWhite
+		if _, err := g.SetViewOnBottom("staging"); err != nil {
+			return err
+		}
+	}
+
 	if v, err := g.SetView("status", 0, 0, leftSideWidth, statusFilesBoundary, gocui.BOTTOM|gocui.RIGHT); err != nil {
 		if err != gocui.ErrUnknownView {
 			return err
diff --git a/pkg/gui/keybindings.go b/pkg/gui/keybindings.go
index 1b656659d..2c47ab4f3 100644
--- a/pkg/gui/keybindings.go
+++ b/pkg/gui/keybindings.go
@@ -213,6 +213,13 @@ func (gui *Gui) GetKeybindings() []*Binding {
 			Handler:     gui.handleResetHard,
 			Description: gui.Tr.SLocalize("resetHard"),
 		}, {
+			ViewName:    "files",
+			Key:         gocui.KeyEnter,
+			Modifier:    gocui.ModNone,
+			Handler:     gui.handleSwitchToStagingPanel,
+			Description: gui.Tr.SLocalize("StageLines"),
+			KeyReadable: "enter",
+		}, {
 			ViewName: "main",
 			Key:      gocui.KeyEsc,
 			Modifier: gocui.ModNone,
@@ -384,6 +391,26 @@ func (gui *Gui) GetKeybindings() []*Binding {
 			Key:      'q',
 			Modifier: gocui.ModNone,
 			Handler:  gui.handleMenuClose,
+		}, {
+			ViewName: "staging",
+			Key:      gocui.KeyEsc,
+			Modifier: gocui.ModNone,
+			Handler:  gui.handleStagingEscape,
+		}, {
+			ViewName: "staging",
+			Key:      gocui.KeyArrowUp,
+			Modifier: gocui.ModNone,
+			Handler:  gui.handleStagingKeyUp,
+		}, {
+			ViewName: "staging",
+			Key:      gocui.KeyArrowDown,
+			Modifier: gocui.ModNone,
+			Handler:  gui.handleStagingKeyDown,
+		}, {
+			ViewName: "staging",
+			Key:      gocui.KeySpace,
+			Modifier: gocui.ModNone,
+			Handler:  gui.handleStageLine,
 		},
 	}
 
diff --git a/pkg/gui/staging_panel.go b/pkg/gui/staging_panel.go
new file mode 100644
index 000000000..be207c2fb
--- /dev/null
+++ b/pkg/gui/staging_panel.go
@@ -0,0 +1,163 @@
+package gui
+
+import (
+	"errors"
+	"io/ioutil"
+
+	"github.com/davecgh/go-spew/spew"
+
+	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/git"
+)
+
+func (gui *Gui) refreshStagingPanel() error {
+	// get the currently selected file. Get the diff of that file directly, not
+	// using any custom diff tools.
+	// parse the file to find out where the chunks and unstaged changes are
+
+	file, err := gui.getSelectedFile(gui.g)
+	if err != nil {
+		if err != gui.Errors.ErrNoFiles {
+			return err
+		}
+		return gui.handleStagingEscape(gui.g, nil)
+	}
+
+	if !file.HasUnstagedChanges {
+		return gui.handleStagingEscape(gui.g, nil)
+	}
+
+	// note for custom diffs, we'll need to send a flag here saying not to use the custom diff
+	diff := gui.GitCommand.Diff(file, true)
+	colorDiff := gui.GitCommand.Diff(file, false)
+
+	gui.Log.WithField("staging", "staging").Info("DIFF IS:")
+	gui.Log.WithField("staging", "staging").Info(spew.Sdump(diff))
+	gui.Log.WithField("staging", "staging").Info("hello")
+
+	if len(diff) < 2 {
+		return gui.handleStagingEscape(gui.g, nil)
+	}
+
+	// parse the diff and store the line numbers of hunks and stageable lines
+	// TODO: maybe instantiate this at application start
+	p, err := git.NewPatchParser(gui.Log)
+	if err != nil {
+		return nil
+	}
+	hunkStarts, stageableLines, err := p.ParsePatch(diff)
+	if err != nil {
+		return nil
+	}
+
+	var currentLineIndex int
+	if gui.State.StagingState != nil {
+		end := len(stageableLines) - 1
+		if end < gui.State.StagingState.CurrentLineIndex {
+			currentLineIndex = end
+		} else {
+			currentLineIndex = gui.State.StagingState.CurrentLineIndex
+		}
+	} else {
+		currentLineIndex = 0
+	}
+
+	gui.State.StagingState = &stagingState{
+		StageableLines:   stageableLines,
+		HunkStarts:       hunkStarts,
+		CurrentLineIndex: currentLineIndex,
+		Diff:             diff,
+	}
+
+	if len(stageableLines) == 0 {
+		return errors.New("No lines to stage")
+	}
+
+	stagingView := gui.getStagingView(gui.g)
+	stagingView.SetCursor(0, stageableLines[currentLineIndex])
+	stagingView.SetOrigin(0, 0)
+	return gui.renderString(gui.g, "staging", colorDiff)
+}
+
+func (gui *Gui) handleStagingEscape(g *gocui.Gui, v *gocui.View) error {
+	if _, err := gui.g.SetViewOnBottom("staging"); err != nil {
+		return err
+	}
+
+	return gui.switchFocus(gui.g, nil, gui.getFilesView(gui.g))
+}
+
+// nextNumber returns the next index, cycling if we reach the end
+func nextIndex(numbers []int, currentNumber int) int {
+	for index, number := range numbers {
+		if number > currentNumber {
+			return index
+		}
+	}
+	return 0
+}
+
+// prevNumber returns the next number, cycling if we reach the end
+func prevIndex(numbers []int, currentNumber int) int {
+	end := len(numbers) - 1
+	for i := end; i >= 0; i -= 1 {
+		if numbers[i] < currentNumber {
+			return i
+		}
+	}
+	return end
+}
+
+func (gui *Gui) handleStagingKeyUp(g *gocui.Gui, v *gocui.View) error {
+	return gui.handleCycleLine(true)
+}
+
+func (gui *Gui) handleStagingKeyDown(g *gocui.Gui, v *gocui.View) error {
+	return gui.handleCycleLine(false)
+}
+
+func (gui *Gui) handleCycleLine(up bool) error {
+	state := gui.State.StagingState
+	lineNumbers := state.StageableLines
+	currentLine := lineNumbers[state.CurrentLineIndex]
+	var newIndex int
+	if up {
+		newIndex = prevIndex(lineNumbers, currentLine)
+	} else {
+		newIndex = nextIndex(lineNumbers, currentLine)
+	}
+
+	state.CurrentLineIndex = newIndex
+	stagingView := gui.getStagingView(gui.g)
+	stagingView.SetCursor(0, lineNumbers[newIndex])
+	stagingView.SetOrigin(0, 0)
+	return nil
+}
+
+func (gui *Gui) handleStageLine(g *gocui.Gui, v *gocui.View) error {
+	state := gui.State.StagingState
+	p, err := git.NewPatchModifier(gui.Log)
+	if err != nil {
+		return err
+	}
+
+	currentLine := state.StageableLines[state.CurrentLineIndex]
+	patch, err := p.ModifyPatch(state.Diff, currentLine)
+	if err != nil {
+		return err
+	}
+
+	// for logging purposes
+	ioutil.WriteFile("patch.diff", []byte(patch), 0600)
+
+	// apply the patch then refresh this panel
+	// create a new temp file with the patch, then call git apply with that patch
+	_, err = gui.GitCommand.ApplyPatch(patch)
+	if err != nil {
+		panic(err)
+	}
+
+	gui.refreshStagingPanel()
+	gui.refreshFiles(gui.g)
+	return nil
+}
diff --git a/pkg/gui/view_helpers.go b/pkg/gui/view_helpers.go
index 6c3e5505c..e6970d92f 100644
--- a/pkg/gui/view_helpers.go
+++ b/pkg/gui/view_helpers.go
@@ -103,6 +103,9 @@ func (gui *Gui) newLineFocused(g *gocui.Gui, v *gocui.View) error {
 		return gui.handleCommitSelect(g, v)
 	case "stash":
 		return gui.handleStashEntrySelect(g, v)
+	case "staging":
+		return nil
+		// return gui.handleStagingSelect(g, v)
 	default:
 		panic(gui.Tr.SLocalize("NoViewMachingNewLineFocusedSwitchStatement"))
 	}
@@ -153,6 +156,10 @@ func (gui *Gui) switchFocus(g *gocui.Gui, oldView, newView *gocui.View) error {
 	if _, err := g.SetCurrentView(newView.Name()); err != nil {
 		return err
 	}
+	if _, err := g.SetViewOnTop(newView.Name()); err != nil {
+		return err
+	}
+
 	g.Cursor = newView.Editable
 
 	return gui.newLineFocused(g, newView)
@@ -293,6 +300,11 @@ func (gui *Gui) getBranchesView(g *gocui.Gui) *gocui.View {
 	return v
 }
 
+func (gui *Gui) getStagingView(g *gocui.Gui) *gocui.View {
+	v, _ := g.View("staging")
+	return v
+}
+
 func (gui *Gui) trimmedContent(v *gocui.View) string {
 	return strings.TrimSpace(v.Buffer())
 }
diff --git a/pkg/i18n/dutch.go b/pkg/i18n/dutch.go
index 9ca18857b..b74d511e4 100644
--- a/pkg/i18n/dutch.go
+++ b/pkg/i18n/dutch.go
@@ -403,6 +403,9 @@ func addDutch(i18nObject *i18n.Bundle) error {
 		}, &i18n.Message{
 			ID:    "NoBranchOnRemote",
 			Other: `Deze branch bestaat niet op de remote. U moet het eerst naar de remote pushen.`,
+		}, &i18n.Message{
+			ID:    "StageLines",
+			Other: `stage individual hunks/lines`,
 		},
 	)
 }
diff --git a/pkg/i18n/english.go b/pkg/i18n/english.go
index 9e0e60166..7a032e105 100644
--- a/pkg/i18n/english.go
+++ b/pkg/i18n/english.go
@@ -411,6 +411,15 @@ func addEnglish(i18nObject *i18n.Bundle) error {
 		}, &i18n.Message{
 			ID:    "NoBranchOnRemote",
 			Other: `This branch doesn't exist on remote. You need to push it to remote first.`,
+		}, &i18n.Message{
+			ID:    "StageLines",
+			Other: `stage individual hunks/lines`,
+		}, &i18n.Message{
+			ID:    "FileStagingRequirements",
+			Other: `Can only stage individual lines for tracked files with unstaged changes`,
+		}, &i18n.Message{
+			ID:    "StagingTitle",
+			Other: `Staging`,
 		},
 	)
 }
diff --git a/pkg/i18n/polish.go b/pkg/i18n/polish.go
index 975035771..27035a18e 100644
--- a/pkg/i18n/polish.go
+++ b/pkg/i18n/polish.go
@@ -386,6 +386,9 @@ func addPolish(i18nObject *i18n.Bundle) error {
 		}, &i18n.Message{
 			ID:    "NoBranchOnRemote",
 			Other: `Ta gałąź nie istnieje na zdalnym. Najpierw musisz go odepchnąć na odległość.`,
+		}, &i18n.Message{
+			ID:    "StageLines",
+			Other: `stage individual hunks/lines`,
 		},
 	)
 }