better keybindings for patch building mode

3 files changed, 16 insertions, 33 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
 
-	// 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
+	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
 	}
 
-	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