Refactor FileTree to contain expansion and traversal logic

Fixes #4.

cmd/twf/main.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"fmt"
-	"strings"
 
 	"github.com/wvanlint/twf/internal/config"
 	"github.com/wvanlint/twf/internal/filetree"
@@ -48,13 +47,16 @@ func main() {
 	if err != nil {
 		panic(err)
 	}
+	tree.Expand()
 	state := state.State{
-		Root:       tree,
-		Cursor:     tree.AbsPath,
-		Expansions: map[string]bool{tree.AbsPath: true},
+		Root:   tree,
+		Cursor: tree,
 	}
 	if config.LocatePath != "" {
-		state.ChangeCursor(config.LocatePath)
+		located, _ := tree.FindPath(config.LocatePath)
+		if located != nil {
+			state.Cursor = located
+		}
 	}
 	views := []terminal.View{
 		views.NewTreeView(config, &state),
@@ -72,8 +74,8 @@ func main() {
 		panic(err)
 	}
 
-	if len(state.Selection) > 0 {
-		fmt.Println(strings.Join(state.Selection, "\n"))
+	for _, node := range state.Selection {
+		fmt.Println(node.AbsPath)
 	}
 	zap.L().Info("Stopping twf.")
 }

internal/config/config.go

@@ -111,6 +111,8 @@ func defaultKeybindings() Keybindings {
 	return map[string][]string{
 		(&term.Event{term.Rune, 'j'}).HashKey():      []string{"tree:next"},
 		(&term.Event{term.Rune, 'k'}).HashKey():      []string{"tree:prev"},
+		(&term.Event{term.Rune, 'h'}).HashKey():      []string{"tree:parent", "tree:close"},
+		(&term.Event{term.Rune, 'l'}).HashKey():      []string{"tree:open", "tree:next"},
 		(&term.Event{Symbol: term.CtrlJ}).HashKey():  []string{"preview:down"},
 		(&term.Event{Symbol: term.CtrlK}).HashKey():  []string{"preview:up"},
 		(&term.Event{term.Rune, 'o'}).HashKey():      []string{"tree:toggle"},

internal/filetree/filetree.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"sort"
 )
 
 type FileTree struct {
@@ -13,6 +14,7 @@ type FileTree struct {
 	parent         *FileTree
 	children       []*FileTree
 	childrenByName map[string]*FileTree
+	expanded       bool
 }
 
 func InitFileTree(p string) (*FileTree, error) {
@@ -52,6 +54,18 @@ func (t *FileTree) Parent() *FileTree {
 	return t.parent
 }
 
+func (t *FileTree) Expanded() bool {
+	return t.expanded
+}
+
+func (t *FileTree) Expand() {
+	t.expanded = true
+}
+
+func (t *FileTree) Collapse() {
+	t.expanded = false
+}
+
 func (t *FileTree) maybeLoadChildren() error {
 	if t.children != nil {
 		return nil
@@ -65,6 +79,7 @@ func (t *FileTree) maybeLoadChildren() error {
 	if err != nil {
 		return err
 	}
+	defer f.Close()
 	contents, err := f.Readdir(0)
 	if err != nil {
 		return err
@@ -88,11 +103,15 @@ func (t *FileTree) maybeLoadChildren() error {
 	return nil
 }
 
-func (t *FileTree) Children() ([]*FileTree, error) {
+func (t *FileTree) Children(order Order) ([]*FileTree, error) {
 	if err := t.maybeLoadChildren(); err != nil {
 		return nil, err
 	}
-	return append(t.children[:0:0], t.children...), nil
+	children := append(t.children[:0:0], t.children...)
+	if order != nil {
+		sort.Slice(children, order(children))
+	}
+	return children, nil
 }
 
 type PathNotFound struct {
@@ -120,7 +139,10 @@ func (t *FileTree) FindPath(origPath string) (*FileTree, error) {
 	currentNode := t
 	var ok bool
 	for _, part := range parts {
-		currentNode.maybeLoadChildren()
+		err := currentNode.maybeLoadChildren()
+		if err != nil {
+			return nil, err
+		}
 		currentNode, ok = currentNode.childrenByName[part]
 		if !ok {
 			return nil, PathNotFound{origPath}
@@ -129,21 +151,93 @@ func (t *FileTree) FindPath(origPath string) (*FileTree, error) {
 	return currentNode, nil
 }
 
-func (t *FileTree) Traverse(f func(*FileTree)) error {
-	queue := []*FileTree{t}
-	var item *FileTree
-	for len(queue) > 0 {
-		item, queue = queue[0], queue[1:]
-		f(item)
-		children, err := item.Children()
-		if err != nil {
-			return err
+func (t *FileTree) Traverse(visibleOnly bool, order Order, f func(*FileTree, int)) error {
+	type treeWithDepth struct {
+		tree  *FileTree
+		depth int
+	}
+	stack := []treeWithDepth{treeWithDepth{t, 0}}
+	for len(stack) > 0 {
+		var current treeWithDepth
+		current, stack = stack[len(stack)-1], stack[:len(stack)-1]
+		f(current.tree, current.depth)
+
+		if !visibleOnly || current.tree.Expanded() {
+			children, err := current.tree.Children(order)
+			if err != nil {
+				return err
+			}
+			for i := len(children) - 1; i >= 0; i-- {
+				stack = append(stack, treeWithDepth{children[i], current.depth + 1})
+			}
 		}
-		queue = append(queue, children...)
 	}
 	return nil
 }
 
+func (t *FileTree) Prev(visibleOnly bool, order Order) (*FileTree, error) {
+	if t.Parent() == nil {
+		return nil, nil
+	}
+	siblings, err := t.Parent().Children(order)
+	if err != nil {
+		return nil, err
+	}
+	if t == siblings[0] {
+		return t.Parent(), nil
+	}
+	var prevSibling *FileTree
+	for i, sibling := range siblings {
+		if sibling == t {
+			prevSibling = siblings[i-1]
+		}
+	}
+	node := prevSibling
+	for {
+		if !node.Expanded() && visibleOnly {
+			return node, nil
+		}
+		children, err := node.Children(order)
+		if err != nil {
+			return nil, err
+		}
+		if len(children) == 0 {
+			return node, nil
+		} else {
+			node = children[len(children)-1]
+		}
+	}
+	return nil, nil
+}
+
+func (t *FileTree) Next(visibleOnly bool, order Order) (*FileTree, error) {
+	if t.Expanded() || !visibleOnly {
+		children, err := t.Children(order)
+		if err != nil {
+			return nil, err
+		}
+		if len(children) > 0 {
+			return children[0], nil
+		}
+	}
+	node := t
+	for node.Parent() != nil {
+		siblings, err := node.Parent().Children(order)
+		if err != nil {
+			return nil, err
+		}
+		for i, sibling := range siblings {
+			if sibling == node && i < len(siblings)-1 {
+				return siblings[i+1], nil
+			}
+		}
+		node = node.Parent()
+	}
+	return nil, nil
+}
+
+type Order func([]*FileTree) func(i, j int) bool
+
 func ByTypeAndName(children []*FileTree) func(i, j int) bool {
 	return func(i, j int) bool {
 		if children[i].IsDir() != children[j].IsDir() {

internal/filetree/filetree_test.go

@@ -24,20 +24,20 @@ func TestIsDir(t *testing.T) {
 func TestChildren(t *testing.T) {
 	tree, err := InitFileTree("testdata")
 	assert.Nil(t, err)
-	children, err := tree.Children()
+	children, err := tree.Children(nil)
 	assert.Nil(t, err)
 	childrenNames := []string{}
 	for _, childTree := range children {
 		childrenNames = append(childrenNames, childTree.Name())
 	}
-	assert.ElementsMatch(t, []string{"dir", "a"}, childrenNames)
+	assert.ElementsMatch(t, []string{"dir1", "dir2", "a"}, childrenNames)
 }
 
 func TestFindPath(t *testing.T) {
 	root, err := InitFileTree("testdata")
 	assert.Nil(t, err)
 
-	node, err := root.FindPath("dir/b")
+	node, err := root.FindPath("dir1/b")
 	assert.Nil(t, err)
 	assert.Equal(t, "b", node.Name())
 
@@ -47,7 +47,7 @@ func TestFindPath(t *testing.T) {
 
 	wd, err := os.Getwd()
 	assert.Nil(t, err)
-	node, err = root.FindPath(filepath.Join(wd, "testdata/dir/b"))
+	node, err = root.FindPath(filepath.Join(wd, "testdata/dir1/b"))
 	assert.Nil(t, err)
 	assert.Equal(t, "b", node.Name())
 
@@ -60,23 +60,48 @@ func TestTraverse(t *testing.T) {
 	names := []string{}
 	root, err := InitFileTree("testdata")
 	assert.Nil(t, err)
-	err = root.Traverse(func(node *FileTree) {
+	err = root.Traverse(false, nil, func(node *FileTree, _ int) {
 		names = append(names, node.Name())
 	})
-	assert.ElementsMatch(t, []string{"testdata", "dir", "a", "b"}, names)
+	assert.ElementsMatch(t, []string{"testdata", "dir1", "dir2", "a", "b", "c"}, names)
 }
 
 func TestByTypeAndName(t *testing.T) {
 	nodes := []*FileTree{}
 	root, err := InitFileTree("testdata")
 	assert.Nil(t, err)
-	err = root.Traverse(func(node *FileTree) {
+	err = root.Traverse(false, nil, func(node *FileTree, _ int) {
 		nodes = append(nodes, node)
 	})
 	sort.Slice(nodes, ByTypeAndName(nodes))
 	names := []string{}
 	for _, node := range nodes {
 		names = append(names, node.Name())
 	}
-	assert.Equal(t, []string{"dir", "testdata", "a", "b"}, names)
+	assert.Equal(t, []string{"dir1", "dir2", "testdata", "a", "b", "c"}, names)
+}
+
+func TestPrevNext(t *testing.T) {
+	tree, err := InitFileTree("testdata")
+	assert.Nil(t, err)
+	names := []string{tree.Name()}
+	for {
+		next, err := tree.Next(false, ByTypeAndName)
+		assert.Nil(t, err)
+		if next == nil {
+			break
+		}
+		tree = next
+		names = append(names, tree.Name())
+	}
+	for {
+		prev, err := tree.Prev(false, ByTypeAndName)
+		assert.Nil(t, err)
+		if prev == nil {
+			break
+		}
+		tree = prev
+		names = append(names, tree.Name())
+	}
+	assert.Equal(t, []string{"testdata", "dir1", "b", "dir2", "c", "a", "c", "dir2", "b", "dir1", "testdata"}, names)
 }

internal/state/state.go

@@ -5,57 +5,7 @@ import (
 )
 
 type State struct {
-	Root       *filetree.FileTree
-	Cursor     string
-	Expansions map[string]bool
-	Selection  []string
-}
-
-func (s *State) ChangeCursor(path string) {
-	node, err := s.Root.FindPath(path)
-	if err != nil {
-		return
-	}
-	s.Cursor = node.AbsPath
-	for node != s.Root {
-		node = node.Parent()
-		s.SetExpansion(node.AbsPath, true)
-	}
-}
-
-func (s *State) MoveCursorToParent() {
-	node, _ := s.Root.FindPath(s.Cursor)
-	parent := node.Parent()
-	if parent != nil {
-		s.Cursor = parent.AbsPath
-	}
-}
-
-func (s *State) SetExpansion(path string, value bool) {
-	if value {
-		s.Expansions[path] = true
-	} else {
-		delete(s.Expansions, path)
-	}
-}
-
-func (s *State) ToggleExpansion(path string) {
-	value, _ := s.Expansions[path]
-	s.SetExpansion(path, !value)
-}
-
-func (s *State) SetExpansionAll(path string, value bool) {
-	tree, _ := s.Root.FindPath(path)
-	tree.Traverse(func(node *filetree.FileTree) {
-		s.SetExpansion(node.AbsPath, value)
-	})
-}
-
-func (s *State) ToggleExpansionAll(path string) {
-	value, _ := s.Expansions[path]
-	s.SetExpansionAll(path, !value)
-}
-
-func (s *State) AddSelection(path string) {
-	s.Selection = append(s.Selection, path)
+	Root      *filetree.FileTree
+	Cursor    *filetree.FileTree
+	Selection []*filetree.FileTree
 }