1use std::collections::HashMap;
3
4use crate::error::{CoreError, Result};
5use crate::node::{Node, NodeDesc, NodeId};
6
7#[derive(Debug, Default)]
12pub struct NodeTree {
13 nodes: HashMap<String, Node>,
14 order: Vec<String>,
15 root: Option<NodeId>,
16}
17
18impl NodeTree {
19 pub fn new() -> Self {
21 Self::default()
22 }
23
24 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 ¤t == 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}