make more use of generics

6 files changed, 49 insertions, 31 deletions

diff --git a/pkg/gui/filetree/collapsed_paths.go b/pkg/gui/filetree/collapsed_paths.go
index 02c0b4303..903999b37 100644
--- a/pkg/gui/filetree/collapsed_paths.go
+++ b/pkg/gui/filetree/collapsed_paths.go
@@ -1,20 +1,38 @@
 package filetree
 
-type CollapsedPaths map[string]bool
+import "github.com/jesseduffield/generics/set"
 
-func (cp CollapsedPaths) ExpandToPath(path string) {
+type CollapsedPaths struct {
+	collapsedPaths *set.Set[string]
+}
+
+func NewCollapsedPaths() *CollapsedPaths {
+	return &CollapsedPaths{
+		collapsedPaths: set.New[string](),
+	}
+}
+
+func (self *CollapsedPaths) ExpandToPath(path string) {
 	// need every directory along the way
 	splitPath := split(path)
 	for i := range splitPath {
 		dir := join(splitPath[0 : i+1])
-		cp[dir] = false
+		self.collapsedPaths.Remove(dir)
 	}
 }
 
-func (cp CollapsedPaths) IsCollapsed(path string) bool {
-	return cp[path]
+func (self *CollapsedPaths) IsCollapsed(path string) bool {
+	return self.collapsedPaths.Includes(path)
+}
+
+func (self *CollapsedPaths) Collapse(path string) {
+	self.collapsedPaths.Add(path)
 }
 
-func (cp CollapsedPaths) ToggleCollapsed(path string) {
-	cp[path] = !cp[path]
+func (self *CollapsedPaths) ToggleCollapsed(path string) {
+	if self.collapsedPaths.Includes(path) {
+		self.collapsedPaths.Remove(path)
+	} else {
+		self.collapsedPaths.Add(path)
+	}
 }
diff --git a/pkg/gui/filetree/commit_file_node.go b/pkg/gui/filetree/commit_file_node.go
index a8f7d0a95..ac2057da5 100644
--- a/pkg/gui/filetree/commit_file_node.go
+++ b/pkg/gui/filetree/commit_file_node.go
@@ -100,7 +100,7 @@ func (s *CommitFileNode) EveryFile(test func(file *models.CommitFile) bool) bool
 	})
 }
 
-func (n *CommitFileNode) Flatten(collapsedPaths map[string]bool) []*CommitFileNode {
+func (n *CommitFileNode) Flatten(collapsedPaths *CollapsedPaths) []*CommitFileNode {
 	results := flatten(n, collapsedPaths)
 	nodes := make([]*CommitFileNode, len(results))
 	for i, result := range results {
@@ -110,7 +110,7 @@ func (n *CommitFileNode) Flatten(collapsedPaths map[string]bool) []*CommitFileNo
 	return nodes
 }
 
-func (node *CommitFileNode) GetNodeAtIndex(index int, collapsedPaths map[string]bool) *CommitFileNode {
+func (node *CommitFileNode) GetNodeAtIndex(index int, collapsedPaths *CollapsedPaths) *CommitFileNode {
 	if node == nil {
 		return nil
 	}
@@ -124,11 +124,11 @@ func (node *CommitFileNode) GetNodeAtIndex(index int, collapsedPaths map[string]
 	return result.(*CommitFileNode)
 }
 
-func (node *CommitFileNode) GetIndexForPath(path string, collapsedPaths map[string]bool) (int, bool) {
+func (node *CommitFileNode) GetIndexForPath(path string, collapsedPaths *CollapsedPaths) (int, bool) {
 	return getIndexForPath(node, path, collapsedPaths)
 }
 
-func (node *CommitFileNode) Size(collapsedPaths map[string]bool) int {
+func (node *CommitFileNode) Size(collapsedPaths *CollapsedPaths) int {
 	if node == nil {
 		return 0
 	}
diff --git a/pkg/gui/filetree/commit_file_tree.go b/pkg/gui/filetree/commit_file_tree.go
index 055e273f3..e539c9dea 100644
--- a/pkg/gui/filetree/commit_file_tree.go
+++ b/pkg/gui/filetree/commit_file_tree.go
@@ -19,7 +19,7 @@ type CommitFileTree struct {
 	tree           *CommitFileNode
 	showTree       bool
 	log            *logrus.Entry
-	collapsedPaths CollapsedPaths
+	collapsedPaths *CollapsedPaths
 }
 
 var _ ICommitFileTree = &CommitFileTree{}
@@ -29,7 +29,7 @@ func NewCommitFileTree(getFiles func() []*models.CommitFile, log *logrus.Entry,
 		getFiles:       getFiles,
 		log:            log,
 		showTree:       showTree,
-		collapsedPaths: CollapsedPaths{},
+		collapsedPaths: NewCollapsedPaths(),
 	}
 }
 
@@ -88,7 +88,7 @@ func (self *CommitFileTree) Tree() INode {
 	return self.tree
 }
 
-func (self *CommitFileTree) CollapsedPaths() CollapsedPaths {
+func (self *CommitFileTree) CollapsedPaths() *CollapsedPaths {
 	return self.collapsedPaths
 }
 
diff --git a/pkg/gui/filetree/file_node.go b/pkg/gui/filetree/file_node.go
index 841f723fc..e73504321 100644
--- a/pkg/gui/filetree/file_node.go
+++ b/pkg/gui/filetree/file_node.go
@@ -87,7 +87,7 @@ func (s *FileNode) Any(test func(node *FileNode) bool) bool {
 	})
 }
 
-func (n *FileNode) Flatten(collapsedPaths map[string]bool) []*FileNode {
+func (n *FileNode) Flatten(collapsedPaths *CollapsedPaths) []*FileNode {
 	results := flatten(n, collapsedPaths)
 	nodes := make([]*FileNode, len(results))
 	for i, result := range results {
@@ -97,7 +97,7 @@ func (n *FileNode) Flatten(collapsedPaths map[string]bool) []*FileNode {
 	return nodes
 }
 
-func (node *FileNode) GetNodeAtIndex(index int, collapsedPaths map[string]bool) *FileNode {
+func (node *FileNode) GetNodeAtIndex(index int, collapsedPaths *CollapsedPaths) *FileNode {
 	if node == nil {
 		return nil
 	}
@@ -111,11 +111,11 @@ func (node *FileNode) GetNodeAtIndex(index int, collapsedPaths map[string]bool)
 	return result.(*FileNode)
 }
 
-func (node *FileNode) GetIndexForPath(path string, collapsedPaths map[string]bool) (int, bool) {
+func (node *FileNode) GetIndexForPath(path string, collapsedPaths *CollapsedPaths) (int, bool) {
 	return getIndexForPath(node, path, collapsedPaths)
 }
 
-func (node *FileNode) Size(collapsedPaths map[string]bool) int {
+func (node *FileNode) Size(collapsedPaths *CollapsedPaths) int {
 	if node == nil {
 		return 0
 	}
diff --git a/pkg/gui/filetree/file_tree.go b/pkg/gui/filetree/file_tree.go
index 0d0524470..47d7f32f2 100644
--- a/pkg/gui/filetree/file_tree.go
+++ b/pkg/gui/filetree/file_tree.go
@@ -27,7 +27,7 @@ type ITree interface {
 	IsCollapsed(path string) bool
 	ToggleCollapsed(path string)
 	Tree() INode
-	CollapsedPaths() CollapsedPaths
+	CollapsedPaths() *CollapsedPaths
 }
 
 type IFileTree interface {
@@ -48,7 +48,7 @@ type FileTree struct {
 	showTree       bool
 	log            *logrus.Entry
 	filter         FileTreeDisplayFilter
-	collapsedPaths CollapsedPaths
+	collapsedPaths *CollapsedPaths
 }
 
 func NewFileTree(getFiles func() []*models.File, log *logrus.Entry, showTree bool) *FileTree {
@@ -57,7 +57,7 @@ func NewFileTree(getFiles func() []*models.File, log *logrus.Entry, showTree boo
 		log:            log,
 		showTree:       showTree,
 		filter:         DisplayAll,
-		collapsedPaths: CollapsedPaths{},
+		collapsedPaths: NewCollapsedPaths(),
 	}
 }
 
@@ -164,7 +164,7 @@ func (self *FileTree) Tree() INode {
 	return self.tree
 }
 
-func (self *FileTree) CollapsedPaths() CollapsedPaths {
+func (self *FileTree) CollapsedPaths() *CollapsedPaths {
 	return self.collapsedPaths
 }
 
diff --git a/pkg/gui/filetree/inode.go b/pkg/gui/filetree/inode.go
index 7d9035fe3..7c8b9fb75 100644
--- a/pkg/gui/filetree/inode.go
+++ b/pkg/gui/filetree/inode.go
@@ -90,11 +90,11 @@ func every(node INode, test func(INode) bool) bool {
 	return true
 }
 
-func flatten(node INode, collapsedPaths map[string]bool) []INode {
+func flatten(node INode, collapsedPaths *CollapsedPaths) []INode {
 	result := []INode{}
 	result = append(result, node)
 
-	if !collapsedPaths[node.GetPath()] {
+	if !collapsedPaths.IsCollapsed(node.GetPath()) {
 		for _, child := range node.GetChildren() {
 			result = append(result, flatten(child, collapsedPaths)...)
 		}
@@ -103,20 +103,20 @@ func flatten(node INode, collapsedPaths map[string]bool) []INode {
 	return result
 }
 
-func getNodeAtIndex(node INode, index int, collapsedPaths map[string]bool) INode {
+func getNodeAtIndex(node INode, index int, collapsedPaths *CollapsedPaths) INode {
 	foundNode, _ := getNodeAtIndexAux(node, index, collapsedPaths)
 
 	return foundNode
 }
 
-func getNodeAtIndexAux(node INode, index int, collapsedPaths map[string]bool) (INode, int) {
+func getNodeAtIndexAux(node INode, index int, collapsedPaths *CollapsedPaths) (INode, int) {
 	offset := 1
 
 	if index == 0 {
 		return node, offset
 	}
 
-	if !collapsedPaths[node.GetPath()] {
+	if !collapsedPaths.IsCollapsed(node.GetPath()) {
 		for _, child := range node.GetChildren() {
 			foundNode, offsetChange := getNodeAtIndexAux(child, index-offset, collapsedPaths)
 			offset += offsetChange
@@ -129,14 +129,14 @@ func getNodeAtIndexAux(node INode, index int, collapsedPaths map[string]bool) (I
 	return nil, offset
 }
 
-func getIndexForPath(node INode, path string, collapsedPaths map[string]bool) (int, bool) {
+func getIndexForPath(node INode, path string, collapsedPaths *CollapsedPaths) (int, bool) {
 	offset := 0
 
 	if node.GetPath() == path {
 		return offset, true
 	}
 
-	if !collapsedPaths[node.GetPath()] {
+	if !collapsedPaths.IsCollapsed(node.GetPath()) {
 		for _, child := range node.GetChildren() {
 			offsetChange, found := getIndexForPath(child, path, collapsedPaths)
 			offset += offsetChange + 1
@@ -149,10 +149,10 @@ func getIndexForPath(node INode, path string, collapsedPaths map[string]bool) (i
 	return offset, false
 }
 
-func size(node INode, collapsedPaths map[string]bool) int {
+func size(node INode, collapsedPaths *CollapsedPaths) int {
 	output := 1
 
-	if !collapsedPaths[node.GetPath()] {
+	if !collapsedPaths.IsCollapsed(node.GetPath()) {
 		for _, child := range node.GetChildren() {
 			output += size(child, collapsedPaths)
 		}