1use serde_json::{Map, Value};
2use teaql_tool_core::{Result, TeaQLToolError};
3
4pub struct TreeTool;
6
7impl TreeTool {
8 pub fn new() -> Self {
9 Self
10 }
11
12 pub fn build(
19 &self,
20 flat_array: Value,
21 id_field: &str,
22 parent_id_field: &str,
23 children_field: &str,
24 root_parent_id: Value,
25 ) -> Result<Value> {
26 let arr = flat_array.as_array().ok_or_else(|| {
27 TeaQLToolError::InvalidArgument("flat_array must be a JSON array".to_string())
28 })?;
29
30 let mut items: Vec<Map<String, Value>> = Vec::new();
32 for item in arr {
33 if let Some(obj) = item.as_object() {
34 items.push(obj.clone());
35 } else {
36 return Err(TeaQLToolError::InvalidArgument(
37 "All elements in flat_array must be objects".to_string(),
38 ));
39 }
40 }
41
42 let mut forest = Vec::new();
43
44 let mut all_nodes = items.clone();
50
51 fn build_children(
52 parent_id: &Value,
53 nodes: &[Map<String, Value>],
54 id_f: &str,
55 pid_f: &str,
56 child_f: &str,
57 ) -> Vec<Value> {
58 let mut children = Vec::new();
59 for node in nodes {
60 if let Some(pid) = node.get(pid_f) {
61 if pid == parent_id {
62 let mut cloned_node = node.clone();
63 let current_id = node.get(id_f).unwrap_or(&Value::Null);
64 let sub_children = build_children(current_id, nodes, id_f, pid_f, child_f);
65 if !sub_children.is_empty() {
66 cloned_node.insert(child_f.to_string(), Value::Array(sub_children));
67 }
68 children.push(Value::Object(cloned_node));
69 }
70 }
71 }
72 children
73 }
74
75 forest = build_children(&root_parent_id, &all_nodes, id_field, parent_id_field, children_field);
76 Ok(Value::Array(forest))
77 }
78}
79
80#[cfg(test)]
81mod tests {
82 use super::*;
83 use serde_json::json;
84
85 #[test]
86 fn test_tree_build() {
87 let tool = TreeTool::new();
88 let flat = json!([
89 {"id": 1, "pid": 0, "name": "Root"},
90 {"id": 2, "pid": 1, "name": "Child 1"},
91 {"id": 3, "pid": 1, "name": "Child 2"},
92 {"id": 4, "pid": 2, "name": "Grandchild"}
93 ]);
94
95 let tree = tool.build(flat, "id", "pid", "children", json!(0)).unwrap();
96 let arr = tree.as_array().unwrap();
97 assert_eq!(arr.len(), 1);
98 let root = arr[0].as_object().unwrap();
99 assert_eq!(root.get("name").unwrap().as_str().unwrap(), "Root");
100
101 let children = root.get("children").unwrap().as_array().unwrap();
102 assert_eq!(children.len(), 2);
103 }
104}