Skip to main content

teaql_tool_std/
tree.rs

1use serde_json::{Map, Value};
2use teaql_tool_core::{Result, TeaQLToolError};
3
4/// Tree structure utility wrapper
5pub struct TreeTool;
6
7impl TreeTool {
8    pub fn new() -> Self {
9        Self
10    }
11
12    /// Build a nested tree from a flat JSON array of objects.
13    /// `flat_array`: A JSON array containing flat objects.
14    /// `id_field`: The property name acting as the unique identifier.
15    /// `parent_id_field`: The property name acting as the parent identifier.
16    /// `children_field`: The property name where children will be inserted.
17    /// `root_parent_id`: The parent ID value that identifies root elements (e.g. `0` or `null`).
18    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        // Extract items as Map
31        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        // Very basic O(N^2) builder for simplicity and avoiding borrow checker self-referencing.
45        // For enterprise sizes usually N is small (menus, org charts < 1000 items).
46        // A recursive approach or hashmap based approach is possible but requires Rc<RefCell> in Rust.
47        
48        // We will do a multi-pass to build the tree.
49        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}