bevy_utils/
map.rs

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

bevy_utils/
map.rs
