package filetree import ( "path" "slices" "strings" "github.com/jesseduffield/lazygit/pkg/commands/models" "github.com/jesseduffield/lazygit/pkg/gui/types" "github.com/samber/lo" ) // Represents a file or directory in a file tree. type Node[T any] struct { // File will be nil if the node is a directory. File *T // If the node is a directory, Children contains the contents of the directory, // otherwise it's nil. Children []*Node[T] // path of the file/directory // private; use either GetPath() or GetInternalPath() to access path string // rather than render a tree as: // a/ // b/ // file.blah // // we instead render it as: // a/b/ // file.blah // This saves vertical space. The CompressionLevel of a node is equal to the // number of times a 'compression' like the above has happened, where two // nodes are squished into one. CompressionLevel int } var _ types.ListItem = &Node[models.File]{} func (self *Node[T]) IsFile() bool { return self.File != nil } func (self *Node[T]) GetFile() *T { return self.File } // This returns the logical path from the user's point of view. It is the // relative path from the root of the repository. // Use this for display, or when you want to perform some action on the path // (e.g. a git command). func (self *Node[T]) GetPath() string { return strings.TrimPrefix(self.path, "./") } // This returns the internal path from the tree's point of view. It's the same // as GetPath(), but prefixed with "./" for the root item. // Use this when interacting with the tree itself, e.g. when calling // ToggleCollapsed. func (self *Node[T]) GetInternalPath() string { return self.path } func (self *Node[T]) Sort() { self.SortChildren() for _, child := range self.Children { child.Sort() } } func (self *Node[T]) ForEachFile(cb func(*T) error) error { if self.IsFile() { if err := cb(self.File); err != nil { return err } } for _, child := range self.Children { if err := child.ForEachFile(cb); err != nil { return err } } return nil } func (self *Node[T]) SortChildren() { if self.IsFile() { return } children := slices.Clone(self.Children) slices.SortFunc(children, func(a, b *Node[T]) int { if !a.IsFile() && b.IsFile() { return -1 } if a.IsFile() && !b.IsFile() { return 1 } return strings.Compare(a.path, b.path) }) // TODO: think about making this in-place self.Children = children } func (self *Node[T]) Some(predicate func(*Node[T]) bool) bool { if predicate(self) { return true } for _, child := range self.Children { if child.Some(predicate) { return true } } return false } func (self *Node[T]) SomeFile(predicate func(*T) bool) bool { if self.IsFile() { if predicate(self.File) { return true } } else { for _, child := range self.Children { if child.SomeFile(predicate) { return true } } } return false } func (self *Node[T]) Every(predicate func(*Node[T]) bool) bool { if !predicate(self) { return false } for _, child := range self.Children { if !child.Every(predicate) { return false } } return true } func (self *Node[T]) EveryFile(predicate func(*T) bool) bool { if self.IsFile() { if !predicate(self.File) { return false } } else { for _, child := range self.Children { if !child.EveryFile(predicate) { return false } } } 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} if len(self.Children) > 0 && !collapsedPaths.IsCollapsed(self.path) { result = append(result, lo.FlatMap(self.Children, func(child *Node[T], _ int) []*Node[T] { return child.Flatten(collapsedPaths) })...) } return result } func (self *Node[T]) GetNodeAtIndex(index int, collapsedPaths *CollapsedPaths) *Node[T] { if self == nil { return nil } node, _ := self.getNodeAtIndexAux(index, collapsedPaths) return node } func (self *Node[T]) getNodeAtIndexAux(index int, collapsedPaths *CollapsedPaths) (*Node[T], int) { offset := 1 if index == 0 { return self, offset } if !collapsedPaths.IsCollapsed(self.path) { for _, child := range self.Children { foundNode, offsetChange := child.getNodeAtIndexAux(index-offset, collapsedPaths) offset += offsetChange if foundNode != nil { return foundNode, offset } } } return nil, offset } func (self *Node[T]) GetIndexForPath(path string, collapsedPaths *CollapsedPaths) (int, bool) { offset := 0 if self.path == path { return offset, true } if !collapsedPaths.IsCollapsed(self.path) { for _, child := range self.Children { offsetChange, found := child.GetIndexForPath(path, collapsedPaths) offset += offsetChange + 1 if found { return offset, true } } } return offset, false } func (self *Node[T]) Size(collapsedPaths *CollapsedPaths) int { if self == nil { return 0 } output := 1 if !collapsedPaths.IsCollapsed(self.path) { for _, child := range self.Children { output += child.Size(collapsedPaths) } } return output } func (self *Node[T]) Compress() { if self == nil { return } self.compressAux() } func (self *Node[T]) compressAux() *Node[T] { if self.IsFile() { return self } children := self.Children for i := range children { grandchildren := children[i].Children for len(grandchildren) == 1 && !grandchildren[0].IsFile() { grandchildren[0].CompressionLevel = children[i].CompressionLevel + 1 children[i] = grandchildren[0] grandchildren = children[i].Children } } for i := range children { children[i] = children[i].compressAux() } self.Children = children return self } func (self *Node[T]) GetPathsMatching(predicate func(*Node[T]) bool) []string { paths := []string{} if predicate(self) { paths = append(paths, self.GetPath()) } for _, child := range self.Children { paths = append(paths, child.GetPathsMatching(predicate)...) } return paths } func (self *Node[T]) GetFilePathsMatching(predicate func(*T) bool) []string { matchingFileNodes := lo.Filter(self.GetLeaves(), func(node *Node[T], _ int) bool { return predicate(node.File) }) return lo.Map(matchingFileNodes, func(node *Node[T], _ int) string { return node.GetPath() }) } func (self *Node[T]) GetLeaves() []*Node[T] { if self.IsFile() { return []*Node[T]{self} } return lo.FlatMap(self.Children, func(child *Node[T], _ int) []*Node[T] { return child.GetLeaves() }) } func (self *Node[T]) ID() string { return self.GetPath() } func (self *Node[T]) Description() string { return self.GetPath() } func (self *Node[T]) Name() string { return path.Base(self.path) }