bevy_utils/
map.rs

1use core::{any::TypeId, hash::Hash};
2
3use bevy_platform::{
4    collections::HashMap,
5    hash::{Hashed, NoOpHash, PassHash},
6};
7use indexmap::map::IndexMap;
8
9/// The [`Entry`][indexmap::map::Entry] type for [`TypeIdMap`].
10pub use indexmap::map::Entry as TypeIdMapEntry;
11
12/// A [`HashMap`] pre-configured to use [`Hashed`] keys and [`PassHash`] passthrough hashing.
13/// Iteration order only depends on the order of insertions and deletions.
14pub type PreHashMap<K, V> = HashMap<Hashed<K>, V, PassHash>;
15
16/// Extension methods intended to add functionality to [`PreHashMap`].
17pub trait PreHashMapExt<K, V> {
18    /// Tries to get or insert the value for the given `key` using the pre-computed hash first.
19    /// If the [`PreHashMap`] does not already contain the `key`, it will clone it and insert
20    /// the value returned by `func`.
21    fn get_or_insert_with<F: FnOnce() -> V>(&mut self, key: &Hashed<K>, func: F) -> &mut V;
22}
23
24impl<K: Hash + Eq + PartialEq + Clone, V> PreHashMapExt<K, V> for PreHashMap<K, V> {
25    #[inline]
26    fn get_or_insert_with<F: FnOnce() -> V>(&mut self, key: &Hashed<K>, func: F) -> &mut V {
27        use bevy_platform::collections::hash_map::RawEntryMut;
28        let entry = self
29            .raw_entry_mut()
30            .from_key_hashed_nocheck(key.hash(), key);
31        match entry {
32            RawEntryMut::Occupied(entry) => entry.into_mut(),
33            RawEntryMut::Vacant(entry) => {
34                let (_, value) = entry.insert_hashed_nocheck(key.hash(), key.clone(), func());
35                value
36            }
37        }
38    }
39}
40
41/// A specialized map type with Key of [`TypeId`]
42/// Iteration order only depends on the order of insertions and deletions.
43pub type TypeIdMap<V> = IndexMap<TypeId, V, NoOpHash>;
44
45/// Extension trait to make use of [`TypeIdMap`] more ergonomic.
46///
47/// Each function on this trait is a trivial wrapper for a function
48/// on [`IndexMap`], replacing a `TypeId` key with a
49/// generic parameter `T`.
50///
51/// # Examples
52///
53/// ```rust
54/// # use std::any::TypeId;
55/// # use bevy_utils::TypeIdMap;
56/// use bevy_utils::TypeIdMapExt;
57///
58/// struct MyType;
59///
60/// // Using the built-in `HashMap` functions requires manually looking up `TypeId`s.
61/// let mut map = TypeIdMap::default();
62/// map.insert(TypeId::of::<MyType>(), 7);
63/// assert_eq!(map.get(&TypeId::of::<MyType>()), Some(&7));
64///
65/// // Using `TypeIdMapExt` functions does the lookup for you.
66/// map.insert_type::<MyType>(7);
67/// assert_eq!(map.get_type::<MyType>(), Some(&7));
68/// ```
69pub trait TypeIdMapExt<V> {
70    /// Inserts a value for the type `T`.
71    ///
72    /// If the map did not previously contain this key then [`None`] is returned,
73    /// otherwise the value for this key is updated and the old value returned.
74    fn insert_type<T: ?Sized + 'static>(&mut self, v: V) -> Option<V>;
75
76    /// Returns a reference to the value for type `T`, if one exists.
77    fn get_type<T: ?Sized + 'static>(&self) -> Option<&V>;
78
79    /// Returns a mutable reference to the value for type `T`, if one exists.
80    fn get_type_mut<T: ?Sized + 'static>(&mut self) -> Option<&mut V>;
81
82    /// Removes type `T` from the map, returning the value for this
83    /// key if it was previously present.
84    fn remove_type<T: ?Sized + 'static>(&mut self) -> Option<V>;
85
86    /// Gets the type `T`'s entry in the map for in-place manipulation.
87    fn entry_type<T: ?Sized + 'static>(&mut self) -> TypeIdMapEntry<'_, TypeId, V>;
88}
89
90impl<V> TypeIdMapExt<V> for TypeIdMap<V> {
91    #[inline]
92    fn insert_type<T: ?Sized + 'static>(&mut self, v: V) -> Option<V> {
93        self.insert(TypeId::of::<T>(), v)
94    }
95
96    #[inline]
97    fn get_type<T: ?Sized + 'static>(&self) -> Option<&V> {
98        self.get(&TypeId::of::<T>())
99    }
100
101    #[inline]
102    fn get_type_mut<T: ?Sized + 'static>(&mut self) -> Option<&mut V> {
103        self.get_mut(&TypeId::of::<T>())
104    }
105
106    #[inline]
107    fn remove_type<T: ?Sized + 'static>(&mut self) -> Option<V> {
108        self.shift_remove(&TypeId::of::<T>())
109    }
110
111    #[inline]
112    fn entry_type<T: ?Sized + 'static>(&mut self) -> TypeIdMapEntry<'_, TypeId, V> {
113        self.entry(TypeId::of::<T>())
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use static_assertions::assert_impl_all;
121
122    // Check that the HashMaps are Clone if the key/values are Clone
123    assert_impl_all!(PreHashMap::<u64, usize>: Clone);
124
125    #[test]
126    fn fast_typeid_hash() {
127        struct Hasher;
128
129        impl core::hash::Hasher for Hasher {
130            fn finish(&self) -> u64 {
131                0
132            }
133            fn write(&mut self, _: &[u8]) {
134                panic!("Hashing of core::any::TypeId changed");
135            }
136            fn write_u64(&mut self, _: u64) {}
137        }
138
139        Hash::hash(&TypeId::of::<()>(), &mut Hasher);
140    }
141
142    crate::cfg::alloc! {
143        #[test]
144        fn stable_hash_within_same_program_execution() {
145            use alloc::vec::Vec;
146
147            let mut map_1 = <HashMap<_, _>>::default();
148            let mut map_2 = <HashMap<_, _>>::default();
149            for i in 1..10 {
150                map_1.insert(i, i);
151                map_2.insert(i, i);
152            }
153            assert_eq!(
154                map_1.iter().collect::<Vec<_>>(),
155                map_2.iter().collect::<Vec<_>>()
156            );
157        }
158    }
159}
bevy_utils/map.rs

bevy_utils/
map.rs
