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