From 91ba4be66df0ae716c45cbfade53b7fede588435 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Wed, 7 Dec 2022 18:32:02 +0800
Subject: [PATCH] [add] utils.tree.Flatten process

---
 helper/array.go    | 11 ++++++++++
 helper/process.go  |  2 +-
 utils/process.go   |  5 +++++
 utils/tree/tree.go | 52 ++++++++++++++++++++++++++++++++++++++++++++++
 utils/tree_test.go | 32 ++++++++++++++++++++++++++++
 5 files changed, 101 insertions(+), 1 deletion(-)
 create mode 100644 utils/tree/tree.go
 create mode 100644 utils/tree_test.go

diff --git a/helper/array.go b/helper/array.go
index ebc29ef9..32f17a2d 100644
--- a/helper/array.go
+++ b/helper/array.go
@@ -162,15 +162,26 @@ func OfArrayPluckValue(any interface{}) ArrayPluckValue {
 
 // NewArrayTreeOption 创建配置
 func NewArrayTreeOption(option map[string]interface{}) ArrayTreeOption {
+
 	new := ArrayTreeOption{
 		Empty:    0,
 		Key:      "id",
 		Parent:   "parent",
 		Children: "children",
 	}
+
+	if v, ok := option["empty"]; ok {
+		new.Empty = v
+	}
+
 	if v, ok := option["parent"].(string); ok {
 		new.Parent = v
 	}
+
+	if v, ok := option["primary"].(string); ok {
+		new.Key = v
+	}
+
 	if v, ok := option["children"].(string); ok {
 		new.Children = v
 	}
diff --git a/helper/process.go b/helper/process.go
index e0dd4817..e5a2a999 100644
--- a/helper/process.go
+++ b/helper/process.go
@@ -17,7 +17,7 @@ func init() {
 	gou.RegisterProcessHandler("xiang.helper.ArrayKeep", ProcessArrayKeep)       // deprecated → utils.arr.Keep  @/utils/process.go
 	gou.RegisterProcessHandler("xiang.helper.ArrayTree", ProcessArrayTree)       // deprecated → utils.arr.Tree  @/utils/process.go
 	gou.RegisterProcessHandler("xiang.helper.ArrayUnique", ProcessArrayUnique)   // deprecated → utils.arr.Unique  @/utils/process.go
-	gou.RegisterProcessHandler("xiang.helper.ArrayMapSet", ProcessArrayMapSet)
+	gou.RegisterProcessHandler("xiang.helper.ArrayMapSet", ProcessArrayMapSet)   // deprecated → utils.arr.MapSet  @/utils/process.go
 
 	gou.RegisterProcessHandler("xiang.helper.MapKeys", ProcessMapKeys)         // deprecated → utils.map.Keys @/utils/process.go
 	gou.RegisterProcessHandler("xiang.helper.MapValues", ProcessMapValues)     // deprecated → utils.map.Values @/utils/process.go
diff --git a/utils/process.go b/utils/process.go
index bc068ec2..db9cd750 100644
--- a/utils/process.go
+++ b/utils/process.go
@@ -4,6 +4,7 @@ import (
 	"github.com/yaoapp/gou"
 	"github.com/yaoapp/yao/utils/datetime"
 	"github.com/yaoapp/yao/utils/str"
+	"github.com/yaoapp/yao/utils/tree"
 )
 
 func init() {
@@ -61,6 +62,10 @@ func init() {
 	gou.AliasProcess("xiang.helper.ArrayGet", "utils.arr.Get")
 	gou.AliasProcess("xiang.helper.ArrayColumn", "utils.arr.Column") // doc
 	gou.AliasProcess("xiang.helper.ArrayKeep", "utils.arr.Keep")
+	gou.AliasProcess("xiang.helper.ArrayMapSet", "utils.arr.MapSet")
+
+	// Tree
+	gou.RegisterProcessHandler("utils.tree.Flatten", tree.ProcessFlatten)
 
 	// Map
 	gou.AliasProcess("xiang.helper.MapGet", "utils.map.Get")
diff --git a/utils/tree/tree.go b/utils/tree/tree.go
new file mode 100644
index 00000000..5906f638
--- /dev/null
+++ b/utils/tree/tree.go
@@ -0,0 +1,52 @@
+package tree
+
+import (
+	"fmt"
+
+	"github.com/yaoapp/gou"
+)
+
+// ProcessFlatten utils.tree.Flatten cast to array
+func ProcessFlatten(process *gou.Process) interface{} {
+	process.ValidateArgNums(1)
+	array := process.ArgsArray(0)
+	option := process.ArgsMap(1, map[string]interface{}{"primary": "id", "children": "children", "parent": "parent"})
+	if _, has := option["primary"]; !has {
+		option["primary"] = "id"
+	}
+
+	if _, has := option["children"]; !has {
+		option["children"] = "children"
+	}
+
+	if _, has := option["parent"]; !has {
+		option["parent"] = "parent"
+	}
+
+	return flatten(array, option, nil)
+}
+
+func flatten(array []interface{}, option map[string]interface{}, id interface{}) []interface{} {
+
+	parent := fmt.Sprintf("%v", option["parent"])
+	primary := fmt.Sprintf("%v", option["primary"])
+	childrenField := fmt.Sprintf("%v", option["children"])
+	res := []interface{}{}
+	for _, v := range array {
+		row, ok := v.(map[string]interface{})
+		if !ok {
+			continue
+		}
+
+		row[parent] = id
+		children, ok := row[childrenField].([]interface{})
+		delete(row, childrenField)
+		res = append(res, row)
+
+		if ok {
+			res = append(res, flatten(children, option, row[primary])...)
+
+		}
+	}
+	return res
+}
diff --git a/utils/tree_test.go b/utils/tree_test.go
new file mode 100644
index 00000000..ccdf712a
--- /dev/null
+++ b/utils/tree_test.go
@@ -0,0 +1,32 @@
+package utils
+
+import (
+	"testing"
+
+	jsoniter "github.com/json-iterator/go"
+	"github.com/stretchr/testify/assert"
+	"github.com/yaoapp/gou"
+)
+
+func TestProcessTreeFlatten(t *testing.T) {
+
+	bytes := []byte(`[
+		{
+		  "id": 1,
+		  "parent": null,
+		  "children": [{ "children": [], "id": 5, "parent": 1 }]
+		},
+		{ "id": 2, "parent": null, "children": [] },
+		{ "id": 3, "parent": null, "children": [] }
+	  ]`)
+
+	var data interface{}
+	err := jsoniter.Unmarshal(bytes, &data)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rows := gou.NewProcess("utils.tree.Flatten", data, map[string]interface{}{"primary": "id", "children": "children", "parent": "parent"}).Run().([]interface{})
+	assert.Equal(t, 4, len(rows))
+	assert.Equal(t, float64(1), rows[1].(map[string]interface{})["parent"])
+}