Repository: Ralith/hecs
Pull request #58 (2020-03-14 02:16:19)

Support querying a single entity
Still considering whether this is worth having.

1tgr: Would be good to have this, just now I wrote a sequence of calls to `get(entity).unwrap()` for individual components.
extrawurst: hm thought there is something for that (see https://github.com/Ralith/hecs/issues/50#issuecomment-581072732):

> It's unlikely to matter for performance in practice, but [`World::entity`](https://docs.rs/hecs/0.2.7/hecs/struct.World.html#method.entity) provides that.
1tgr: Agree that `World::entity` allows me to access a single entity, but I still have to issue multiple calls to `get` or `get_mut` afterward. With this `World::query_one` it's a one-liner regardless of the number of components involved.

I'm not so concerned about performance in this case as convenience... the `query_` methods nicely handle the component checks for me, and they extend to `with`, `without` and `Option<T>` for more complex cases.
extrawurst: Agreed, I think it would be a nice addition to have `query_one`
Ralith: Thanks for the feedback! I'll move this forward. My only reservation was reluctance to increase the API surface area without anyone having actually asked for it.
Ralith: Revised the API a bit for more consistency with `query`; in particular `with`/`without` work, the error cases can be distinguished, and there's one less layer of indirection. Any final feedback?
1tgr: LGTM!

Patch:
diff --git a/src/lib.rs b/src/lib.rs
index 5130847..5e771e4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -67,6 +67,7 @@ mod bundle;
 mod entities;
 mod entity_builder;
 mod query;
+mod query_one;
 mod world;
 
 pub use archetype::Archetype;
@@ -75,6 +76,7 @@ pub use bundle::{Bundle, DynamicBundle, MissingComponent};
 pub use entities::{Entity, NoSuchEntity};
 pub use entity_builder::{BuiltEntity, EntityBuilder};
 pub use query::{Access, BatchedIter, Query, QueryBorrow, QueryIter, With, Without};
+pub use query_one::QueryOne;
 pub use world::{ArchetypesGeneration, Component, ComponentError, Iter, SpawnBatchIter, World};
 
 // Unstable implementation details needed by the macros
diff --git a/src/query_one.rs b/src/query_one.rs
new file mode 100644
index 0000000..2645d74
--- /dev/null
+++ b/src/query_one.rs
@@ -0,0 +1,81 @@
+use core::marker::PhantomData;
+
+use crate::query::{Fetch, With, Without};
+use crate::{Archetype, Component, Query};
+
+/// A borrow of a `World` sufficient to execute the query `Q` on a single entity
+pub struct QueryOne<'a, Q: Query> {
+    archetype: &'a Archetype,
+    index: u32,
+    borrowed: bool,
+    _marker: PhantomData<Q>,
+}
+
+impl<'a, Q: Query> QueryOne<'a, Q> {
+    /// Construct a query accessing the entity in `archetype` at `index`
+    ///
+    /// # Safety
+    ///
+    /// `index` must be in-bounds for `archetype`
+    pub(crate) unsafe fn new(archetype: &'a Archetype, index: u32) -> Self {
+        Self {
+            archetype,
+            index,
+            borrowed: false,
+            _marker: PhantomData,
+        }
+    }
+
+    /// Get the query result, or `None` if the entity does not satisfy the query
+    ///
+    /// Must be called at most once.
+    ///
+    /// Panics if called more than once or if it would construct a borrow that clashes with another
+    /// pre-existing borrow.
+    pub fn get(&mut self) -> Option<<Q::Fetch as Fetch<'_>>::Item> {
+        if self.borrowed {
+            panic!("called QueryOnce::get twice; construct a new query instead");
+        }
+        unsafe {
+            let mut fetch = Q::Fetch::get(self.archetype, self.index as usize)?;
+            self.borrowed = true;
+            Q::Fetch::borrow(self.archetype);
+            Some(fetch.next())
+        }
+    }
+
+    /// Transform the query into one that requires a certain component without borrowing it
+    ///
+    /// See `QueryBorrow::with` for details.
+    pub fn with<T: Component>(self) -> QueryOne<'a, With<T, Q>> {
+        self.transform()
+    }
+
+    /// Transform the query into one that skips entities having a certain component
+    ///
+    /// See `QueryBorrow::without` for details.
+    pub fn without<T: Component>(self) -> QueryOne<'a, Without<T, Q>> {
+        self.transform()
+    }
+
+    /// Helper to change the type of the query
+    fn transform<R: Query>(mut self) -> QueryOne<'a, R> {
+        let x = QueryOne {
+            archetype: self.archetype,
+            index: self.index,
+            borrowed: self.borrowed,
+            _marker: PhantomData,
+        };
+        // Ensure `Drop` won't fire redundantly
+        self.borrowed = false;
+        x
+    }
+}
+
+impl<Q: Query> Drop for QueryOne<'_, Q> {
+    fn drop(&mut self) {
+        if self.borrowed {
+            Q::Fetch::release(self.archetype);
+        }
+    }
+}
diff --git a/src/world.rs b/src/world.rs
index cd458c9..6882080 100644
--- a/src/world.rs
+++ b/src/world.rs
@@ -26,7 +26,7 @@ use crate::archetype::Archetype;
 use crate::entities::{Entities, Location};
 use crate::{
     Bundle, DynamicBundle, Entity, EntityRef, MissingComponent, NoSuchEntity, Query, QueryBorrow,
-    Ref, RefMut,
+    QueryOne, Ref, RefMut,
 };
 
 /// An unordered collection of entities, each having any number of distinctly typed components
@@ -248,6 +248,27 @@ impl World {
         QueryBorrow::new(&self.entities.meta, &self.archetypes)
     }
 
+    /// Prepare a query against a single entity
+    ///
+    /// Call `get` on the resulting `QueryOne` to actually execute the query.
+    ///
+    /// Handy for accessing multiple components simultaneously.
+    ///
+    /// # Example
+    /// ```
+    /// # use hecs::*;
+    /// let mut world = World::new();
+    /// let a = world.spawn((123, true, "abc"));
+    /// assert_eq!(
+    ///     world.query_one::<(&i32, &bool)>(a).unwrap().get().unwrap(),
+    ///     (&123, &true)
+    /// );
+    /// ```
+    pub fn query_one<Q: Query>(&self, entity: Entity) -> Result<QueryOne<'_, Q>, NoSuchEntity> {
+        let loc = self.entities.get(entity)?;
+        Ok(unsafe { QueryOne::new(&self.archetypes[loc.archetype as usize], loc.index) })
+    }
+
     /// Borrow the `T` component of `entity`
     ///
     /// Panics if the component is already uniquely borrowed from another entity with the same


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 5538241..ffb0750 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -341,3 +341,20 @@ fn spawn_batch() {
         .collect::<Vec<_>>();
     assert_eq!(entities.len(), 100);
 }
+
+#[test]
+fn query_one() {
+    let mut world = World::new();
+    let a = world.spawn(("abc", 123));
+    let b = world.spawn(("def", 456));
+    let c = world.spawn(("ghi", 789, true));
+    assert_eq!(world.query_one::<&i32>(a).unwrap().get(), Some(&123));
+    assert_eq!(world.query_one::<&i32>(b).unwrap().get(), Some(&456));
+    assert!(world.query_one::<(&i32, &bool)>(a).unwrap().get().is_none());
+    assert_eq!(
+        world.query_one::<(&i32, &bool)>(c).unwrap().get(),
+        Some((&789, &true))
+    );
+    world.despawn(a).unwrap();
+    assert!(world.query_one::<&i32>(a).is_err());
+}