Skip to main content

pebble/ecs/
query.rs

1use crate::ecs::{resources::Resources, system_param::SystemParam};
2
3pub struct Query<'a, Q: hecs::Query> {
4    world: &'a hecs::World,
5    borrow: hecs::QueryBorrow<'a, Q>,
6    scratch: Option<hecs::QueryOne<'a, Q>>,
7}
8
9impl<'a, Q: hecs::Query> IntoIterator for &'a mut Query<'a, Q> {
10    type Item = Q::Item<'a>;
11    type IntoIter = hecs::QueryIter<'a, Q>;
12
13    fn into_iter(self) -> Self::IntoIter {
14        (&mut self.borrow).into_iter()
15    }
16}
17
18impl<'a, Q: hecs::Query> Query<'a, Q> {
19    pub fn iter(&mut self) -> impl Iterator<Item = Q::Item<'_>> {
20        self.borrow.iter()
21    }
22
23    pub fn get(&mut self, entity: hecs::Entity) -> Option<Q::Item<'_>> {
24        self.scratch = Some(self.world.query_one::<Q>(entity));
25        self.scratch.as_mut().unwrap().get().ok()
26    }
27
28    pub fn with<R: hecs::Query>(self) -> Query<'a, hecs::With<Q, R>> {
29        Query {
30            world: self.world,
31            borrow: self.borrow.with::<R>(),
32            scratch: None,
33        }
34    }
35
36    pub fn without<R: hecs::Query>(self) -> Query<'a, hecs::Without<Q, R>> {
37        Query {
38            world: self.world,
39            borrow: self.borrow.without::<R>(),
40            scratch: None,
41        }
42    }
43
44    pub fn single(&mut self) -> Q::Item<'_> {
45        self.get_single()
46            .expect("Query::single: expected exactly one matching entity")
47    }
48
49    pub fn get_single(&mut self) -> Option<Q::Item<'_>> {
50        let mut iter = self.borrow.iter();
51        let first = iter.next()?;
52        if iter.next().is_some() {
53            return None;
54        }
55        Some(first)
56    }
57}
58
59impl<Q> SystemParam for Query<'static, Q>
60where
61    Q: hecs::Query + 'static,
62{
63    type Item<'a> = Query<'a, Q>;
64    type State = ();
65
66    fn fetch<'a>(
67        world: &'a hecs::World,
68        _resources: &'a Resources,
69        _state: &'a mut Self::State,
70    ) -> Self::Item<'a> {
71        Query {
72            world,
73            borrow: world.query::<Q>(),
74            scratch: None,
75        }
76    }
77}