From 1fa9ea7f0401b518a1b3661ede82bc903d203323 Mon Sep 17 00:00:00 2001
From: Brandon <brandondong96@gmail.com>
Date: Sat, 22 Feb 2025 14:09:01 -0800
Subject: [PATCH] Fix discarding submodule changes in nested folders

The current rules for discarding submodule changes is that no other changed item
must be also selected. There are some bugs with the current implementation when
submodules are in folders.

For example, selecting and discarding a folder with only a nested submodule
change will currently do nothing. The submodule changes should be discarded. The
folder only contains submodule changes so it should be no different than
pressing discard on the submodule entry itself.

Also, I noticed range selecting both the folder and the submodule and then
pressing discard would be incorrectly disallowed.
---
 pkg/gui/controllers/files_controller.go | 46 +++++++++++++++++++------
 pkg/gui/filetree/node.go                | 16 +++++++++
 2 files changed, 51 insertions(+), 11 deletions(-)

diff --git a/pkg/gui/controllers/files_controller.go b/pkg/gui/controllers/files_controller.go
index 5f4a3bae0..e37127181 100644
--- a/pkg/gui/controllers/files_controller.go
+++ b/pkg/gui/controllers/files_controller.go
@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/jesseduffield/generics/set"
 	"github.com/jesseduffield/gocui"
 	"github.com/jesseduffield/lazygit/pkg/commands/git_commands"
 	"github.com/jesseduffield/lazygit/pkg/commands/models"
@@ -1194,13 +1195,36 @@ func filterNodesHaveUnstagedChanges(nodes []*filetree.FileNode) []*filetree.File
 	})
 }
 
+func findSubmoduleNode(nodes []*filetree.FileNode, submodules []*models.SubmoduleConfig) *models.File {
+	for _, node := range nodes {
+		submoduleNode := node.FindFirstFileBy(func(f *models.File) bool {
+			return f.IsSubmodule(submodules)
+		})
+		if submoduleNode != nil {
+			return submoduleNode
+		}
+	}
+	return nil
+}
+
 func (self *FilesController) canRemove(selectedNodes []*filetree.FileNode) *types.DisabledReason {
+	// Return disabled if the selection contains multiple changed items and includes a submodule change.
 	submodules := self.c.Model().Submodules
-	submoduleCount := lo.CountBy(selectedNodes, func(node *filetree.FileNode) bool {
-		return node.File != nil && node.File.IsSubmodule(submodules)
-	})
-	if submoduleCount > 0 && len(selectedNodes) > 1 {
-		return &types.DisabledReason{Text: self.c.Tr.RangeSelectNotSupportedForSubmodules}
+	hasFiles := false
+	uniqueSelectedSubmodules := set.New[*models.SubmoduleConfig]()
+
+	for _, node := range selectedNodes {
+		_ = node.ForEachFile(func(f *models.File) error {
+			if submodule := f.SubmoduleConfig(submodules); submodule != nil {
+				uniqueSelectedSubmodules.Add(submodule)
+			} else {
+				hasFiles = true
+			}
+			return nil
+		})
+		if uniqueSelectedSubmodules.Len() > 0 && (hasFiles || uniqueSelectedSubmodules.Len() > 1) {
+			return &types.DisabledReason{Text: self.c.Tr.RangeSelectNotSupportedForSubmodules}
+		}
 	}
 
 	return nil
@@ -1209,11 +1233,13 @@ func (self *FilesController) canRemove(selectedNodes []*filetree.FileNode) *type
 func (self *FilesController) remove(selectedNodes []*filetree.FileNode) error {
 	submodules := self.c.Model().Submodules
 
+	selectedNodes = normalisedSelectedNodes(selectedNodes)
+
 	// If we have one submodule then we must only have one submodule or `canRemove` would have
 	// returned an error
-	firstNode := selectedNodes[0]
-	if firstNode.File != nil && firstNode.File.IsSubmodule(submodules) {
-		submodule := firstNode.File.SubmoduleConfig(submodules)
+	submoduleNode := findSubmoduleNode(selectedNodes, submodules)
+	if submoduleNode != nil {
+		submodule := submoduleNode.SubmoduleConfig(submodules)
 
 		menuItems := []*types.MenuItem{
 			{
@@ -1224,11 +1250,9 @@ func (self *FilesController) remove(selectedNodes []*filetree.FileNode) error {
 			},
 		}
 
-		return self.c.Menu(types.CreateMenuOptions{Title: firstNode.GetPath(), Items: menuItems})
+		return self.c.Menu(types.CreateMenuOptions{Title: submoduleNode.GetPath(), Items: menuItems})
 	}
 
-	selectedNodes = normalisedSelectedNodes(selectedNodes)
-
 	discardAllChangesItem := types.MenuItem{
 		Label: self.c.Tr.DiscardAllChanges,
 		OnPress: func() error {
diff --git a/pkg/gui/filetree/node.go b/pkg/gui/filetree/node.go
index ce3bac566..97d5232b5 100644
--- a/pkg/gui/filetree/node.go
+++ b/pkg/gui/filetree/node.go
@@ -169,6 +169,22 @@ func (self *Node[T]) EveryFile(predicate func(*T) bool) bool {
 	return true
 }
 
+func (self *Node[T]) FindFirstFileBy(predicate func(*T) bool) *T {
+	if self.IsFile() {
+		if predicate(self.File) {
+			return self.File
+		}
+	} else {
+		for _, child := range self.Children {
+			if file := child.FindFirstFileBy(predicate); file != nil {
+				return file
+			}
+		}
+	}
+
+	return nil
+}
+
 func (self *Node[T]) Flatten(collapsedPaths *CollapsedPaths) []*Node[T] {
 	result := []*Node[T]{self}