FsUtils/mountree/mnt_tree.go

package mountree

import (
	"strings"
	"sync"
)

type TreeNode[T PayloadType] struct {
	Path     []string
	Children map[string]*TreeNode[T]
	Parent   *TreeNode[T]
	Payload  *T
}

type Tree[T PayloadType] struct {
	root   *TreeNode[T]
	rwlock sync.RWMutex
}

func NewTree[T PayloadType]() *Tree[T] {
	tree := &Tree[T]{
		root: &TreeNode[T]{
			Path:     []string{},
			Children: make(map[string]*TreeNode[T]),
			Parent:   nil,
			Payload:  nil,
		},
		rwlock: sync.RWMutex{},
	}
	return tree
}

func (t *Tree[T]) search(
	pathSeq []string,
	found func(fnode *TreeNode[T], remainPath []string, pnode *TreeNode[T]) error,
	readonly bool,
) error {
	fspnode := t.root
	if readonly {
		defer t.rwlock.RUnlock()
		t.rwlock.RLock()
	} else {
		defer t.rwlock.Unlock()
		t.rwlock.Lock()
	}
	total := len(pathSeq)
	if total <= 0 {
		return found(t.root, []string{}, fspnode)
	}
	current := t.root
	for i, v := range pathSeq {
		if current.Children == nil || len(current.Children) == 0 {
			return found(current, pathSeq[i:], fspnode)
		}
		if child, ok := current.Children[v]; ok {
			current = child
			if current.Payload != nil {
				fspnode = current
			}
		} else {
			return found(current, pathSeq[i:], fspnode)
		}
	}
	return found(current, []string{}, fspnode)
}

func (t *Tree[T]) Mount(pathSeq []string, payload T, forceRemount bool) error {
	return t.search(
		pathSeq,
		func(fnode *TreeNode[T], remainPath []string, pnode *TreeNode[T]) error {
			if len(remainPath) == 0 {
				if fnode.Payload != nil && !forceRemount {
					return ErrMountPointAlreadyExists
				}
				fnode.Payload = &payload
				return nil
			}
			currentPath := fnode.Path
			currentNode := fnode
			for _, v := range remainPath {
				currentPath = append(currentPath, v)
				newNode := &TreeNode[T]{
					Path:     currentPath,
					Children: make(map[string]*TreeNode[T]),
					Parent:   currentNode,
					Payload:  nil,
				}
				currentNode.Children[v] = newNode
				currentNode = newNode
			}
			currentNode.Payload = &payload
			return nil
		},
		false,
	)
}

func (t *Tree[T]) Umount(pathSeq []string) error {
	if len(pathSeq) == 0 {
		t.root.Payload = nil
		return nil
	}
	return t.search(
		pathSeq,
		func(fnode *TreeNode[T], remainPath []string, pnode *TreeNode[T]) error {
			if len(remainPath) != 0 {
				return ErrMountPointNotExists
			}
			if len(fnode.Children) > 0 {
				fnode.Payload = nil
				return nil
			} else {
				pnode := fnode.Parent
				delete(pnode.Children, fnode.Path[len(fnode.Path)-1])
				return nil
			}
		},
		false,
	)
}

func (t *Tree[T]) GetPayload(pathSeq []string) (payload *T, RemainPath []string, rerr error) {
	payload = nil
	RemainPath = pathSeq
	rerr = t.search(
		pathSeq,
		func(fnode *TreeNode[T], remainPath []string, pnode *TreeNode[T]) error {
			if fnode.Payload == nil {
				if pnode.Payload == nil {
					return ErrNoAvailableMountPointForThisPath
				} else {
					payload = pnode.Payload
					xpl := len(pnode.Path)
					RemainPath = pathSeq[xpl:]
					return nil
				}
			} else {
				RemainPath = remainPath
				payload = fnode.Payload
				return nil
			}
		},
		true,
	)
	return
}

func (t *Tree[T]) ListAllMount(prt func(path string, payload T), sep string) {
	stack := []*TreeNode[T]{t.root}

	for len(stack) > 0 {
		node := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		if node.Payload != nil {
			prtPath := sep + strings.Join(node.Path, sep)
			prt(prtPath, *node.Payload)
		}

		for _, child := range node.Children {
			stack = append(stack, child)
		}
	}
}