feat: single & first entity queries (take 2) (#454)

Closes #289

Hello 👋 it's been a long while. This is take 2 of #386 which was
reverted in cfce040 due to CI failures.

## Changes

- Add `Entities::(get_)single_with` methods to get a single entity by a
query
- Add `Entities::(get_)first_with_bitset` methods to get the first
entity in the given bitset
- Add `Entities::(get_)first_with` methods to get the first entity in
the given query

## Summary

Add convenience methods to get a single entity (and one or more of its
components) when there should only be one of the thing you're looking
for. This could be useful, for example, to check if there is only one
player alive in the game. You could use `Entities::get_single_with`
which would allow you to detect when there is only one player (`Ok`) or
when there are none/multiple (`Err`).

All methods have a panicking and non-panicking variant -- e.g.
`get_single_with` returns a `Result<_, QuerySingleError>`, while
`single_with` panics if the return value is an `Err`. The `*first_with*`
methods either return an `Option` or panic.

framework_crates/bones_ecs/src/components/iterator.rs

@@ -38,13 +38,13 @@ impl<'a> Iterator for UntypedComponentBitsetIterator<'a> {
     type Item = SchemaRef<'a>;
     fn next(&mut self) -> Option<Self::Item> {
         let max_id = self.components.max_id;
-        while !(self.bitset.bit_test(self.current_id)
-            && self.components.bitset.bit_test(self.current_id))
-            && self.current_id <= max_id
+        while self.current_id < max_id
+            && !(self.bitset.bit_test(self.current_id)
+                && self.components.bitset.bit_test(self.current_id))
         {
             self.current_id += 1;
         }
-        let ret = if self.current_id <= max_id {
+        let ret = if self.current_id < max_id {
             // SAFE: Here we are just getting a pointer, not doing anything unsafe with it.
             Some(unsafe {
                 SchemaRef::from_ptr_schema(
@@ -71,11 +71,11 @@ impl<'a> Iterator for UntypedComponentOptionalBitsetIterator<'a> {
     fn next(&mut self) -> Option<Self::Item> {
         // We stop iterating at bitset length, not component store length, as we want to iterate over
         // whole bitset and return None for entities that don't have this optional component.
-        let max_id = self.bitset.bit_len() - 1;
-        while !self.bitset.bit_test(self.current_id) && self.current_id <= max_id {
+        let max_id = self.bitset.bit_len();
+        while self.current_id < max_id && !self.bitset.bit_test(self.current_id) {
             self.current_id += 1;
         }
-        let ret = if self.current_id <= max_id {
+        let ret = if self.current_id < max_id {
             // SAFE: Here we are just getting a pointer, not doing anything unsafe with it.
             if self.components.bitset.bit_test(self.current_id) {
                 Some(Some(unsafe {
@@ -110,13 +110,13 @@ impl<'a> Iterator for UntypedComponentBitsetIteratorMut<'a> {
     type Item = SchemaRefMut<'a>;
     fn next(&mut self) -> Option<Self::Item> {
         let max_id = self.components.max_id;
-        while !(self.bitset.bit_test(self.current_id)
-            && self.components.bitset.bit_test(self.current_id))
-            && self.current_id <= max_id
+        while self.current_id < max_id
+            && !(self.bitset.bit_test(self.current_id)
+                && self.components.bitset.bit_test(self.current_id))
         {
             self.current_id += 1;
         }
-        let ret = if self.current_id <= max_id {
+        let ret = if self.current_id < max_id {
             // SAFE: We know that the index is within bounds, and we know that the pointer will be
             // valid for the new lifetime.
             Some(unsafe {
@@ -144,11 +144,11 @@ impl<'a> Iterator for UntypedComponentOptionalBitsetIteratorMut<'a> {
     fn next(&mut self) -> Option<Self::Item> {
         // We do not stop iterating at component store length, as we want to iterate over
         // whole bitset and return None for entities that don't have this optional component.
-        let max_id = self.bitset.bit_len() - 1;
-        while !self.bitset.bit_test(self.current_id) && self.current_id <= max_id {
+        let max_id = self.bitset.bit_len();
+        while self.current_id < max_id && !self.bitset.bit_test(self.current_id) {
             self.current_id += 1;
         }
-        let ret = if self.current_id <= max_id {
+        let ret = if self.current_id < max_id {
             // SAFE: Here we are just getting a pointer, not doing anything unsafe with it.
             if self.components.bitset.bit_test(self.current_id) {
                 Some(Some(unsafe {

framework_crates/bones_ecs/src/components/typed.rs

@@ -83,7 +83,7 @@ impl<T: HasSchema> ComponentStore<T> {
         self.untyped.get_mut_or_insert(entity, f)
     }
 
-    /// Get mutable references s to the component data for multiple entities at the same time.
+    /// Get mutable references to the component data for multiple entities at the same time.
     ///
     /// # Panics
     ///
@@ -110,6 +110,27 @@ impl<T: HasSchema> ComponentStore<T> {
         self.untyped.remove(entity)
     }
 
+    /// Gets an immutable reference to the component if there is exactly one instance of it.
+    #[inline]
+    pub fn get_single_with_bitset(&self, bitset: Rc<BitSetVec>) -> Result<&T, QuerySingleError> {
+        // SOUND: we know the schema matches.
+        self.untyped
+            .get_single_with_bitset(bitset)
+            .map(|x| unsafe { x.cast_into_unchecked() })
+    }
+
+    /// Gets a mutable reference to the component if there is exactly one instance of it.
+    #[inline]
+    pub fn get_single_with_bitset_mut(
+        &mut self,
+        bitset: Rc<BitSetVec>,
+    ) -> Result<&mut T, QuerySingleError> {
+        // SOUND: we know the schema matches.
+        self.untyped
+            .get_single_with_bitset_mut(bitset)
+            .map(|x| unsafe { x.cast_into_mut_unchecked() })
+    }
+
     /// Iterates immutably over all components of this type.
     /// Very fast but doesn't allow joining with other component types.
     #[inline]
@@ -137,6 +158,15 @@ impl<T: HasSchema> ComponentStore<T> {
 ///
 /// Automatically implemented for [`ComponentStore`].
 pub trait ComponentIterBitset<'a, T: HasSchema> {
+    /// Gets an immutable reference to the component if there is exactly one instance of it.
+    fn get_single_with_bitset(&self, bitset: Rc<BitSetVec>) -> Result<&T, QuerySingleError>;
+
+    /// Gets a mutable reference to the component if there is exactly one instance of it.
+    fn get_single_mut_with_bitset(
+        &mut self,
+        bitset: Rc<BitSetVec>,
+    ) -> Result<&mut T, QuerySingleError>;
+
     /// Iterates immutably over the components of this type where `bitset`
     /// indicates the indices of entities.
     /// Slower than `iter()` but allows joining between multiple component types.
@@ -174,6 +204,27 @@ pub trait ComponentIterBitset<'a, T: HasSchema> {
 }
 
 impl<'a, T: HasSchema> ComponentIterBitset<'a, T> for ComponentStore<T> {
+    /// Gets an immutable reference to the component if there is exactly one instance of it.
+    fn get_single_with_bitset(&self, bitset: Rc<BitSetVec>) -> Result<&T, QuerySingleError> {
+        // SOUND: we know the schema matches.
+        fn map<T>(r: SchemaRef) -> &T {
+            unsafe { r.cast_into_unchecked() }
+        }
+        self.untyped.get_single_with_bitset(bitset).map(map)
+    }
+
+    /// Gets a mutable reference to the component if there is exactly one instance of it.
+    fn get_single_mut_with_bitset(
+        &mut self,
+        bitset: Rc<BitSetVec>,
+    ) -> Result<&mut T, QuerySingleError> {
+        // SOUND: we know the schema matches.
+        fn map<T>(r: SchemaRefMut) -> &mut T {
+            unsafe { r.cast_into_mut_unchecked() }
+        }
+        self.untyped.get_single_with_bitset_mut(bitset).map(map)
+    }
+
     /// Iterates immutably over the components of this type where `bitset`
     /// indicates the indices of entities.
     /// Slower than `iter()` but allows joining between multiple component types.
@@ -207,7 +258,7 @@ impl<'a, T: HasSchema> ComponentIterBitset<'a, T> for ComponentStore<T> {
     #[inline]
     fn iter_mut_with_bitset(&mut self, bitset: Rc<BitSetVec>) -> ComponentBitsetIteratorMut<T> {
         // SOUND: we know the schema matches.
-        fn map<T>(r: SchemaRefMut<'_>) -> &mut T {
+        fn map<T>(r: SchemaRefMut) -> &mut T {
             unsafe { r.cast_into_mut_unchecked() }
         }
 
@@ -250,14 +301,16 @@ impl<'a, T: HasSchema> ComponentIterBitset<'a, T> for ComponentStore<T> {
 
 #[cfg(test)]
 mod tests {
+    use std::rc::Rc;
+
     use crate::prelude::*;
 
+    #[derive(Debug, Clone, PartialEq, Eq, HasSchema, Default)]
+    #[repr(C)]
+    struct A(String);
+
     #[test]
     fn create_remove_components() {
-        #[derive(Debug, Clone, PartialEq, Eq, HasSchema, Default)]
-        #[repr(C)]
-        struct A(String);
-
         let mut entities = Entities::default();
         let e1 = entities.create();
         let e2 = entities.create();
@@ -276,10 +329,6 @@ mod tests {
 
     #[test]
     fn get_mut_or_insert() {
-        #[derive(Debug, Clone, PartialEq, Eq, HasSchema, Default)]
-        #[repr(C)]
-        struct A(String);
-
         let mut entities = Entities::default();
         let e1 = entities.create();
 
@@ -299,4 +348,129 @@ mod tests {
         // Test that existing component is retrieved
         assert_eq!(comp.0, "Test2");
     }
+
+    #[test]
+    fn single_returns_none_when_empty() {
+        let storage = ComponentStore::<A>::default();
+        let bitset = Rc::new({
+            let mut entities = Entities::default();
+            entities.create();
+            entities.bitset().clone()
+        });
+
+        let maybe_comp = storage.get_single_with_bitset(bitset);
+
+        assert_eq!(maybe_comp, Err(QuerySingleError::NoEntities));
+    }
+
+    #[test]
+    fn single_returns_some_single() {
+        let mut storage = ComponentStore::<A>::default();
+        let mut entities = Entities::default();
+
+        // Create some dummies so that the target entity isn't 0
+        (0..3).map(|_| entities.create()).count();
+
+        let e = entities.create();
+        let a = A("a".to_string());
+        storage.insert(e, a.clone());
+
+        let bitset = Rc::new(entities.bitset().clone());
+
+        let maybe_comp = storage.get_single_with_bitset(bitset);
+
+        assert_eq!(maybe_comp, Ok(&a));
+    }
+
+    #[test]
+    fn single_returns_none_when_more_than_1() {
+        let mut entities = Entities::default();
+        let mut storage = ComponentStore::<A>::default();
+
+        (0..3)
+            .map(|i| storage.insert(entities.create(), A(i.to_string())))
+            .count();
+
+        let bitset = Rc::new(entities.bitset().clone());
+
+        let maybe_comp = storage.get_single_with_bitset(bitset);
+
+        assert_eq!(maybe_comp, Err(QuerySingleError::MultipleEntities));
+    }
+
+    #[test]
+    fn iter_with_bitset() {
+        let mut entities = Entities::default();
+        let mut storage = ComponentStore::<A>::default();
+
+        {
+            let bitset = Rc::new(entities.bitset().clone());
+
+            let mut comp_iter = storage.iter_with_bitset(bitset.clone());
+            assert_eq!(comp_iter.next(), None);
+
+            let mut comp_mut_iter = storage.iter_mut_with_bitset(bitset);
+            assert_eq!(comp_mut_iter.next(), None);
+        }
+
+        {
+            let e = entities.create();
+            let mut a = A("e".to_string());
+            storage.insert(e, a.clone());
+
+            let bitset = Rc::new(entities.bitset().clone());
+
+            let mut comp_iter = storage.iter_with_bitset(bitset.clone());
+            assert_eq!(comp_iter.next(), Some(&a));
+
+            let mut comp_mut_iter = storage.iter_mut_with_bitset(bitset);
+            assert_eq!(comp_mut_iter.next(), Some(&mut a));
+
+            entities.kill(e);
+        }
+    }
+
+    #[test]
+    fn iter_with_bitset_optional() {
+        let mut entities = Entities::default();
+        let mut storage = ComponentStore::<A>::default();
+
+        {
+            let bitset = Rc::new(entities.bitset().clone());
+
+            let mut comp_iter = storage.iter_with_bitset_optional(bitset.clone());
+            assert_eq!(comp_iter.next(), None);
+
+            let mut comp_mut_iter = storage.iter_mut_with_bitset_optional(bitset);
+            assert_eq!(comp_mut_iter.next(), None);
+        }
+
+        {
+            let e = entities.create();
+            let bitset = Rc::new(entities.bitset().clone());
+
+            let mut comp_iter = storage.iter_with_bitset_optional(bitset.clone());
+            assert_eq!(comp_iter.next(), Some(None));
+
+            let mut comp_mut_iter = storage.iter_mut_with_bitset_optional(bitset);
+            assert_eq!(comp_mut_iter.next(), Some(None));
+
+            entities.kill(e);
+        }
+
+        {
+            let e = entities.create();
+            let mut a = A("e".to_string());
+            storage.insert(e, a.clone());
+            let bitset = Rc::new(entities.bitset().clone());
+
+            let mut comp_iter = storage.iter_with_bitset_optional(bitset.clone());
+            assert_eq!(comp_iter.next(), Some(Some(&a)));
+
+            let mut comp_mut_iter = storage.iter_mut_with_bitset_optional(bitset);
+            assert_eq!(comp_mut_iter.next(), Some(Some(&mut a)));
+
+            entities.kill(e);
+        }
+    }
 }

framework_crates/bones_ecs/src/components/untyped.rs

@@ -295,12 +295,10 @@ impl UntypedComponentStore {
         entity: Entity,
         f: impl FnOnce() -> T,
     ) -> &mut T {
-        if self.bitset.bit_test(entity.index() as usize) {
-            return self.get_mut(entity).unwrap();
-        } else {
+        if !self.bitset.bit_test(entity.index() as usize) {
             self.insert(entity, f());
-            self.get_mut(entity).unwrap()
         }
+        self.get_mut(entity).unwrap()
     }
 
     /// Get a [`SchemaRefMut`] to the component for the given [`Entity`]
@@ -492,6 +490,37 @@ impl UntypedComponentStore {
         }
     }
 
+    /// Get a reference to the component store if there is exactly one instance of the component.
+    pub fn get_single_with_bitset(
+        &self,
+        bitset: Rc<BitSetVec>,
+    ) -> Result<SchemaRef, QuerySingleError> {
+        let len = self.bitset().bit_len();
+        let mut iter = (0..len).filter(|&i| bitset.bit_test(i) && self.bitset().bit_test(i));
+        let i = iter.next().ok_or(QuerySingleError::NoEntities)?;
+        if iter.next().is_some() {
+            return Err(QuerySingleError::MultipleEntities);
+        }
+        // TODO: add unchecked variant to avoid redundant validation
+        self.get_idx(i).ok_or(QuerySingleError::NoEntities)
+    }
+
+    /// Get a mutable reference to the component store if there is exactly one instance of the
+    /// component.
+    pub fn get_single_with_bitset_mut(
+        &mut self,
+        bitset: Rc<BitSetVec>,
+    ) -> Result<SchemaRefMut, QuerySingleError> {
+        let len = self.bitset().bit_len();
+        let mut iter = (0..len).filter(|&i| bitset.bit_test(i) && self.bitset().bit_test(i));
+        let i = iter.next().ok_or(QuerySingleError::NoEntities)?;
+        if iter.next().is_some() {
+            return Err(QuerySingleError::MultipleEntities);
+        }
+        // TODO: add unchecked variant to avoid redundant validation
+        self.get_idx_mut(i).ok_or(QuerySingleError::NoEntities)
+    }
+
     /// Iterates immutably over all components of this type.
     ///
     /// Very fast but doesn't allow joining with other component types.
@@ -600,9 +629,8 @@ impl<'a> Iterator for UntypedComponentStoreIter<'a> {
                 if let Some(ptr) = self.store.get_idx(self.idx) {
                     self.idx += 1;
                     break Some(ptr);
-                } else {
-                    self.idx += 1;
                 }
+                self.idx += 1;
             } else {
                 break None;
             }
@@ -627,9 +655,8 @@ impl<'a> Iterator for UntypedComponentStoreIterMut<'a> {
                     break Some(unsafe {
                         SchemaRefMut::from_ptr_schema(ptr.as_ptr(), ptr.schema())
                     });
-                } else {
-                    self.idx += 1;
                 }
+                self.idx += 1;
             } else {
                 break None;
             }