|
@@ -0,0 +1,73 @@
|
|
|
|
+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
|
|
|
|
+}
|