use serde_json::{Map, Value};
use teaql_tool_core::{Result, TeaQLToolError};
pub struct TreeTool;
impl TreeTool {
pub fn new() -> Self {
Self
}
pub fn build(
&self,
flat_array: Value,
id_field: &str,
parent_id_field: &str,
children_field: &str,
root_parent_id: Value,
) -> Result<Value> {
let arr = flat_array.as_array().ok_or_else(|| {
TeaQLToolError::InvalidArgument("flat_array must be a JSON array".to_string())
})?;
let mut items: Vec<Map<String, Value>> = Vec::new();
for item in arr {
if let Some(obj) = item.as_object() {
items.push(obj.clone());
} else {
return Err(TeaQLToolError::InvalidArgument(
"All elements in flat_array must be objects".to_string(),
));
}
}
let mut forest = Vec::new();
let mut all_nodes = items.clone();
fn build_children(
parent_id: &Value,
nodes: &[Map<String, Value>],
id_f: &str,
pid_f: &str,
child_f: &str,
) -> Vec<Value> {
let mut children = Vec::new();
for node in nodes {
if let Some(pid) = node.get(pid_f) {
if pid == parent_id {
let mut cloned_node = node.clone();
let current_id = node.get(id_f).unwrap_or(&Value::Null);
let sub_children = build_children(current_id, nodes, id_f, pid_f, child_f);
if !sub_children.is_empty() {
cloned_node.insert(child_f.to_string(), Value::Array(sub_children));
}
children.push(Value::Object(cloned_node));
}
}
}
children
}
forest = build_children(&root_parent_id, &all_nodes, id_field, parent_id_field, children_field);
Ok(Value::Array(forest))
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_tree_build() {
let tool = TreeTool::new();
let flat = json!([
{"id": 1, "pid": 0, "name": "Root"},
{"id": 2, "pid": 1, "name": "Child 1"},
{"id": 3, "pid": 1, "name": "Child 2"},
{"id": 4, "pid": 2, "name": "Grandchild"}
]);
let tree = tool.build(flat, "id", "pid", "children", json!(0)).unwrap();
let arr = tree.as_array().unwrap();
assert_eq!(arr.len(), 1);
let root = arr[0].as_object().unwrap();
assert_eq!(root.get("name").unwrap().as_str().unwrap(), "Root");
let children = root.get("children").unwrap().as_array().unwrap();
assert_eq!(children.len(), 2);
}
}