From 52eb28ee4a8e53baa4cc919ebc36c2520b44437e Mon Sep 17 00:00:00 2001
From: Tiago Lam <tiagolam@gmail.com>
Date: Fri, 25 Jun 2021 12:07:48 +0100
Subject: [PATCH] bevy_utils: Re-introduce with_capacity().

PR #1235 had removed the AHashExt trait and respective
`with_capacity()`s implementations, leaving only the less ergonomic
`HashMap::with_capacity_and_hasher(size, Default::default())` option
available.

This re-introduces `AHashExt` and respective `with_capacity()`
implementations to give a more ergonimic way to set a HashMap / HashSet
capacity.

As a note, this has also been discussed and agreed on issue #2115, which
this PR addresses.

Fixes #2115.

Signed-off-by: Tiago Lam <tiagolam@gmail.com>
---
 crates/bevy_utils/src/lib.rs | 78 ++++++++++++++++++++++++++++++++++++
 1 file changed, 78 insertions(+)

diff --git a/crates/bevy_utils/src/lib.rs b/crates/bevy_utils/src/lib.rs
index 9d66179df301ce..f6ba4ea3c41d73 100644
--- a/crates/bevy_utils/src/lib.rs
+++ b/crates/bevy_utils/src/lib.rs
@@ -37,6 +37,28 @@ impl std::hash::BuildHasher for FixedState {
 /// AHash is designed for performance and is NOT cryptographically secure.
 pub type HashMap<K, V> = std::collections::HashMap<K, V, RandomState>;
 
+pub trait AHashExt {
+    fn with_capacity(capacity: usize) -> Self;
+}
+
+impl<K, V> AHashExt for HashMap<K, V> {
+    /// Creates an empty `HashMap` with the specified capacity with AHash.
+    ///
+    /// The hash map will be able to hold at least `capacity` elements without
+    /// reallocating. If `capacity` is 0, the hash map will not allocate.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bevy_utils::{HashMap, AHashExt};
+    /// let mut map: HashMap<&str, i32> = HashMap::with_capacity(10);
+    /// ```
+    #[inline]
+    fn with_capacity(capacity: usize) -> Self {
+        HashMap::with_capacity_and_hasher(capacity, RandomState::default())
+    }
+}
+
 /// A stable std hash map implementing AHash, a high speed keyed hashing algorithm
 /// intended for use in in-memory hashmaps.
 ///
@@ -46,12 +68,49 @@ pub type HashMap<K, V> = std::collections::HashMap<K, V, RandomState>;
 /// AHash is designed for performance and is NOT cryptographically secure.
 pub type StableHashMap<K, V> = std::collections::HashMap<K, V, FixedState>;
 
+impl<K, V> AHashExt for StableHashMap<K, V> {
+    /// Creates an empty `StableHashMap` with the specified capacity with AHash.
+    ///
+    /// The hash map will be able to hold at least `capacity` elements without
+    /// reallocating. If `capacity` is 0, the hash map will not allocate.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bevy_utils::{StableHashMap, AHashExt};
+    /// let mut map: StableHashMap<&str, i32> = StableHashMap::with_capacity(10);
+    /// ```
+    #[inline]
+    fn with_capacity(capacity: usize) -> Self {
+        StableHashMap::with_capacity_and_hasher(capacity, FixedState::default())
+    }
+}
+
 /// A std hash set implementing AHash, a high speed keyed hashing algorithm
 /// intended for use in in-memory hashmaps.
 ///
 /// AHash is designed for performance and is NOT cryptographically secure.
 pub type HashSet<K> = std::collections::HashSet<K, RandomState>;
 
+impl<K> AHashExt for HashSet<K> {
+    /// Creates an empty `HashSet` with the specified capacity with AHash.
+    ///
+    /// The hash set will be able to hold at least `capacity` elements without
+    /// reallocating. If `capacity` is 0, the hash set will not allocate.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bevy_utils::{HashSet, AHashExt};
+    /// let set: HashSet<i32> = HashSet::with_capacity(10);
+    /// assert!(set.capacity() >= 10);
+    /// ```
+    #[inline]
+    fn with_capacity(capacity: usize) -> Self {
+        HashSet::with_capacity_and_hasher(capacity, RandomState::default())
+    }
+}
+
 /// A stable std hash set implementing AHash, a high speed keyed hashing algorithm
 /// intended for use in in-memory hashmaps.
 ///
@@ -60,3 +119,22 @@ pub type HashSet<K> = std::collections::HashSet<K, RandomState>;
 ///
 /// AHash is designed for performance and is NOT cryptographically secure.
 pub type StableHashSet<K> = std::collections::HashSet<K, FixedState>;
+
+impl<K> AHashExt for StableHashSet<K> {
+    /// Creates an empty `StableHashSet` with the specified capacity with AHash.
+    ///
+    /// The hash set will be able to hold at least `capacity` elements without
+    /// reallocating. If `capacity` is 0, the hash set will not allocate.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bevy_utils::{StableHashSet, AHashExt};
+    /// let set: StableHashSet<i32> = StableHashSet::with_capacity(10);
+    /// assert!(set.capacity() >= 10);
+    /// ```
+    #[inline]
+    fn with_capacity(capacity: usize) -> Self {
+        StableHashSet::with_capacity_and_hasher(capacity, FixedState::default())
+    }
+}