1use std::collections::HashMap;
3
4use crate::error::{CoreError, Result};
5use crate::node::{Node, NodeDesc, NodeId};
6
7#[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 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 ¤t == 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}