From 79299be3b23725e4055ca264e6fa0a0c16af430f Mon Sep 17 00:00:00 2001
From: Jesse Duffield <jessedduffield@gmail.com>
Date: Sat, 29 Feb 2020 18:44:08 +1100
Subject: [PATCH] better keybindings for patch building mode

---
 pkg/gui/keybindings.go          | 12 ++---------
 pkg/gui/line_by_line_panel.go   |  6 +++++-
 pkg/gui/patch_building_panel.go | 35 +++++++++++----------------------
 pkg/i18n/english.go             |  3 +++
 4 files changed, 21 insertions(+), 35 deletions(-)

diff --git a/pkg/gui/keybindings.go b/pkg/gui/keybindings.go
index ab470a626..ee673839c 100644
--- a/pkg/gui/keybindings.go
+++ b/pkg/gui/keybindings.go
@@ -1137,16 +1137,8 @@ func (gui *Gui) GetInitialKeybindings() []*Binding {
 			Contexts:    []string{"patch-building"},
 			Key:         gui.getKey("universal.select"),
 			Modifier:    gocui.ModNone,
-			Handler:     gui.handleAddSelectionToPatch,
-			Description: gui.Tr.SLocalize("StageSelection"),
-		},
-		{
-			ViewName:    "main",
-			Contexts:    []string{"patch-building"},
-			Key:         gui.getKey("universal.remove"),
-			Modifier:    gocui.ModNone,
-			Handler:     gui.handleRemoveSelectionFromPatch,
-			Description: gui.Tr.SLocalize("ResetSelection"),
+			Handler:     gui.handleToggleSelectionForPatch,
+			Description: gui.Tr.SLocalize("ToggleSelectionForPatch"),
 		},
 		{
 			ViewName:    "main",
diff --git a/pkg/gui/line_by_line_panel.go b/pkg/gui/line_by_line_panel.go
index c9e06eeba..4a0f1258d 100644
--- a/pkg/gui/line_by_line_panel.go
+++ b/pkg/gui/line_by_line_panel.go
@@ -220,6 +220,10 @@ func (gui *Gui) handleMouseScrollDown(g *gocui.Gui, v *gocui.View) error {
 	return gui.handleCycleLine(1)
 }
 
+func (gui *Gui) getSelectedCommitFileName() string {
+	return gui.State.CommitFiles[gui.State.Panels.CommitFiles.SelectedLine].Name
+}
+
 func (gui *Gui) refreshMainView() error {
 	state := gui.State.Panels.LineByLine
 
@@ -227,7 +231,7 @@ func (gui *Gui) refreshMainView() error {
 	// I'd prefer not to have knowledge of contexts using this file but I'm not sure
 	// how to get around this
 	if gui.State.MainContext == "patch-building" {
-		filename := gui.State.CommitFiles[gui.State.Panels.CommitFiles.SelectedLine].Name
+		filename := gui.getSelectedCommitFileName()
 		includedLineIndices = gui.GitCommand.PatchManager.GetFileIncLineIndices(filename)
 	}
 	colorDiff := state.PatchParser.Render(state.FirstLineIdx, state.LastLineIdx, includedLineIndices)
diff --git a/pkg/gui/patch_building_panel.go b/pkg/gui/patch_building_panel.go
index 22e7c4b4a..f4e2eea1d 100644
--- a/pkg/gui/patch_building_panel.go
+++ b/pkg/gui/patch_building_panel.go
@@ -2,6 +2,7 @@ package gui
 
 import (
 	"github.com/jesseduffield/gocui"
+	"github.com/jesseduffield/lazygit/pkg/utils"
 )
 
 func (gui *Gui) refreshPatchBuildingPanel(selectedLineIdx int) error {
@@ -42,38 +43,24 @@ func (gui *Gui) refreshPatchBuildingPanel(selectedLineIdx int) error {
 	return nil
 }
 
-func (gui *Gui) handleAddSelectionToPatch(g *gocui.Gui, v *gocui.View) error {
+func (gui *Gui) handleToggleSelectionForPatch(g *gocui.Gui, v *gocui.View) error {
 	state := gui.State.Panels.LineByLine
 
+	toggleFunc := gui.GitCommand.PatchManager.AddFileLineRange
+	filename := gui.getSelectedCommitFileName()
+	includedLineIndices := gui.GitCommand.PatchManager.GetFileIncLineIndices(filename)
+	currentLineIsStaged := utils.IncludesInt(includedLineIndices, state.SelectedLineIdx)
+	if currentLineIsStaged {
+		toggleFunc = gui.GitCommand.PatchManager.RemoveFileLineRange
+	}
+
 	// add range of lines to those set for the file
 	commitFile := gui.getSelectedCommitFile(gui.g)
 	if commitFile == nil {
 		return gui.renderString(gui.g, "commitFiles", gui.Tr.SLocalize("NoCommiteFiles"))
 	}
 
-	gui.GitCommand.PatchManager.AddFileLineRange(commitFile.Name, state.FirstLineIdx, state.LastLineIdx)
-
-	if err := gui.refreshCommitFilesView(); err != nil {
-		return err
-	}
-
-	if err := gui.refreshPatchBuildingPanel(-1); err != nil {
-		return err
-	}
-
-	return nil
-}
-
-func (gui *Gui) handleRemoveSelectionFromPatch(g *gocui.Gui, v *gocui.View) error {
-	state := gui.State.Panels.LineByLine
-
-	// add range of lines to those set for the file
-	commitFile := gui.getSelectedCommitFile(gui.g)
-	if commitFile == nil {
-		return gui.renderString(gui.g, "commitFiles", gui.Tr.SLocalize("NoCommiteFiles"))
-	}
-
-	gui.GitCommand.PatchManager.RemoveFileLineRange(commitFile.Name, state.FirstLineIdx, state.LastLineIdx)
+	toggleFunc(commitFile.Name, state.FirstLineIdx, state.LastLineIdx)
 
 	if err := gui.refreshCommitFilesView(); err != nil {
 		return err
diff --git a/pkg/i18n/english.go b/pkg/i18n/english.go
index 96f0f4986..a4d16b4d3 100644
--- a/pkg/i18n/english.go
+++ b/pkg/i18n/english.go
@@ -511,6 +511,9 @@ func addEnglish(i18nObject *i18n.Bundle) error {
 		}, &i18n.Message{
 			ID:    "ToggleSelectHunk",
 			Other: `toggle select hunk`,
+		}, &i18n.Message{
+			ID:    "ToggleSelectionForPatch",
+			Other: `add/remove line(s) to patch`,
 		},
 		&i18n.Message{
 			ID:    "TogglePanel",