diff --git a/crates/bevy_utils/src/map.rs b/crates/bevy_utils/src/map.rs index ca74e34dbb..3b54a357aa 100644 --- a/crates/bevy_utils/src/map.rs +++ b/crates/bevy_utils/src/map.rs @@ -1,7 +1,7 @@ use core::{any::TypeId, hash::Hash}; use bevy_platform::{ - collections::HashMap, + collections::{hash_map::Entry, HashMap}, hash::{Hashed, NoOpHash, PassHash}, }; @@ -38,6 +38,78 @@ impl PreHashMapExt for PreHashMap = HashMap; +/// Extension trait to make use of [`TypeIdMap`] more ergonomic. +/// +/// Each function on this trait is a trivial wrapper for a function +/// on [`HashMap`], replacing a `TypeId` key with a +/// generic parameter `T`. +/// +/// # Examples +/// +/// ```rust +/// # use std::any::TypeId; +/// # use bevy_utils::TypeIdMap; +/// use bevy_utils::TypeIdMapExt; +/// +/// struct MyType; +/// +/// // Using the built-in `HashMap` functions requires manually looking up `TypeId`s. +/// let mut map = TypeIdMap::default(); +/// map.insert(TypeId::of::(), 7); +/// assert_eq!(map.get(&TypeId::of::()), Some(&7)); +/// +/// // Using `TypeIdMapExt` functions does the lookup for you. +/// map.insert_type::(7); +/// assert_eq!(map.get_type::(), Some(&7)); +/// ``` +pub trait TypeIdMapExt { + /// Inserts a value for the type `T`. + /// + /// If the map did not previously contain this key then [`None`] is returned, + /// otherwise the value for this key is updated and the old value returned. + fn insert_type(&mut self, v: V) -> Option; + + /// Returns a reference to the value for type `T`, if one exists. + fn get_type(&self) -> Option<&V>; + + /// Returns a mutable reference to the value for type `T`, if one exists. + fn get_type_mut(&mut self) -> Option<&mut V>; + + /// Removes type `T` from the map, returning the value for this + /// key if it was previously present. + fn remove_type(&mut self) -> Option; + + /// Gets the type `T`'s entry in the map for in-place manipulation. + fn entry_type(&mut self) -> Entry<'_, TypeId, V, NoOpHash>; +} + +impl TypeIdMapExt for TypeIdMap { + #[inline] + fn insert_type(&mut self, v: V) -> Option { + self.insert(TypeId::of::(), v) + } + + #[inline] + fn get_type(&self) -> Option<&V> { + self.get(&TypeId::of::()) + } + + #[inline] + fn get_type_mut(&mut self) -> Option<&mut V> { + self.get_mut(&TypeId::of::()) + } + + #[inline] + fn remove_type(&mut self) -> Option { + self.remove(&TypeId::of::()) + } + + #[inline] + fn entry_type(&mut self) -> Entry<'_, TypeId, V, NoOpHash> { + self.entry(TypeId::of::()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -67,7 +139,7 @@ mod tests { #[test] fn stable_hash_within_same_program_execution() { use alloc::vec::Vec; - + let mut map_1 = >::default(); let mut map_2 = >::default(); for i in 1..10 {