utils/tree/tree.go

package tree

type Node[T any] struct {
	Value    T         `json:"value"`
	Children []Node[T] `json:"children"`
}

func SetChildren[T any](node *Node[T], getChildValues func(node *Node[T]) ([]T, error)) error {
	childValues, err := getChildValues(node)
	if err != nil {
		return err
	}
	if len(childValues) == 0 {
		node.Children = nil
		return nil
	}

	childNodes := make([]Node[T], 0)
	for _, childValue := range childValues {
		childNode := Node[T]{
			Value: childValue,
		}
		err := SetChildren[T](&childNode, getChildValues)
		if err != nil {
			return err
		}
		childNodes = append(childNodes, childNode)
	}
	node.Children = childNodes
	return nil
}

func QueryTree[I, V any](
	input I,
	getRoots func(inputMsg I) ([]V, error),
	getChildren func(e *V) ([]V, error),
) ([]Node[V], error) {
	//顶部
	nRootValues, err := getRoots(input)
	if err != nil {
		return nil, err
	}

	//里面
	rootNodeSlice := make([]Node[V], 0)
	for _, nRootValue := range nRootValues {
		nRoot := Node[V]{
			Value: nRootValue,
		}
		err := SetChildren[V](&nRoot, func(node *Node[V]) ([]V, error) {
			nChildValues, err := getChildren(&node.Value)
			if err != nil {
				return nil, err
			}
			if len(nChildValues) == 0 {
				node.Children = nil
				return nil, nil
			}
			result := make([]V, 0)
			for _, nChildValue := range nChildValues {
				result = append(result, nChildValue)
			}
			return result, nil
		})

		if err != nil {
			return nil, err
		}
		rootNodeSlice = append(rootNodeSlice, nRoot)
	}

	return rootNodeSlice, nil
}