diff --git a/crates/bevy_reflect/src/map.rs b/crates/bevy_reflect/src/map.rs index ad1fcdd73e..ee6fab4ff0 100644 --- a/crates/bevy_reflect/src/map.rs +++ b/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`]. -/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding `#[reflect(Hash)]` -/// to the entire struct or enum. -/// This is true even for manual implementors who do not use the hashed value, +/// All keys are expected to return a valid hash value from [`PartialReflect::reflect_hash`] and be +/// comparable using [`PartialReflect::reflect_partial_eq`]. +/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding +/// `#[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::(), 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, Box)>, - indices: HashMap, + indices: HashTable, } impl DynamicMap { @@ -225,20 +229,38 @@ impl DynamicMap { pub fn insert(&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, Box)], + ) -> 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))) - .map(|index| &*self.values.get(*index).unwrap().1) + .find(hash, eq) + .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))) - .cloned() - .map(move |index| &mut *self.values.get_mut(index).unwrap().1) + .find(hash, eq) + .map(|&index| &mut *self.values[index].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, - mut value: Box, + value: Box, ) -> Option> { - match self - .indices - .entry(key.reflect_hash().expect(hash_error!(key))) - { - Entry::Occupied(entry) => { - let (_old_key, old_value) = self.values.get_mut(*entry.get()).unwrap(); - core::mem::swap(old_value, &mut value); - Some(value) + assert_eq!( + key.reflect_partial_eq(&*key), + Some(true), + "keys inserted in `Map`-like types are expected to reflect `PartialEq`" + ); + + let hash = Self::internal_hash(&*key); + 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) => { - entry.insert(self.values.len()); + None => { + 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> { - let index = self - .indices - .remove(&key.reflect_hash().expect(hash_error!(key)))?; - let (_key, value) = self.values.remove(index); - Some(value) + let hash = Self::internal_hash(key); + let eq = Self::internal_eq(key, &self.values); + match self.indices.find_entry(hash, eq) { + Ok(entry) => { + 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()); + } } diff --git a/crates/bevy_reflect/src/set.rs b/crates/bevy_reflect/src/set.rs index 85bc57ca84..f15730071a 100644 --- a/crates/bevy_reflect/src/set.rs +++ b/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 -/// order of these entries is not guaranteed by this trait. +/// Sets contain zero or more entries of a fixed type, and correspond to types +/// 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`]. -/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding `#[reflect(Hash)]` -/// to the entire struct or enum. -/// This is true even for manual implementors who do not use the hashed value, +/// All values are expected to return a valid hash value from [`PartialReflect::reflect_hash`] and be +/// comparable using [`PartialReflect::reflect_partial_eq`]. +/// If using the [`#[derive(Reflect)]`](derive@crate::Reflect) macro, this can be done by adding +/// `#[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::(), 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 {