1use crate::{
2 ecs::{
3 commands::{DespawnEntity, UniverseCommands},
4 components::Name,
5 life_cycle::EntityChanges,
6 Comp, Entity, Universe, WorldRef,
7 },
8 prefab::{Prefab, PrefabError, PrefabProxy},
9 state::StateToken,
10};
11use serde::{Deserialize, Serialize};
12use std::collections::HashMap;
13
14#[derive(Debug, Clone, PartialEq)]
15pub struct Parent(pub Entity);
16
17impl PrefabProxy<ParentPrefabProxy> for Parent {
18 fn from_proxy_with_extras(
19 proxy: ParentPrefabProxy,
20 named_entities: &HashMap<String, Entity>,
21 _: StateToken,
22 ) -> Result<Self, PrefabError> {
23 if let Some(entity) = named_entities.get(&proxy.0) {
24 Ok(Self(*entity))
25 } else {
26 Err(PrefabError::Custom(format!(
27 "Could not find entity named: {}",
28 proxy.0
29 )))
30 }
31 }
32}
33
34#[derive(Debug, Default, Serialize, Deserialize)]
35pub struct ParentPrefabProxy(pub String);
36
37impl Prefab for ParentPrefabProxy {}
38
39#[derive(Debug, Default)]
40pub struct Hierarchy {
41 roots: Vec<Entity>,
42 child_parent_relations: HashMap<Entity, Entity>,
43 parent_children_relations: HashMap<Entity, Vec<Entity>>,
44 entity_names_map: HashMap<Entity, String>,
45 name_entities_map: HashMap<String, Entity>,
46}
47
48impl Hierarchy {
49 pub fn roots(&self) -> impl Iterator<Item = Entity> + '_ {
50 self.roots.iter().copied()
51 }
52
53 pub fn parents(&self) -> impl Iterator<Item = Entity> + '_ {
54 self.parent_children_relations.keys().copied()
55 }
56
57 pub fn childs(&self) -> impl Iterator<Item = Entity> + '_ {
58 self.child_parent_relations.keys().copied()
59 }
60
61 pub fn parent(&self, child: Entity) -> Option<Entity> {
62 self.child_parent_relations.get(&child).copied()
63 }
64
65 pub fn children(&self, parent: Entity) -> Option<impl Iterator<Item = Entity> + '_> {
66 self.parent_children_relations
67 .get(&parent)
68 .map(|list| list.iter().copied())
69 }
70
71 pub fn entity_by_name(&self, name: &str) -> Option<Entity> {
72 self.name_entities_map.get(name).copied()
73 }
74
75 pub fn name_by_entity(&self, entity: Entity) -> Option<&str> {
76 self.entity_names_map.get(&entity).map(|name| name.as_str())
77 }
78
79 pub fn find(&self, root: Option<Entity>, mut path: &str) -> Option<Entity> {
80 let mut root = match root {
81 Some(root) => root,
82 None => {
83 let part = match path.find('/') {
84 Some(found) => {
85 let part = &path[..found];
86 path = &path[(found + 1)..];
87 part
88 }
89 None => {
90 let part = path;
91 path = "";
92 part
93 }
94 };
95 match self.entity_by_name(part) {
96 Some(root) => root,
97 None => return None,
98 }
99 }
100 };
101 for part in path.split('/') {
102 match part {
103 "" | "." => {}
104 ".." => match self.parent(root) {
105 Some(parent) => root = parent,
106 None => return None,
107 },
108 part => match self.children(root) {
109 Some(mut iter) => match iter.find(|child| {
110 self.name_by_entity(*child)
111 .map(|name| name == part)
112 .unwrap_or_default()
113 }) {
114 Some(child) => root = child,
115 None => return None,
116 },
117 None => return None,
118 },
119 }
120 }
121 Some(root)
122 }
123
124 pub fn iter(&self) -> HierarchyIter {
125 HierarchyIter {
126 hierarchy: self,
127 iter_stack: vec![Box::new(self.roots.iter().copied())],
128 parent_stack: vec![],
129 }
130 }
131}
132
133pub struct HierarchyIter<'a> {
134 hierarchy: &'a Hierarchy,
135 iter_stack: Vec<Box<dyn Iterator<Item = Entity> + 'a>>,
136 parent_stack: Vec<Entity>,
137}
138
139impl<'a> Iterator for HierarchyIter<'a> {
140 type Item = (Entity, Option<Entity>);
141
142 fn next(&mut self) -> Option<Self::Item> {
143 loop {
144 if let Some(mut iter) = self.iter_stack.pop() {
145 if let Some(entity) = iter.next() {
146 self.iter_stack.push(iter);
147 if let Some(children) = self.hierarchy.children(entity) {
148 self.iter_stack.push(Box::new(children));
149 }
150 return Some((entity, self.parent_stack.last().copied()));
151 } else {
152 continue;
153 }
154 } else {
155 return None;
156 }
157 }
158 }
159}
160
161pub type HierarchySystemResources<'a> = (
162 WorldRef,
163 &'a mut UniverseCommands,
164 &'a EntityChanges,
165 &'a mut Hierarchy,
166 Comp<&'a Parent>,
167 Comp<&'a Name>,
168);
169
170pub fn hierarchy_system(universe: &mut Universe) {
171 let (world, mut commands, changes, mut hierarchy, ..) =
172 universe.query_resources::<HierarchySystemResources>();
173
174 if changes.has_changed() {
175 despawn(&mut commands, &changes, &hierarchy);
176
177 hierarchy.roots = Vec::with_capacity(world.len() as usize);
178 hierarchy.child_parent_relations = HashMap::with_capacity(world.len() as usize);
179 hierarchy.parent_children_relations = HashMap::with_capacity(world.len() as usize / 10);
180 hierarchy.entity_names_map = HashMap::with_capacity(world.len() as usize / 10);
181 hierarchy.name_entities_map = HashMap::with_capacity(world.len() as usize / 10);
182 } else {
183 hierarchy.roots.clear();
184 hierarchy.child_parent_relations.clear();
185 hierarchy.parent_children_relations.clear();
186 hierarchy.entity_names_map.clear();
187 hierarchy.name_entities_map.clear();
188 }
189
190 for (entity, _) in world.query::<()>().without::<&Parent>().iter() {
191 hierarchy.roots.push(entity);
192 }
193
194 for (child, (parent, name)) in world.query::<(Option<&Parent>, Option<&Name>)>().iter() {
195 if let Some(parent) = parent {
196 hierarchy.child_parent_relations.insert(child, parent.0);
197 let list: &mut Vec<Entity> = hierarchy
198 .parent_children_relations
199 .entry(parent.0)
200 .or_default();
201 list.push(child);
202 }
203 if let Some(name) = name {
204 let name: String = name.0.clone().into();
205 hierarchy.entity_names_map.insert(child, name.to_owned());
206 hierarchy.name_entities_map.insert(name, child);
207 }
208 }
209
210 for (child, name) in world.query::<&Name>().iter() {
211 hierarchy
212 .entity_names_map
213 .insert(child, name.0.as_ref().to_owned());
214 hierarchy
215 .name_entities_map
216 .insert(name.0.as_ref().to_owned(), child);
217 }
218}
219
220fn despawn(commands: &mut UniverseCommands, changes: &EntityChanges, hierarchy: &Hierarchy) {
221 for entity in changes.despawned() {
222 despawn_children(commands, entity, hierarchy);
223 }
224}
225
226fn despawn_children(commands: &mut UniverseCommands, parent: Entity, hierarchy: &Hierarchy) {
227 if let Some(iter) = hierarchy.children(parent) {
228 for entity in iter {
229 commands.schedule(DespawnEntity(entity));
230 despawn_children(commands, entity, hierarchy);
231 }
232 }
233}
234
235#[cfg(test)]
236mod tests {
237 use super::*;
238 use hecs::Component;
239
240 #[test]
241 fn test_send_sync() {
242 fn foo<T>()
243 where
244 T: Component + Send + Sync,
245 {
246 println!("{} is Component", std::any::type_name::<T>());
247 }
248
249 foo::<Parent>();
250 }
251}