Skip to main content

pite_core/
tree.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2use std::collections::HashMap;
3
4use crate::error::{CoreError, Result};
5use crate::node::{Node, NodeDesc, NodeId};
6
7/// Hierarchical node tree managing node lifecycles and parent-child relationships.
8///
9/// Nodes are stored with ownership in the tree; removing or dropping a node subtree
10/// automatically cleans up associated child nodes and script attachments.
11#[derive(Debug, Default)]
12pub struct NodeTree {
13    nodes: HashMap<String, Node>,
14    order: Vec<String>,
15    root: Option<NodeId>,
16}
17
18impl NodeTree {
19    /// Creates an empty `NodeTree`.
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Inserts a node into the tree from a [`NodeDesc`].
25    ///
26    /// Validates node ID uniqueness and updates parent-child links. If no root has been
27    /// established and the node has no parent, it is designated as the root.
28    pub fn insert(&mut self, desc: NodeDesc) -> Result<()> {
29        if desc.id.is_empty() {
30            return Err(CoreError::EmptyNodeId);
31        }
32        if self.nodes.contains_key(desc.id.as_str()) {
33            return Err(CoreError::DuplicateNode(desc.id.to_string()));
34        }
35        if self.root.is_none() && desc.parent.is_none() {
36            self.root = Some(desc.id.clone());
37        }
38        let node = Node {
39            id: desc.id.clone(),
40            type_name: desc.type_name,
41            name: desc.name,
42            parent: desc.parent.clone(),
43            children: Vec::new(),
44            props: desc.props,
45            script: desc.script,
46        };
47        if let Some(parent_id) = &desc.parent {
48            if let Some(parent) = self.nodes.get_mut(parent_id.as_str()) {
49                parent.children.push(desc.id.clone());
50            }
51        }
52        self.order.push(desc.id.to_string());
53        self.nodes.insert(desc.id.to_string(), node);
54        Ok(())
55    }
56
57    pub fn get(&self, id: &NodeId) -> Option<&Node> {
58        self.nodes.get(id.as_str())
59    }
60
61    pub fn contains(&self, id: &NodeId) -> bool {
62        self.nodes.contains_key(id.as_str())
63    }
64
65    pub fn children_of(&self, id: &NodeId) -> Vec<NodeId> {
66        self.nodes
67            .get(id.as_str())
68            .map(|n| n.children.clone())
69            .unwrap_or_default()
70    }
71
72    pub fn get_mut(&mut self, id: &NodeId) -> Option<&mut Node> {
73        self.nodes.get_mut(id.as_str())
74    }
75
76    pub fn iter(&self) -> impl Iterator<Item = &Node> {
77        self.order.iter().filter_map(|k| self.nodes.get(k))
78    }
79
80    pub fn root(&self) -> Option<&NodeId> {
81        self.root.as_ref()
82    }
83
84    pub fn len(&self) -> usize {
85        self.nodes.len()
86    }
87
88    pub fn is_empty(&self) -> bool {
89        self.nodes.is_empty()
90    }
91
92    pub fn remove(&mut self, id: &NodeId) -> Result<()> {
93        let node = self
94            .nodes
95            .get(id.as_str())
96            .ok_or_else(|| CoreError::UnknownNode { id: id.to_string() })?;
97        if self.root.as_ref() == Some(id) {
98            return Err(CoreError::RemoveRoot);
99        }
100        let parent = node.parent.clone();
101        let mut doomed = vec![id.to_string()];
102        let mut i = 0;
103        while i < doomed.len() {
104            let key = doomed[i].clone();
105            if let Some(n) = self.nodes.get(&key) {
106                doomed.extend(n.children.iter().map(ToString::to_string));
107            }
108            i += 1;
109        }
110        if let Some(parent_id) = parent {
111            if let Some(p) = self.nodes.get_mut(parent_id.as_str()) {
112                p.children.retain(|c| c.as_str() != id.as_str());
113            }
114        }
115        for key in doomed {
116            self.nodes.remove(&key);
117            self.order.retain(|k| k != &key);
118        }
119        Ok(())
120    }
121
122    pub fn reparent(&mut self, id: &NodeId, new_parent: Option<NodeId>) -> Result<()> {
123        if !self.nodes.contains_key(id.as_str()) {
124            return Err(CoreError::UnknownNode { id: id.to_string() });
125        }
126        if self.root.as_ref() == Some(id) && new_parent.is_some() {
127            return Err(CoreError::RemoveRoot);
128        }
129        if let Some(parent_id) = &new_parent {
130            if !self.nodes.contains_key(parent_id.as_str()) {
131                return Err(CoreError::MissingParent {
132                    child: id.to_string(),
133                    parent: parent_id.to_string(),
134                });
135            }
136            let mut cursor = Some(parent_id.clone());
137            while let Some(current) = cursor {
138                if &current == id {
139                    return Err(CoreError::Cycle {
140                        child: id.to_string(),
141                        parent: parent_id.to_string(),
142                    });
143                }
144                cursor = self
145                    .nodes
146                    .get(current.as_str())
147                    .and_then(|n| n.parent.clone());
148            }
149        }
150        let old_parent = self.nodes.get(id.as_str()).and_then(|n| n.parent.clone());
151        if let Some(old_id) = old_parent {
152            if let Some(old) = self.nodes.get_mut(old_id.as_str()) {
153                old.children.retain(|c| c.as_str() != id.as_str());
154            }
155        }
156        if let Some(parent_id) = &new_parent {
157            if let Some(new) = self.nodes.get_mut(parent_id.as_str()) {
158                new.children.push(id.clone());
159            }
160        }
161        if let Some(node) = self.nodes.get_mut(id.as_str()) {
162            node.parent = new_parent;
163        }
164        Ok(())
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    use crate::NodeDesc;
172
173    fn desc(id: &str, parent: Option<&str>) -> NodeDesc {
174        let mut d = NodeDesc::new(NodeId::from(id.to_string()), "Node2D");
175        d.name = id.to_string();
176        d.parent = parent.map(|p| NodeId::from(p.to_string()));
177        d
178    }
179
180    fn sample_tree() -> NodeTree {
181        let mut tree = NodeTree::new();
182        tree.insert(desc("root", None)).unwrap();
183        tree.insert(desc("a", Some("root"))).unwrap();
184        tree.insert(desc("b", Some("a"))).unwrap();
185        tree
186    }
187
188    #[test]
189    fn remove_drops_subtree_and_unlinks_parent() {
190        let mut tree = sample_tree();
191        tree.remove(&NodeId::from("a".to_string())).unwrap();
192        assert_eq!(tree.len(), 1);
193        assert!(tree.get(&NodeId::from("b".to_string())).is_none());
194        assert!(tree
195            .children_of(&NodeId::from("root".to_string()))
196            .is_empty());
197    }
198
199    #[test]
200    fn remove_root_is_forbidden() {
201        let mut tree = sample_tree();
202        let err = tree.remove(&NodeId::from("root".to_string())).unwrap_err();
203        assert!(matches!(err, CoreError::RemoveRoot));
204        assert_eq!(tree.len(), 3);
205    }
206
207    #[test]
208    fn remove_unknown_errors() {
209        let mut tree = sample_tree();
210        let err = tree.remove(&NodeId::from("ghost".to_string())).unwrap_err();
211        assert!(matches!(err, CoreError::UnknownNode { .. }));
212    }
213
214    #[test]
215    fn reparent_moves_linkage() {
216        let mut tree = sample_tree();
217        tree.insert(desc("c", Some("root"))).unwrap();
218        tree.reparent(
219            &NodeId::from("b".to_string()),
220            Some(NodeId::from("c".to_string())),
221        )
222        .unwrap();
223        let b = tree.get(&NodeId::from("b".to_string())).unwrap();
224        assert_eq!(b.parent, Some(NodeId::from("c".to_string())));
225        assert!(tree.children_of(&NodeId::from("a".to_string())).is_empty());
226        assert_eq!(tree.children_of(&NodeId::from("c".to_string())).len(), 1);
227    }
228
229    #[test]
230    fn reparent_missing_parent_errors() {
231        let mut tree = sample_tree();
232        let err = tree
233            .reparent(
234                &NodeId::from("b".to_string()),
235                Some(NodeId::from("ghost".to_string())),
236            )
237            .unwrap_err();
238        assert!(matches!(err, CoreError::MissingParent { .. }));
239    }
240
241    #[test]
242    fn reparent_into_descendant_is_cycle() {
243        let mut tree = sample_tree();
244        let err = tree
245            .reparent(
246                &NodeId::from("a".to_string()),
247                Some(NodeId::from("b".to_string())),
248            )
249            .unwrap_err();
250        assert!(matches!(err, CoreError::Cycle { .. }));
251    }
252}