save work

pkg/remote/fileshare/pathtree/pathtree.go

@@ -0,0 +1,93 @@
+package pathtree
+
+import (
+	"log"
+	"strings"
+)
+
+type WalkFunc func(path string, isLeaf bool) error
+
+type Tree struct {
+	Root      *Node
+	RootPath  string
+	nodes     map[string]*Node
+	delimiter string
+}
+
+type Node struct {
+	Children map[string]*Node
+}
+
+func (n *Node) Walk(curPath string, walkFunc WalkFunc, delimiter string) error {
+	if err := walkFunc(curPath, len(n.Children) == 0); err != nil {
+		return err
+	}
+	for name, child := range n.Children {
+		if err := child.Walk(curPath+delimiter+name, walkFunc, delimiter); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func NewTree(path string, delimiter string) *Tree {
+	if !strings.HasSuffix(path, delimiter) {
+		path += delimiter
+	}
+	return &Tree{
+		Root: &Node{
+			Children: make(map[string]*Node),
+		},
+		nodes:     make(map[string]*Node),
+		RootPath:  path,
+		delimiter: delimiter,
+	}
+}
+
+func (t *Tree) Add(path string) {
+	relativePath := strings.TrimPrefix(path, t.RootPath)
+	log.Printf("relativePath: %s", relativePath)
+
+	// If the path is not a child of the root path, ignore it
+	if relativePath == path {
+		return
+	}
+
+	// If the path is already in the tree, ignore it
+	if t.nodes[relativePath] != nil {
+		return
+	}
+
+	components := strings.Split(relativePath, t.delimiter)
+
+	// Quick check to see if the parent path is already in the tree, in which case we can skip the loop
+	if len(components) > 1 {
+		parentPath := strings.Join(components[:len(components)-1], t.delimiter)
+		log.Printf("parentPath: %s", parentPath)
+		if t.nodes[parentPath] != nil {
+			lastPathComponent := components[len(components)-1]
+			t.nodes[parentPath].Children[lastPathComponent] = &Node{
+				Children: make(map[string]*Node),
+			}
+			t.nodes[relativePath] = t.nodes[parentPath].Children[lastPathComponent]
+			return
+		}
+	}
+
+	currentNode := t.Root
+	for i, component := range components {
+		if _, ok := currentNode.Children[component]; !ok {
+			currentNode.Children[component] = &Node{
+				Children: make(map[string]*Node),
+			}
+			curPath := strings.Join(components[:i+1], t.delimiter)
+			log.Printf("curPath: %s", curPath)
+			t.nodes[curPath] = currentNode.Children[component]
+		}
+		currentNode = currentNode.Children[component]
+	}
+}
+
+func (t *Tree) Walk(walkFunc WalkFunc) error {
+	return t.Root.Walk(strings.TrimSuffix(t.RootPath, t.delimiter), walkFunc, t.delimiter)
+}

pkg/remote/fileshare/pathtree/pathtree_test.go

@@ -0,0 +1,117 @@
+package pathtree_test
+
+import (
+	"errors"
+	"log"
+	"testing"
+
+	"github.com/wavetermdev/waveterm/pkg/remote/fileshare/pathtree"
+)
+
+func TestAdd(t *testing.T) {
+	t.Parallel()
+
+	tree := initializeTree()
+
+	// Check that the tree has the expected structure
+	if len(tree.Root.Children) != 3 {
+		t.Errorf("expected 3 children, got %d", len(tree.Root.Children))
+	}
+
+	if len(tree.Root.Children["a"].Children) != 3 {
+		t.Errorf("expected 3 children, got %d", len(tree.Root.Children["a"].Children))
+	}
+
+	if len(tree.Root.Children["b"].Children) != 1 {
+		t.Errorf("expected 1 child, got %d", len(tree.Root.Children["b"].Children))
+	}
+
+	if len(tree.Root.Children["b"].Children["g"].Children) != 1 {
+		t.Errorf("expected 1 child, got %d", len(tree.Root.Children["b"].Children["g"].Children))
+	}
+
+	if len(tree.Root.Children["b"].Children["g"].Children["h"].Children) != 0 {
+		t.Errorf("expected 0 children, got %d", len(tree.Root.Children["b"].Children["g"].Children["h"].Children))
+	}
+
+	if len(tree.Root.Children["c"].Children) != 0 {
+		t.Errorf("expected 0 children, got %d", len(tree.Root.Children["c"].Children))
+	}
+
+	// Check that adding the same path again does not change the tree
+	tree.Add("root/a/d")
+	if len(tree.Root.Children["a"].Children) != 3 {
+		t.Errorf("expected 3 children, got %d", len(tree.Root.Children["a"].Children))
+	}
+
+	// Check that adding a path that is not a child of the root path does not change the tree
+	tree.Add("etc/passwd")
+	if len(tree.Root.Children) != 3 {
+		t.Errorf("expected 3 children, got %d", len(tree.Root.Children))
+	}
+}
+
+func TestWalk(t *testing.T) {
+	t.Parallel()
+
+	tree := initializeTree()
+
+	// Check that the tree traverses all nodes and identifies leaf nodes correctly
+	pathMap := make(map[string]bool)
+	err := tree.Walk(func(path string, isLeaf bool) error {
+		pathMap[path] = isLeaf
+		return nil
+	})
+
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	expectedPathMap := map[string]bool{
+		"root/":      false,
+		"root/a":     false,
+		"root/a/d":   true,
+		"root/a/e":   true,
+		"root/a/f":   true,
+		"root/b":     false,
+		"root/b/g":   false,
+		"root/b/g/h": true,
+		"root/c":     true,
+	}
+
+	log.Printf("pathMap: %v", pathMap)
+
+	for path, isLeaf := range expectedPathMap {
+		if pathMap[path] != isLeaf {
+			if isLeaf {
+				t.Errorf("expected %s to be a leaf", path)
+			} else {
+				t.Errorf("expected %s to not be a leaf", path)
+			}
+		}
+	}
+
+	expectedError := errors.New("test error")
+
+	// Check that the walk function returns an error if it is returned by the walk function
+	err = tree.Walk(func(path string, isLeaf bool) error {
+		return expectedError
+	})
+	if err != expectedError {
+		t.Errorf("expected error %v, got %v", expectedError, err)
+	}
+}
+
+func initializeTree() *pathtree.Tree {
+	tree := pathtree.NewTree("root/", "/")
+	tree.Add("root/a")
+	tree.Add("root/b")
+	tree.Add("root/c")
+	tree.Add("root/a/d")
+	tree.Add("root/a/e")
+	tree.Add("root/a/f")
+	tree.Add("root/b/g")
+	tree.Add("root/b/g/h")
+	log.Printf("tree: %v", tree)
+	return tree
+}