use std::mem;
use hecs::{Component, DynamicBundle, Entity, QueryBorrow, Without, World};
use hecs_schedule::{error::Result, GenericWorld};
use crate::{
AncestorIter, BreadthFirstIterator, Child, ChildrenIter, DepthFirstIterator, DepthFirstVisitor,
Parent,
};
pub trait HierarchyMut {
fn attach<T: Component>(&mut self, child: Entity, parent: Entity) -> Result<Entity>;
fn attach_new<T: Component, C: DynamicBundle>(
&mut self,
parent: Entity,
components: C,
) -> Result<Entity>;
fn detach_all<T: Component>(&mut self, entity: Entity) -> Result<()>;
fn detach_children<T: Component>(&mut self, parent: Entity) -> Result<Vec<Entity>>;
fn despawn_children<T: Component>(&mut self, parent: Entity) -> Result<()>;
fn detach<T: Component>(&mut self, child: Entity) -> Result<()>;
fn despawn_all<T: Component>(&mut self, parent: Entity);
}
pub trait Hierarchy
where
Self: Sized,
{
fn parent<T: Component>(&self, child: Entity) -> Result<Entity>;
fn root<T: Component>(&self, child: Entity) -> Result<Entity>;
fn children<T: Component>(&self, parent: Entity) -> ChildrenIter<T>;
fn ancestors<T: Component>(&self, child: Entity) -> AncestorIter<T>;
fn descendants_depth_first<T: Component>(&self, root: Entity) -> DepthFirstIterator<T>;
fn visit<T: Component, F: Fn(&Self, Entity) -> bool + Component>(
&self,
root: Entity,
accept: F,
) -> DepthFirstVisitor<Self, T, F>;
fn descendants_breadth_first<T: Component>(
&self,
root: Entity,
) -> BreadthFirstIterator<Self, T>;
fn roots<T: Component>(&self) -> Result<QueryBorrow<Without<&Parent<T>, &Child<T>>>>;
}
impl HierarchyMut for World {
fn attach<T: Component>(&mut self, child: Entity, parent: Entity) -> Result<Entity> {
let mut maybe_p = self.try_get_mut::<Parent<T>>(parent);
if let Ok(ref mut p) = maybe_p {
p.num_children += 1;
let prev = p.last_child;
p.last_child = child;
let mut prev_data = self.try_get_mut::<Child<T>>(prev)?;
let next = prev_data.next;
prev_data.next = child;
mem::drop(prev_data);
mem::drop(maybe_p);
{
let mut next_data = self.try_get_mut::<Child<T>>(next)?;
next_data.prev = child;
}
self.try_insert(child, (Child::<T>::new(parent, next, prev),))?;
return Ok(child);
}
mem::drop(maybe_p);
self.try_insert(parent, (Parent::<T>::new(1, child),))?;
self.try_insert(child, (Child::<T>::new(parent, child, child),))?;
Ok(child)
}
fn attach_new<T: Component, C: DynamicBundle>(
&mut self,
parent: Entity,
components: C,
) -> Result<Entity> {
let child = self.spawn(components);
self.attach::<T>(child, parent)
}
fn detach_all<T: Component>(&mut self, entity: Entity) -> Result<()> {
self.detach_children::<T>(entity)?;
self.detach::<T>(entity)?;
Ok(())
}
fn detach_children<T: Component>(&mut self, parent: Entity) -> Result<Vec<Entity>> {
let children = self.children::<T>(parent).collect::<Vec<Entity>>();
children.iter().try_for_each(|child| -> Result<_> {
self.try_remove_one::<Child<T>>(*child)?;
Ok(())
})?;
self.remove_one::<Parent<T>>(parent).unwrap();
Ok(children)
}
fn despawn_children<T: Component>(&mut self, parent: Entity) -> Result<()> {
let children = self.children::<T>(parent).collect::<Vec<Entity>>();
children
.iter()
.for_each(|child| self.despawn_all::<Child<T>>(*child));
self.remove_one::<Parent<T>>(parent).unwrap();
Ok(())
}
fn detach<T: Component>(&mut self, child: Entity) -> Result<()> {
let data = self.try_get_mut::<Child<T>>(child)?;
let parent = data.parent;
let prev = data.prev;
let next = data.next;
mem::drop(data);
self.try_get_mut::<Child<T>>(prev)?.next = next;
self.try_get_mut::<Child<T>>(next)?.prev = prev;
let mut parent = self.try_get_mut::<Parent<T>>(parent)?;
parent.num_children -= 1;
if parent.last_child == child {
parent.last_child = prev;
}
Ok(())
}
fn despawn_all<T: Component>(&mut self, parent: Entity) {
let to_despawn = self
.descendants_depth_first::<T>(parent)
.collect::<Vec<_>>();
let _ = self.detach::<T>(parent);
to_despawn.iter().for_each(|entity| {
let _ = self.despawn(*entity);
});
let _ = self.despawn(parent);
}
}
impl<W: GenericWorld> Hierarchy for W {
fn parent<T: Component>(&self, child: Entity) -> Result<Entity> {
self.try_get::<Child<T>>(child).map(|child| child.parent)
}
fn root<T: Component>(&self, child: Entity) -> Result<Entity> {
let mut cur = child;
loop {
match self.parent::<T>(cur) {
Ok(val) => cur = val,
Err(hecs_schedule::Error::MissingComponent(_, _)) => break,
Err(val) => return Err(val),
}
}
Ok(cur)
}
fn children<T: Component>(&self, parent: Entity) -> ChildrenIter<T> {
self.try_get::<Parent<T>>(parent)
.and_then(|parent| {
let first_child = parent.first_child(self)?;
Ok(ChildrenIter::new(
self,
parent.num_children,
Some(first_child),
))
})
.unwrap_or_else(move |_| {
ChildrenIter::new(self, 0, None)
})
}
fn ancestors<T: Component>(&self, child: Entity) -> AncestorIter<T> {
AncestorIter::new(self, child)
}
fn descendants_depth_first<T: Component>(&self, root: Entity) -> DepthFirstIterator<T> {
DepthFirstIterator::new(self, root)
}
fn descendants_breadth_first<T: Component>(
&self,
root: Entity,
) -> BreadthFirstIterator<Self, T> {
BreadthFirstIterator::new(self, root)
}
fn visit<T: Component, F: Fn(&Self, Entity) -> bool + Component>(
&self,
root: Entity,
accept: F,
) -> DepthFirstVisitor<Self, T, F> {
DepthFirstVisitor::new(self, root, accept)
}
fn roots<T: Component>(&self) -> Result<QueryBorrow<Without<&Parent<T>, &Child<T>>>> {
Ok(self.try_query::<&Parent<T>>()?.without::<&Child<T>>())
}
}
trait WorldExt {
fn try_insert(&mut self, e: Entity, c: impl DynamicBundle) -> Result<()>;
fn try_remove_one<C: Component>(&mut self, e: Entity) -> Result<C>;
}
impl WorldExt for World {
fn try_insert(&mut self, e: Entity, c: impl DynamicBundle) -> Result<()> {
self.insert(e, c)
.map_err(|_| hecs_schedule::Error::NoSuchEntity(e))
}
fn try_remove_one<C: Component>(&mut self, e: Entity) -> Result<C> {
self.remove_one::<C>(e)
.map_err(|_| hecs_schedule::Error::NoSuchEntity(e))
}
}
pub type HierarchyQuery<'a, T> = (&'a Parent<T>, &'a Child<T>);