oxygengine_core/ecs/
hierarchy.rs

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}