feat: add support for alternative ordering strategies (#424)

8 files changed, 155 insertions, 63 deletions

diff --git a/cmd/root.go b/cmd/root.go
index 27412d9..5074bbe 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -86,6 +86,7 @@ func initConfig() {
 	// keybindings: filetree view
 	viper.SetDefault("keybinding.toggle-collapse-dir", "space")
 	viper.SetDefault("keybinding.toggle-collapse-all-dir", "ctrl+space")
+	viper.SetDefault("keybinding.toggle-sort-order", "ctrl+o")
 	viper.SetDefault("keybinding.toggle-filetree-attributes", "ctrl+b")
 	viper.SetDefault("keybinding.toggle-added-files", "ctrl+a")
 	viper.SetDefault("keybinding.toggle-removed-files", "ctrl+r")
diff --git a/dive/filetree/efficiency.go b/dive/filetree/efficiency.go
index f45d281..d1244d1 100644
--- a/dive/filetree/efficiency.go
+++ b/dive/filetree/efficiency.go
@@ -79,7 +79,7 @@ func Efficiency(trees []*FileTree) (float64, EfficiencySlice) {
 			}
 
 			if previousTreeNode.Data.FileInfo.IsDir {
-				err = previousTreeNode.VisitDepthChildFirst(sizer, nil)
+				err = previousTreeNode.VisitDepthChildFirst(sizer, nil, nil)
 				if err != nil {
 					logrus.Errorf("unable to propagate whiteout dir: %+v", err)
 					return err
diff --git a/dive/filetree/file_node.go b/dive/filetree/file_node.go
index f99bc32..2cb1bb7 100644
--- a/dive/filetree/file_node.go
+++ b/dive/filetree/file_node.go
@@ -3,7 +3,6 @@ package filetree
 import (
 	"archive/tar"
 	"fmt"
-	"sort"
 	"strings"
 
 	"github.com/dustin/go-humanize"
@@ -27,6 +26,7 @@ var diffTypeColor = map[DiffType]*color.Color{
 type FileNode struct {
 	Tree     *FileTree
 	Parent   *FileNode
+	Size     int64 // memoized total size of file or directory
 	Name     string
 	Data     NodeData
 	Children map[string]*FileNode
@@ -39,6 +39,7 @@ func NewNode(parent *FileNode, name string, data FileInfo) (node *FileNode) {
 	node.Name = name
 	node.Data = *NewNodeData()
 	node.Data.FileInfo = *data.Copy()
+	node.Size = -1 // signal lazy load later
 
 	node.Children = make(map[string]*FileNode)
 	node.Parent = parent
@@ -149,41 +150,49 @@ func (node *FileNode) MetadataString() string {
 	group := node.Data.FileInfo.Gid
 	userGroup := fmt.Sprintf("%d:%d", user, group)
 
+	// don't include file sizes of children that have been removed (unless the node in question is a removed dir,
+	// then show the accumulated size of removed files)
+	sizeBytes := node.GetSize()
+
+	size := humanize.Bytes(uint64(sizeBytes))
+
+	return diffTypeColor[node.Data.DiffType].Sprint(fmt.Sprintf(AttributeFormat, dir, fileMode, userGroup, size))
+}
+
+func (node *FileNode) GetSize() int64 {
+	if 0 <= node.Size {
+		return node.Size
+	}
 	var sizeBytes int64
 
 	if node.IsLeaf() {
 		sizeBytes = node.Data.FileInfo.Size
 	} else {
 		sizer := func(curNode *FileNode) error {
-			// don't include file sizes of children that have been removed (unless the node in question is a removed dir,
-			// then show the accumulated size of removed files)
+
 			if curNode.Data.DiffType != Removed || node.Data.DiffType == Removed {
 				sizeBytes += curNode.Data.FileInfo.Size
 			}
 			return nil
 		}
-
-		err := node.VisitDepthChildFirst(sizer, nil)
+		err := node.VisitDepthChildFirst(sizer, nil, nil)
 		if err != nil {
 			logrus.Errorf("unable to propagate node for metadata: %+v", err)
 		}
 	}
-
-	size := humanize.Bytes(uint64(sizeBytes))
-
-	return diffTypeColor[node.Data.DiffType].Sprint(fmt.Sprintf(AttributeFormat, dir, fileMode, userGroup, size))
+	node.Size = sizeBytes
+	return node.Size
 }
 
 // VisitDepthChildFirst iterates a tree depth-first (starting at this FileNode), evaluating the deepest depths first (visit on bubble up)
-func (node *FileNode) VisitDepthChildFirst(visitor Visitor, evaluator VisitEvaluator) error {
-	var keys []string
-	for key := range node.Children {
-		keys = append(keys, key)
+func (node *FileNode) VisitDepthChildFirst(visitor Visitor, evaluator VisitEvaluator, sorter OrderStrategy) error {
+	if sorter == nil {
+		sorter = GetSortOrderStrategy(ByName)
 	}
-	sort.Strings(keys)
+	keys := sorter.orderKeys(node.Children)
 	for _, name := range keys {
 		child := node.Children[name]
-		err := child.VisitDepthChildFirst(visitor, evaluator)
+		err := child.VisitDepthChildFirst(visitor, evaluator, sorter)
 		if err != nil {
 			return err
 		}
@@ -199,7 +208,7 @@ func (node *FileNode) VisitDepthChildFirst(visitor Visitor, evaluator VisitEvalu
 }
 
 // VisitDepthParentFirst iterates a tree depth-first (starting at this FileNode), evaluating the shallowest depths first (visit while sinking down)
-func (node *FileNode) VisitDepthParentFirst(visitor Visitor, evaluator VisitEvaluator) error {
+func (node *FileNode) VisitDepthParentFirst(visitor Visitor, evaluator VisitEvaluator, sorter OrderStrategy) error {
 	var err error
 
 	doVisit := evaluator != nil && evaluator(node) || evaluator == nil
@@ -216,14 +225,13 @@ func (node *FileNode) VisitDepthParentFirst(visitor Visitor, evaluator VisitEval
 		}
 	}
 
-	var keys []string
-	for key := range node.Children {
-		keys = append(keys, key)
+	if sorter == nil {
+		sorter = GetSortOrderStrategy(ByName)
 	}
-	sort.Strings(keys)
+	keys := sorter.orderKeys(node.Children)
 	for _, name := range keys {
 		child := node.Children[name]
-		err = child.VisitDepthParentFirst(visitor, evaluator)
+		err = child.VisitDepthParentFirst(visitor, evaluator, sorter)
 		if err != nil {
 			return err
 		}
diff --git a/dive/filetree/file_tree.go b/dive/filetree/file_tree.go
index ede52b9..b6b0668 100644
--- a/dive/filetree/file_tree.go
+++ b/dive/filetree/file_tree.go
@@ -3,7 +3,6 @@ package filetree
 import (
 	"fmt"
 	"path"
-	"sort"
 	"strings"
 
 	"github.com/google/uuid"
@@ -24,11 +23,12 @@ const (
 
 // FileTree represents a set of files, directories, and their relations.
 type FileTree struct {
-	Root     *FileNode
-	Size     int
-	FileSize uint64
-	Name     string
-	Id       uuid.UUID
+	Root      *FileNode
+	Size      int
+	FileSize  uint64
+	Name      string
+	Id        uuid.UUID
+	SortOrder SortOrder
 }
 
 // NewFileTree creates an empty FileTree
@@ -39,6 +39,7 @@ func NewFileTree() (tree *FileTree) {
 	tree.Root.Tree = tree
 	tree.Root.Children = make(map[string]*FileNode)
 	tree.Id = uuid.New()
+	tree.SortOrder = ByName
 	return tree
 }
 
@@ -67,12 +68,8 @@ func (tree *FileTree) renderStringTreeBetween(startRow, stopRow int, showAttribu
 		currentParams, paramsToVisit = paramsToVisit[0], paramsToVisit[1:]
 
 		// take note of the next nodes to visit later
-		var keys []string
-		for key := range currentParams.node.Children {
-			keys = append(keys, key)
-		}
-		// we should always visit nodes in order
-		sort.Strings(keys)
+		sorter := GetSortOrderStrategy(tree.SortOrder)
+		keys := sorter.orderKeys(currentParams.node.Children)
 
 		var childParams = make([]renderParams, 0)
 		for idx, name := range keys {
@@ -174,6 +171,7 @@ func (tree *FileTree) Copy() *FileTree {
 	newTree.Size = tree.Size
 	newTree.FileSize = tree.FileSize
 	newTree.Root = tree.Root.Copy(newTree.Root)
+	newTree.SortOrder = tree.SortOrder
 
 	// update the tree pointers
 	err := newTree.VisitDepthChildFirst(func(node *FileNode) error {
@@ -196,12 +194,14 @@ type VisitEvaluator func(*FileNode) bool
 
 // VisitDepthChildFirst iterates the given tree depth-first, evaluating the deepest depths first (visit on bubble up)
 func (tree *FileTree) VisitDepthChildFirst(visitor Visitor, evaluator VisitEvaluator) error {
-	return tree.Root.VisitDepthChildFirst(visitor, evaluator)
+	sorter := GetSortOrderStrategy(tree.SortOrder)
+	return tree.Root.VisitDepthChildFirst(visitor, evaluator, sorter)
 }
 
 // VisitDepthParentFirst iterates the given tree depth-first, evaluating the shallowest depths first (visit while sinking down)
 func (tree *FileTree) VisitDepthParentFirst(visitor Visitor, evaluator VisitEvaluator) error {
-	return tree.Root.VisitDepthParentFirst(visitor, evaluator)
+	sorter := GetSortOrderStrategy(tree.SortOrder)
+	return tree.Root.VisitDepthParentFirst(visitor, evaluator, sorter)
 }
 
 // Stack takes two trees and combines them together. This is done by "stacking" the given tree on top of the owning tree.
diff --git a/dive/filetree/order_strategy.go b/dive/filetree/order_strategy.go
new file mode 100644
index 0000000..1838dd8
--- /dev/null
+++ b/dive/filetree/order_strategy.go
@@ -0,0 +1,61 @@
+package filetree
+
+import (
+	"sort"
+)
+
+type SortOrder int
+
+const (
+	ByName = iota
+	BySizeDesc
+
+	NumSortOrderConventions
+)
+
+type OrderStrategy interface {
+	orderKeys(files map[string]*FileNode) []string
+}
+
+func GetSortOrderStrategy(sortOrder SortOrder) OrderStrategy {
+	switch sortOrder {
+	case ByName:
+		return orderByNameStrategy{}
+	case BySizeDesc:
+		return orderBySizeDescStrategy{}
+	}
+	return orderByNameStrategy{}
+}
+
+type orderByNameStrategy struct{}
+
+func (orderByNameStrategy) orderKeys(files map[string]*FileNode) []string {
+	var keys []string
+	for key := range files {
+		keys = append(keys, key)
+	}
+
+	sort.Strings(keys)
+
+	return keys
+}
+
+type orderBySizeDescStrategy struct{}
+
+func (orderBySizeDescStrategy) orderKeys(files map[string]*FileNode) []string {
+	var keys []string
+	for key := range files {
+		keys = append(keys, key)
+	}
+
+	sort.Slice(keys, func(i, j int) bool {
+		ki, kj := keys[i], keys[j]
+		ni, nj := files[ki], files[kj]
+		if ni.GetSize() == nj.GetSize() {
+			return ki < kj
+		}
+		return ni.GetSize() > nj.GetSize()
+	})
+
+	return keys
+}
diff --git a/runtime/ui/view/filetree.go b/runtime/ui/view/filetree.go
index 62439e9..e92be37 100644
--- a/runtime/ui/view/filetree.go
+++ b/runtime/ui/view/filetree.go
@@ -24,7 +24,7 @@ type FileTree struct {
 	gui    *gocui.Gui
 	view   *gocui.View
 	header *gocui.View
-	vm     *viewmodel.FileTree
+	vm     *viewmodel.FileTreeViewModel
 	title  string
 
 	filterRegex         *regexp.Regexp
@@ -99,6 +99,11 @@ func (v *FileTree) Setup(view, header *gocui.View) error {
 			Display:    "Collapse all dir",
 		},
 		{
+			ConfigKeys: []string{"keybinding.toggle-sort-order"},
+			OnAction:   v.toggleSortOrder,
+			Display:    "Toggle sort order",
+		},
+		{
 			ConfigKeys: []string{"keybinding.toggle-added-files"},
 			OnAction:   func() error { return v.toggleShowDiffType(filetree.Added) },
 			IsSelected: func() bool { return !v.vm.HiddenDiffTypes[filetree.Added] },
@@ -288,6 +293,16 @@ func (v *FileTree) toggleCollapseAll() error {
 	return v.Render()
 }
 
+func (v *FileTree) toggleSortOrder() error {
+	err := v.vm.ToggleSortOrder()
+	if err != nil {
+		return err
+	}
+	v.resetCursor()
+	_ = v.Update()
+	return v.Render()
+}
+
 func (v *FileTree) toggleWrapTree() error {
 	v.view.Wrap = !v.view.Wrap
 	return nil
diff --git a/runtime/ui/viewmodel/filetree.go b/runtime/ui/viewmodel/filetree.go
index d0172d4..8734eaf 100644
--- a/runtime/ui/viewmodel/filetree.go
+++ b/runtime/ui/viewmodel/filetree.go
@@ -16,7 +16,7 @@ import (
 
 // FileTreeViewModel holds the UI objects and data models for populating the right pane. Specifically the pane that
 // shows selected layer or aggregate file ASCII tree.
-type FileTree struct {
+type FileTreeViewModel struct {
 	ModelTree *filetree.FileTree
 	ViewTree  *filetree.FileTree
 	RefTrees  []*filetree.FileTree
@@ -39,8 +39,8 @@ type FileTree struct {
 }
 
 // NewFileTreeViewModel creates a new view object attached the the global [gocui] screen object.
-func NewFileTreeViewModel(tree *filetree.FileTree, refTrees []*filetree.FileTree, cache filetree.Comparer) (treeViewModel *FileTree, err error) {
-	treeViewModel = new(FileTree)
+func NewFileTreeViewModel(tree *filetree.FileTree, refTrees []*filetree.FileTree, cache filetree.Comparer) (treeViewModel *FileTreeViewModel, err error) {
+	treeViewModel = new(FileTreeViewModel)
 
 	// populate main fields
 	treeViewModel.ShowAttributes = viper.GetBool("filetree.show-attributes")
@@ -71,13 +71,13 @@ func NewFileTreeViewModel(tree *filetree.FileTree, refTrees []*filetree.FileTree
 }
 
 // Setup initializes the UI concerns within the context of a global [gocui] view object.
-func (vm *FileTree) Setup(lowerBound, height int) {
+func (vm *FileTreeViewModel) Setup(lowerBound, height int) {
 	vm.bufferIndexLowerBound = lowerBound
 	vm.refHeight = height
 }
 
 // height returns the current height and considers the header
-func (vm *FileTree) height() int {
+func (vm *FileTreeViewModel) height() int {
 	if vm.ShowAttributes {
 		return vm.refHeight - 1
 	}
@@ -85,24 +85,24 @@ func (vm *FileTree) height() int {
 }
 
 // bufferIndexUpperBound returns the current upper bounds for the view
-func (vm *FileTree) bufferIndexUpperBound() int {
+func (vm *FileTreeViewModel) bufferIndexUpperBound() int {
 	return vm.bufferIndexLowerBound + vm.height()
 }
 
 // IsVisible indicates if the file tree view pane is currently initialized
-func (vm *FileTree) IsVisible() bool {
+func (vm *FileTreeViewModel) IsVisible() bool {
 	return vm != nil
 }
 
 // ResetCursor moves the cursor back to the top of the buffer and translates to the top of the buffer.
-func (vm *FileTree) ResetCursor() {
+func (vm *FileTreeViewModel) ResetCursor() {
 	vm.TreeIndex = 0
 	vm.bufferIndex = 0
 	vm.bufferIndexLowerBound = 0
 }
 
 // SetTreeByLayer populates the view model by stacking the indicated image layer file trees.
-func (vm *FileTree) SetTreeByLayer(bottomTreeStart, bottomTreeStop, topTreeStart, topTreeStop int) error {
+func (vm *FileTreeViewModel) SetTreeByLayer(bottomTreeStart, bottomTreeStop, topTreeStart, topTreeStop int) error {
 	if topTreeStop > len(vm.RefTrees)-1 {
 		return fmt.Errorf("invalid layer index given: %d of %d", topTreeStop, len(vm.RefTrees)-1)
 	}
@@ -131,7 +131,7 @@ func (vm *FileTree) SetTreeByLayer(bottomTreeStart, bottomTreeStop, topTreeStart
 }
 
 // doCursorUp performs the internal view's buffer adjustments on cursor up. Note: this is independent of the gocui buffer.
-func (vm *FileTree) CursorUp() bool {
+func (vm *FileTreeViewModel) CursorUp() bool {
 	if vm.TreeIndex <= 0 {
 		return false
 	}
@@ -146,7 +146,7 @@ func (vm *FileTree) CursorUp() bool {
 }
 
 // doCursorDown performs the internal view's buffer adjustments on cursor down. Note: this is independent of the gocui buffer.
-func (vm *FileTree) CursorDown() bool {
+func (vm *FileTreeViewModel) CursorDown() bool {
 	if vm.TreeIndex >= vm.ModelTree.VisibleSize() {
 		return false
 	}
@@ -162,7 +162,7 @@ func (vm *FileTree) CursorDown() bool {
 }
 
 // CursorLeft moves the cursor up until we reach the Parent Node or top of the tree
-func (vm *FileTree) CursorLeft(filterRegex *regexp.Regexp) error {
+func (vm *FileTreeViewModel) CursorLeft(filterRegex *regexp.Regexp) error {
 	var visitor func(*filetree.FileNode) error
 	var evaluator func(*filetree.FileNode) bool
 	var dfsCounter, newIndex int
@@ -213,7 +213,7 @@ func (vm *FileTree) CursorLeft(filterRegex *regexp.Regexp) error {
 }
 
 // CursorRight descends into directory expanding it if needed
-func (vm *FileTree) CursorRight(filterRegex *regexp.Regexp) error {
+func (vm *FileTreeViewModel) CursorRight(filterRegex *regexp.Regexp) error {
 	node := vm.getAbsPositionNode(filterRegex)
 	if node == nil {
 		return nil
@@ -245,7 +245,7 @@ func (vm *FileTree) CursorRight(filterRegex *regexp.Regexp) error {
 }
 
 // PageDown moves to next page putting the cursor on top
-func (vm *FileTree) PageDown() error {
+func (vm *FileTreeViewModel) PageDown() error {
 	nextBufferIndexLowerBound := vm.bufferIndexLowerBound + vm.height()
 	nextBufferIndexUpperBound := nextBufferIndexLowerBound + vm.height()
 
@@ -271,7 +271,7 @@ func (vm *FileTree) PageDown() error {
 }
 
 // PageUp moves to previous page putting the cursor on top
-func (vm *FileTree) PageUp() error {
+func (vm *FileTreeViewModel) PageUp() error {
 	nextBufferIndexLowerBound := vm.bufferIndexLowerBound - vm.height()
 	nextBufferIndexUpperBound := nextBufferIndexLowerBound + vm.height()
 
@@ -296,7 +296,7 @@ func (vm *FileTree) PageUp() error {
 }
 
 // getAbsPositionNode determines the selected screen cursor's location in the file tree, returning the selected FileNode.
-func (vm *FileTree) getAbsPositionNode(filterRegex *regexp.Regexp) (node *filetree.FileNode) {
+func (vm *FileTreeViewModel) getAbsPositionNode(filterRegex *regexp.Regexp) (node *filetree.FileNode) {
 	var visitor func(*filetree.FileNode) error
 	var evaluator func(*filetree.FileNode) bool
 	var dfsCounter int
@@ -327,7 +327,7 @@ func (vm *FileTree) getAbsPositionNode(filterRegex *regexp.Regexp) (node *filetr
 }
 
 // ToggleCollapse will collapse/expand the selected FileNode.
-func (vm *FileTree) ToggleCollapse(filterRegex *regexp.Regexp) error {
+func (vm *FileTreeViewModel) ToggleCollapse(filterRegex *regexp.Regexp) error {
 	node := vm.getAbsPositionNode(filterRegex)
 	if node != nil && node.Data.FileInfo.IsDir {
 		node.Data.ViewInfo.Collapsed = !node.Data.ViewInfo.Collapsed
@@ -336,7 +336,7 @@ func (vm *FileTree) ToggleCollapse(filterRegex *regexp.Regexp) error {
 }
 
 // ToggleCollapseAll will collapse/expand the all directories.
-func (vm *FileTree) ToggleCollapseAll() error {
+func (vm *FileTreeViewModel) ToggleCollapseAll() error {
 	vm.CollapseAll = !vm.CollapseAll
 
 	visitor := func(curNode *filetree.FileNode) error {
@@ -356,7 +356,14 @@ func (vm *FileTree) ToggleCollapseAll() error {
 	return nil
 }
 
-func (vm *FileTree) ConstrainLayout() {
+// ToggleSortOrder will toggle the sort order in which files are displayed
+func (vm *FileTreeViewModel) ToggleSortOrder() error {
+	vm.ModelTree.SortOrder = (vm.ModelTree.SortOrder + 1) % filetree.NumSortOrderConventions
+
+	return nil
+}
+
+func (vm *FileTreeViewModel) ConstrainLayout() {
 	if !vm.constrainedRealEstate {
 		logrus.Debugf("constraining filetree layout")
 		vm.constrainedRealEstate = true
@@ -365,7 +372,7 @@ func (vm *FileTree) ConstrainLayout() {
 	}
 }
 
-func (vm *FileTree) ExpandLayout() {
+func (vm *FileTreeViewModel) ExpandLayout() {
 	if vm.constrainedRealEstate {
 		logrus.Debugf("expanding filetree layout")
 		vm.ShowAttributes = vm.unconstrainedShowAttributes
@@ -374,7 +381,7 @@ func (vm *FileTree) ExpandLayout() {
 }
 
 // ToggleCollapse will collapse/expand the selected FileNode.
-func (vm *FileTree) ToggleAttributes() error {
+func (vm *FileTreeViewModel) ToggleAttributes() error {
 	// ignore any attempt to show the attributes when the layout is constrained
 	if vm.constrainedRealEstate {
 		return nil
@@ -384,12 +391,12 @@ func (vm *FileTree) ToggleAttributes() error {
 }
 
 // ToggleShowDiffType will show/hide the selected DiffType in the filetree pane.
-func (vm *FileTree) ToggleShowDiffType(diffType filetree.DiffType) {
+func (vm *FileTreeViewModel) ToggleShowDiffType(diffType filetree.DiffType) {
 	vm.HiddenDiffTypes[diffType] = !vm.HiddenDiffTypes[diffType]
 }
 
 // Update refreshes the state objects for future rendering.
-func (vm *FileTree) Update(filterRegex *regexp.Regexp, width, height int) error {
+func (vm *FileTreeViewModel) Update(filterRegex *regexp.Regexp, width, height int) error {
 	vm.refWidth = width
 	vm.refHeight = height
 
@@ -437,7 +444,7 @@ func (vm *FileTree) Update(filterRegex *regexp.Regexp, width, height int) error
 }
 
 // Render flushes the state objects (file tree) to the pane.
-func (vm *FileTree) Render() error {
+func (vm *FileTreeViewModel) Render() error {
 	treeString := vm.ViewTree.StringBetween(vm.bufferIndexLowerBound, vm.bufferIndexUpperBound(), vm.ShowAttributes)
 	lines := strings.Split(treeString, "\n")
 
diff --git a/runtime/ui/viewmodel/filetree_test.go b/runtime/ui/viewmodel/filetree_test.go
index 668eb1e..f315b41 100644
--- a/runtime/ui/viewmodel/filetree_test.go
+++ b/runtime/ui/viewmodel/filetree_test.go
@@ -73,7 +73,7 @@ func assertTestData(t *testing.T, actualBytes []byte) {
 	helperCheckDiff(t, expectedBytes, actualBytes)
 }
 
-func initializeTestViewModel(t *testing.T) *FileTree {
+func initializeTestViewModel(t *testing.T) *FileTreeViewModel {
 	result := docker.TestAnalysisFromArchive(t, "../../../.data/test-docker-image.tar")
 
 	cache := filetree.NewComparer(result.RefTrees)
@@ -98,7 +98,7 @@ func initializeTestViewModel(t *testing.T) *FileTree {
 	return vm
 }
 
-func runTestCase(t *testing.T, vm *FileTree, width, height int, filterRegex *regexp.Regexp) {
+func runTestCase(t *testing.T, vm *FileTreeViewModel, width, height int, filterRegex *regexp.Regexp) {
 	err := vm.Update(filterRegex, width, height)
 	if err != nil {
 		t.Errorf("failed to update viewmodel: %v", err)