Skip to main content

bevy_ecs/entity/
hash_map.rs

1//! Contains the [`EntityEquivalentHashMap`] type, a [`HashMap`] pre-configured to use [`EntityHash`] hashing.
2//!
3//! This module is a lightweight wrapper around Bevy's [`HashMap`] that is more performant for [`Entity`] keys.
4
5use core::{
6    fmt::{self, Debug, Formatter},
7    hash::Hash,
8    iter::FusedIterator,
9    marker::PhantomData,
10    ops::{Deref, DerefMut, Index},
11};
12
13use bevy_platform::collections::hash_map::{self, HashMap};
14#[cfg(feature = "bevy_reflect")]
15use bevy_reflect::Reflect;
16
17use super::{Entity, EntityEquivalent, EntityHash, EntitySetIterator};
18
19/// A [`HashMap`] pre-configured to use [`EntityHash`] hashing.
20#[cfg_attr(feature = "bevy_reflect", derive(Reflect))]
21#[cfg_attr(feature = "serialize", derive(serde::Deserialize, serde::Serialize))]
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct EntityEquivalentHashMap<K: EntityEquivalent + Hash, V>(HashMap<K, V, EntityHash>);
24
25/// A [`HashMap`] pre-configured to use [`EntityHash`] hashing with an [`Entity`].
26pub type EntityHashMap<V> = EntityEquivalentHashMap<Entity, V>;
27
28impl<K: EntityEquivalent + Hash, V> EntityEquivalentHashMap<K, V> {
29    /// Creates an empty `EntityEquivalentHashMap`.
30    ///
31    /// Equivalent to [`HashMap::with_hasher(EntityHash)`].
32    ///
33    /// [`HashMap::with_hasher(EntityHash)`]: HashMap::with_hasher
34    pub const fn new() -> Self {
35        Self(HashMap::with_hasher(EntityHash))
36    }
37
38    /// Creates an empty `EntityEquivalentHashMap` with the specified capacity.
39    ///
40    /// Equivalent to [`HashMap::with_capacity_and_hasher(n, EntityHash)`].
41    ///
42    /// [`HashMap::with_capacity_and_hasher(n, EntityHash)`]: HashMap::with_capacity_and_hasher
43    pub fn with_capacity(n: usize) -> Self {
44        Self(HashMap::with_capacity_and_hasher(n, EntityHash))
45    }
46
47    /// Constructs an `EntityEquivalentHashMap` from an [`HashMap`].
48    pub const fn from_index_map(set: HashMap<K, V, EntityHash>) -> Self {
49        Self(set)
50    }
51
52    /// Returns the inner [`HashMap`].
53    pub fn into_inner(self) -> HashMap<K, V, EntityHash> {
54        self.0
55    }
56
57    /// An iterator visiting all keys in arbitrary order.
58    /// The iterator element type is `&'a K`.
59    ///
60    /// Equivalent to [`HashMap::keys`].
61    pub fn keys(&self) -> Keys<'_, K, V> {
62        Keys(self.0.keys(), PhantomData)
63    }
64
65    /// Creates a consuming iterator visiting all the keys in arbitrary order.
66    /// The map cannot be used after calling this.
67    /// The iterator element type is [`Entity`].
68    ///
69    /// Equivalent to [`HashMap::into_keys`].
70    pub fn into_keys(self) -> IntoKeys<K, V> {
71        IntoKeys(self.0.into_keys(), PhantomData)
72    }
73}
74
75impl<K: EntityEquivalent + Hash, V> Default for EntityEquivalentHashMap<K, V> {
76    fn default() -> Self {
77        Self(Default::default())
78    }
79}
80
81impl<K: EntityEquivalent + Hash, V> Deref for EntityEquivalentHashMap<K, V> {
82    type Target = HashMap<K, V, EntityHash>;
83
84    fn deref(&self) -> &Self::Target {
85        &self.0
86    }
87}
88
89impl<K: EntityEquivalent + Hash, V> DerefMut for EntityEquivalentHashMap<K, V> {
90    fn deref_mut(&mut self) -> &mut Self::Target {
91        &mut self.0
92    }
93}
94
95impl<'a, K: EntityEquivalent + Hash + Copy, V: Copy> Extend<&'a (K, V)>
96    for EntityEquivalentHashMap<K, V>
97{
98    fn extend<I: IntoIterator<Item = &'a (K, V)>>(&mut self, iter: I) {
99        self.0.extend(iter);
100    }
101}
102
103impl<'a, K: EntityEquivalent + Hash + Copy, V: Copy> Extend<(&'a K, &'a V)>
104    for EntityEquivalentHashMap<K, V>
105{
106    fn extend<I: IntoIterator<Item = (&'a K, &'a V)>>(&mut self, iter: I) {
107        self.0.extend(iter);
108    }
109}
110
111impl<K: EntityEquivalent + Hash, V> Extend<(K, V)> for EntityEquivalentHashMap<K, V> {
112    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
113        self.0.extend(iter);
114    }
115}
116
117impl<K: EntityEquivalent + Hash, V, const N: usize> From<[(K, V); N]>
118    for EntityEquivalentHashMap<K, V>
119{
120    fn from(value: [(K, V); N]) -> Self {
121        Self(HashMap::from_iter(value))
122    }
123}
124
125impl<K: EntityEquivalent + Hash, V> FromIterator<(K, V)> for EntityEquivalentHashMap<K, V> {
126    fn from_iter<I: IntoIterator<Item = (K, V)>>(iterable: I) -> Self {
127        Self(HashMap::from_iter(iterable))
128    }
129}
130
131impl<K: EntityEquivalent + Hash, V> From<HashMap<K, V, EntityHash>>
132    for EntityEquivalentHashMap<K, V>
133{
134    fn from(value: HashMap<K, V, EntityHash>) -> Self {
135        Self(value)
136    }
137}
138
139// `EntityEquivalent` does not guarantee maintained equality on conversions from one implementor to another,
140// so we restrict this impl to only keys of type `Entity`.
141impl<V, Q: EntityEquivalent + Hash + ?Sized> Index<&Q> for EntityHashMap<V> {
142    type Output = V;
143
144    fn index(&self, key: &Q) -> &V {
145        self.0.index(&key.entity())
146    }
147}
148
149impl<'a, K: EntityEquivalent + Hash, V> IntoIterator for &'a EntityEquivalentHashMap<K, V> {
150    type Item = (&'a K, &'a V);
151    type IntoIter = hash_map::Iter<'a, K, V>;
152
153    fn into_iter(self) -> Self::IntoIter {
154        self.0.iter()
155    }
156}
157
158impl<'a, K: EntityEquivalent + Hash, V> IntoIterator for &'a mut EntityEquivalentHashMap<K, V> {
159    type Item = (&'a K, &'a mut V);
160    type IntoIter = hash_map::IterMut<'a, K, V>;
161
162    fn into_iter(self) -> Self::IntoIter {
163        self.0.iter_mut()
164    }
165}
166
167impl<K: EntityEquivalent + Hash, V> IntoIterator for EntityEquivalentHashMap<K, V> {
168    type Item = (K, V);
169    type IntoIter = hash_map::IntoIter<K, V>;
170
171    fn into_iter(self) -> Self::IntoIter {
172        self.0.into_iter()
173    }
174}
175
176/// An iterator over the keys of a [`EntityEquivalentHashMap`] in arbitrary order.
177/// The iterator element type is `&'a K`.
178///
179/// This struct is created by the [`keys`] method on [`EntityEquivalentHashMap`]. See its documentation for more.
180///
181/// [`keys`]: EntityEquivalentHashMap::keys
182pub struct Keys<'a, K: EntityEquivalent + Hash, V, S = EntityHash>(
183    hash_map::Keys<'a, K, V>,
184    PhantomData<S>,
185);
186
187impl<'a, K: EntityEquivalent + Hash, V> Keys<'a, K, V> {
188    /// Constructs a [`Keys<'a, K, V, S>`] from a [`hash_map::Keys<'a, K, V>`] unsafely.
189    ///
190    /// # Safety
191    ///
192    /// `keys` must either be empty, or have been obtained from a
193    /// [`hash_map::HashMap`] using the `S` hasher.
194    pub const unsafe fn from_keys_unchecked<S>(
195        keys: hash_map::Keys<'a, K, V>,
196    ) -> Keys<'a, K, V, S> {
197        Keys(keys, PhantomData)
198    }
199
200    /// Returns the inner [`Keys`](hash_map::Keys).
201    pub const fn into_inner(self) -> hash_map::Keys<'a, K, V> {
202        self.0
203    }
204}
205
206impl<'a, K: EntityEquivalent + Hash, V> Deref for Keys<'a, K, V> {
207    type Target = hash_map::Keys<'a, K, V>;
208
209    fn deref(&self) -> &Self::Target {
210        &self.0
211    }
212}
213
214impl<'a, K: EntityEquivalent + Hash, V> Iterator for Keys<'a, K, V> {
215    type Item = &'a K;
216
217    fn next(&mut self) -> Option<Self::Item> {
218        self.0.next()
219    }
220
221    fn size_hint(&self) -> (usize, Option<usize>) {
222        self.0.size_hint()
223    }
224
225    fn fold<B, F>(self, init: B, f: F) -> B
226    where
227        Self: Sized,
228        F: FnMut(B, Self::Item) -> B,
229    {
230        self.0.fold(init, f)
231    }
232}
233
234impl<K: EntityEquivalent + Hash, V> ExactSizeIterator for Keys<'_, K, V> {}
235
236impl<K: EntityEquivalent + Hash, V> FusedIterator for Keys<'_, K, V> {}
237
238impl<K: EntityEquivalent + Hash, V> Clone for Keys<'_, K, V> {
239    fn clone(&self) -> Self {
240        // SAFETY: We are cloning an already valid `Keys`.
241        unsafe { Self::from_keys_unchecked(self.0.clone()) }
242    }
243}
244
245impl<K: EntityEquivalent + Hash + Debug, V: Debug> Debug for Keys<'_, K, V> {
246    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
247        f.debug_tuple("Keys").field(&self.0).field(&self.1).finish()
248    }
249}
250
251impl<K: EntityEquivalent + Hash, V> Default for Keys<'_, K, V> {
252    fn default() -> Self {
253        // SAFETY: `Keys` is empty.
254        unsafe { Self::from_keys_unchecked(Default::default()) }
255    }
256}
257
258// SAFETY: Keys stems from a correctly behaving `HashMap<K, V, EntityHash>`.
259unsafe impl<K: EntityEquivalent + Hash, V> EntitySetIterator for Keys<'_, K, V> {}
260
261/// An owning iterator over the keys of a [`EntityEquivalentHashMap`] in arbitrary order.
262/// The iterator element type is [`Entity`].
263///
264/// This struct is created by the [`into_keys`] method on [`EntityEquivalentHashMap`].
265/// See its documentation for more.
266/// The map cannot be used after calling that method.
267///
268/// [`into_keys`]: EntityEquivalentHashMap::into_keys
269pub struct IntoKeys<K: EntityEquivalent + Hash, V, S = EntityHash>(
270    hash_map::IntoKeys<K, V>,
271    PhantomData<S>,
272);
273
274impl<K: EntityEquivalent + Hash, V> IntoKeys<K, V> {
275    /// Constructs a [`IntoKeys<K, V, S>`] from a [`hash_map::IntoKeys<K, V>`] unsafely.
276    ///
277    /// # Safety
278    ///
279    /// `into_keys` must either be empty, or have been obtained from a
280    /// [`hash_map::HashMap`] using the `S` hasher.
281    pub const unsafe fn from_into_keys_unchecked<S>(
282        into_keys: hash_map::IntoKeys<K, V>,
283    ) -> IntoKeys<K, V, S> {
284        IntoKeys(into_keys, PhantomData)
285    }
286
287    /// Returns the inner [`IntoKeys`](hash_map::IntoKeys).
288    pub fn into_inner(self) -> hash_map::IntoKeys<K, V> {
289        self.0
290    }
291}
292
293impl<K: EntityEquivalent + Hash, V> Deref for IntoKeys<K, V> {
294    type Target = hash_map::IntoKeys<K, V>;
295
296    fn deref(&self) -> &Self::Target {
297        &self.0
298    }
299}
300
301impl<K: EntityEquivalent + Hash, V> Iterator for IntoKeys<K, V> {
302    type Item = K;
303
304    fn next(&mut self) -> Option<Self::Item> {
305        self.0.next()
306    }
307
308    fn size_hint(&self) -> (usize, Option<usize>) {
309        self.0.size_hint()
310    }
311
312    fn fold<B, F>(self, init: B, f: F) -> B
313    where
314        Self: Sized,
315        F: FnMut(B, Self::Item) -> B,
316    {
317        self.0.fold(init, f)
318    }
319}
320
321impl<K: EntityEquivalent + Hash, V> ExactSizeIterator for IntoKeys<K, V> {}
322
323impl<K: EntityEquivalent + Hash, V> FusedIterator for IntoKeys<K, V> {}
324
325impl<K: EntityEquivalent + Hash + Debug, V: Debug> Debug for IntoKeys<K, V> {
326    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
327        f.debug_tuple("IntoKeys")
328            .field(&self.0)
329            .field(&self.1)
330            .finish()
331    }
332}
333
334impl<K: EntityEquivalent + Hash, V> Default for IntoKeys<K, V> {
335    fn default() -> Self {
336        // SAFETY: `IntoKeys` is empty.
337        unsafe { Self::from_into_keys_unchecked(Default::default()) }
338    }
339}
340
341// SAFETY: IntoKeys stems from a correctly behaving `HashMap<K, V, EntityHash>`.
342unsafe impl<K: EntityEquivalent + Hash, V> EntitySetIterator for IntoKeys<K, V> {}
343
344#[cfg(test)]
345mod tests {
346    use super::*;
347    use bevy_reflect::Reflect;
348    use static_assertions::assert_impl_all;
349
350    // Check that the HashMaps are Clone if the key/values are Clone
351    assert_impl_all!(EntityHashMap::<usize>: Clone);
352    // EntityEquivalentHashMap should implement Reflect
353    #[cfg(feature = "bevy_reflect")]
354    assert_impl_all!(EntityHashMap::<i32>: Reflect);
355}