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}