Skip to main content

bevy_ecs/entity/
hash_map.rs

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