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.

Author: brandondong

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 {

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}