Use HashTable in DynamicMap and fix bug in remove (#15158)

# Objective

- `DynamicMap` currently uses an `HashMap` from a `u64` hash to the
entry index in a `Vec`. This is incorrect in the presence of hash
collisions, so let's fix it;
- `DynamicMap::remove` was also buggy, as it didn't fix up the indexes
of the other elements after removal. Fix that up as well and add a
regression test.

## Solution

- Use `HashTable` in `DynamicMap` to distinguish entries that have the
same hash by using `reflect_partial_eq`, bringing it more in line with
what `DynamicSet` does;
- Reimplement `DynamicMap::remove` to properly fix up the index of moved
elements after the removal.

## Testing

- A regression test was added for the `DynamicMap::remove` issue.

---

Some kinda related considerations: the use of a separate `Vec` for
storing the entries adds some complications that I'm not sure are worth.
This is mainly used to implement an efficient `get_at`, which is relied
upon by `MapIter`. However both `HashMap` and `BTreeMap` implement
`get_at` inefficiently (and cannot do so efficiently), leading to a
`O(N^2)` complexity for iterating them. This could be removed in favor
of a `Box<dyn Iterator>` like it's done in `DynamicSet`.

crates/bevy_reflect/src/map.rs

@@ -1,7 +1,7 @@
 use core::fmt::{Debug, Formatter};
 
 use bevy_reflect_derive::impl_type_path;
-use bevy_utils::{Entry, HashMap};
+use bevy_utils::hashbrown::HashTable;
 
 use crate::{
     self as bevy_reflect, type_info::impl_type_methods, ApplyError, MaybeTyped, PartialReflect,
@@ -11,15 +11,17 @@ use crate::{
 /// A trait used to power [map-like] operations via [reflection].
 ///
 /// Maps contain zero or more entries of a key and its associated value,
-/// and correspond to types like [`HashMap`].
+/// and correspond to types like [`HashMap`] and [`BTreeMap`].
 /// The order of these entries is not guaranteed by this trait.
 ///
-/// # Hashing
+/// # Hashing and equality
 ///
-/// All keys are expected to return a valid hash value from [`PartialReflect::reflect_hash`].
+/// All keys are expected to return a valid hash value from [`PartialReflect::reflect_hash`] and be
-/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding `#[reflect(Hash)]`
+/// comparable using [`PartialReflect::reflect_partial_eq`].
-/// to the entire struct or enum.
+/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding
-/// This is true even for manual implementors who do not use the hashed value,
+/// `#[reflect(Hash, PartialEq)]` to the entire struct or enum.
+/// The ordering is expected to be total, that is as if the reflected type implements the [`Eq`] trait.
+/// This is true even for manual implementors who do not hash or compare values,
 /// as it is still relied on by [`DynamicMap`].
 ///
 /// # Example
@@ -37,6 +39,8 @@ use crate::{
 /// assert_eq!(field.try_downcast_ref::<bool>(), Some(&true));
 /// ```
 ///
+/// [`HashMap`]: std::collections::HashMap
+/// [`BTreeMap`]: std::collections::BTreeMap
 /// [map-like]: https://doc.rust-lang.org/book/ch08-03-hash-maps.html
 /// [reflection]: crate
 pub trait Map: PartialReflect {
@@ -198,7 +202,7 @@ macro_rules! hash_error {
 pub struct DynamicMap {
     represented_type: Option<&'static TypeInfo>,
     values: Vec<(Box<dyn PartialReflect>, Box<dyn PartialReflect>)>,
-    indices: HashMap<u64, usize>,
+    indices: HashTable<usize>,
 }
 
 impl DynamicMap {
@@ -225,20 +229,38 @@ impl DynamicMap {
     pub fn insert<K: PartialReflect, V: PartialReflect>(&mut self, key: K, value: V) {
         self.insert_boxed(Box::new(key), Box::new(value));
     }
+
+    fn internal_hash(value: &dyn PartialReflect) -> u64 {
+        value.reflect_hash().expect(hash_error!(value))
+    }
+
+    fn internal_eq<'a>(
+        value: &'a dyn PartialReflect,
+        values: &'a [(Box<dyn PartialReflect>, Box<dyn PartialReflect>)],
+    ) -> impl FnMut(&usize) -> bool + 'a {
+        |&index| {
+            value
+            .reflect_partial_eq(&*values[index].0)
+            .expect("underlying type does not reflect `PartialEq` and hence doesn't support equality checks")
+        }
+    }
 }
 
 impl Map for DynamicMap {
     fn get(&self, key: &dyn PartialReflect) -> Option<&dyn PartialReflect> {
+        let hash = Self::internal_hash(key);
+        let eq = Self::internal_eq(key, &self.values);
         self.indices
-            .get(&key.reflect_hash().expect(hash_error!(key)))
+            .find(hash, eq)
-            .map(|index| &*self.values.get(*index).unwrap().1)
+            .map(|&index| &*self.values[index].1)
     }
 
     fn get_mut(&mut self, key: &dyn PartialReflect) -> Option<&mut dyn PartialReflect> {
+        let hash = Self::internal_hash(key);
+        let eq = Self::internal_eq(key, &self.values);
         self.indices
-            .get(&key.reflect_hash().expect(hash_error!(key)))
+            .find(hash, eq)
-            .cloned()
+            .map(|&index| &mut *self.values[index].1)
-            .map(move |index| &mut *self.values.get_mut(index).unwrap().1)
     }
 
     fn get_at(&self, index: usize) -> Option<(&dyn PartialReflect, &dyn PartialReflect)> {
@@ -283,31 +305,60 @@ impl Map for DynamicMap {
     fn insert_boxed(
         &mut self,
         key: Box<dyn PartialReflect>,
-        mut value: Box<dyn PartialReflect>,
+        value: Box<dyn PartialReflect>,
     ) -> Option<Box<dyn PartialReflect>> {
-        match self
+        assert_eq!(
-            .indices
+            key.reflect_partial_eq(&*key),
-            .entry(key.reflect_hash().expect(hash_error!(key)))
+            Some(true),
-        {
+            "keys inserted in `Map`-like types are expected to reflect `PartialEq`"
-            Entry::Occupied(entry) => {
+        );
-                let (_old_key, old_value) = self.values.get_mut(*entry.get()).unwrap();
+
-                core::mem::swap(old_value, &mut value);
+        let hash = Self::internal_hash(&*key);
-                Some(value)
+        let eq = Self::internal_eq(&*key, &self.values);
+        match self.indices.find(hash, eq) {
+            Some(&index) => {
+                let (key_ref, value_ref) = &mut self.values[index];
+                *key_ref = key;
+                let old_value = core::mem::replace(value_ref, value);
+                Some(old_value)
             }
-            Entry::Vacant(entry) => {
+            None => {
-                entry.insert(self.values.len());
+                let index = self.values.len();
                 self.values.push((key, value));
+                self.indices.insert_unique(hash, index, |&index| {
+                    Self::internal_hash(&*self.values[index].0)
+                });
                 None
             }
         }
     }
 
     fn remove(&mut self, key: &dyn PartialReflect) -> Option<Box<dyn PartialReflect>> {
-        let index = self
+        let hash = Self::internal_hash(key);
-            .indices
+        let eq = Self::internal_eq(key, &self.values);
-            .remove(&key.reflect_hash().expect(hash_error!(key)))?;
+        match self.indices.find_entry(hash, eq) {
-        let (_key, value) = self.values.remove(index);
+            Ok(entry) => {
-        Some(value)
+                let (index, _) = entry.remove();
+                let (_, old_value) = self.values.swap_remove(index);
+
+                // The `swap_remove` might have moved the last element of `values`
+                // to `index`, so we might need to fix up its index in `indices`.
+                // If the removed element was also the last element there's nothing to
+                // fixup and this will return `None`, otherwise it returns the key
+                // whose index needs to be fixed up.
+                if let Some((moved_key, _)) = self.values.get(index) {
+                    let hash = Self::internal_hash(&**moved_key);
+                    let moved_index = self
+                        .indices
+                        .find_mut(hash, |&moved_index| moved_index == self.values.len())
+                        .expect("key inserted in a `DynamicMap` is no longer present, this means its reflected `Hash` might be incorrect");
+                    *moved_index = index;
+                }
+
+                Some(old_value)
+            }
+            Err(_) => None,
+        }
     }
 }
 
@@ -667,4 +718,21 @@ mod tests {
             assert_eq!(size, iter.index);
         }
     }
+
+    #[test]
+    fn remove() {
+        let mut map = DynamicMap::default();
+        map.insert(0, 0);
+        map.insert(1, 1);
+
+        assert_eq!(map.remove(&0).unwrap().try_downcast_ref(), Some(&0));
+        assert!(map.get(&0).is_none());
+        assert_eq!(map.get(&1).unwrap().try_downcast_ref(), Some(&1));
+
+        assert_eq!(map.remove(&1).unwrap().try_downcast_ref(), Some(&1));
+        assert!(map.get(&1).is_none());
+
+        assert!(map.remove(&1).is_none());
+        assert!(map.get(&1).is_none());
+    }
 }

crates/bevy_reflect/src/set.rs

@@ -10,15 +10,18 @@ use crate::{
 
 /// A trait used to power [set-like] operations via [reflection].
 ///
-/// Sets contain zero or more entries of a fixed type, and correspond to types like [`HashSet`](std::collections::HashSet). The
+/// Sets contain zero or more entries of a fixed type, and correspond to types
-/// order of these entries is not guaranteed by this trait.
+/// like [`HashSet`] and [`BTreeSet`].
+/// The order of these entries is not guaranteed by this trait.
 ///
-/// # Hashing
+/// # Hashing and equality
 ///
-/// All values are expected to return a valid hash value from [`PartialReflect::reflect_hash`].
+/// All values are expected to return a valid hash value from [`PartialReflect::reflect_hash`] and be
-/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding `#[reflect(Hash)]`
+/// comparable using [`PartialReflect::reflect_partial_eq`].
-/// to the entire struct or enum.
+/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding
-/// This is true even for manual implementors who do not use the hashed value,
+/// `#[reflect(Hash, PartialEq)]` to the entire struct or enum.
+/// The ordering is expected to be total, that is as if the reflected type implements the [`Eq`] trait.
+/// This is true even for manual implementors who do not hash or compare values,
 /// as it is still relied on by [`DynamicSet`].
 ///
 /// # Example
@@ -36,6 +39,8 @@ use crate::{
 /// assert_eq!(field.try_downcast_ref::<u32>(), Some(&123_u32));
 /// ```
 ///
+/// [`HashSet`]: std::collections::HashSet
+/// [`BTreeSet`]: std::collections::BTreeSet
 /// [set-like]: https://doc.rust-lang.org/stable/std/collections/struct.HashSet.html
 /// [reflection]: crate
 pub trait Set: PartialReflect {