Merge pull request #245 from cuviper/more-reserves

Add more reservation abilities

Author: cuviper

RELEASES.md

@@ -24,6 +24,11 @@
     `par_sort_by_cached_key` methods which perform stable sorts in place
     using a key extraction function.
 
+  - `IndexMap` and `IndexSet` now have `reserve_exact`, `try_reserve`, and
+    `try_reserve_exact` methods that correspond to the same methods on `Vec`.
+    However, exactness only applies to the direct capacity for items, while the
+    raw hash table still follows its own rules for capacity and load factor.
+
   - The `hashbrown` dependency has been updated to version 0.13.
 
 - 1.9.1

src/lib.rs

@@ -191,3 +191,58 @@ trait Entries {
     where
         F: FnOnce(&mut [Self::Entry]);
 }
+
+/// The error type for `try_reserve` methods.
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct TryReserveError {
+    kind: TryReserveErrorKind,
+}
+
+#[derive(Clone, PartialEq, Eq, Debug)]
+enum TryReserveErrorKind {
+    // The standard library's kind is currently opaque to us, otherwise we could unify this.
+    Std(alloc::collections::TryReserveError),
+    CapacityOverflow,
+    AllocError { layout: alloc::alloc::Layout },
+}
+
+// These are not `From` so we don't expose them in our public API.
+impl TryReserveError {
+    fn from_alloc(error: alloc::collections::TryReserveError) -> Self {
+        Self {
+            kind: TryReserveErrorKind::Std(error),
+        }
+    }
+
+    fn from_hashbrown(error: hashbrown::TryReserveError) -> Self {
+        Self {
+            kind: match error {
+                hashbrown::TryReserveError::CapacityOverflow => {
+                    TryReserveErrorKind::CapacityOverflow
+                }
+                hashbrown::TryReserveError::AllocError { layout } => {
+                    TryReserveErrorKind::AllocError { layout }
+                }
+            },
+        }
+    }
+}
+
+impl core::fmt::Display for TryReserveError {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        let reason = match &self.kind {
+            TryReserveErrorKind::Std(e) => return core::fmt::Display::fmt(e, f),
+            TryReserveErrorKind::CapacityOverflow => {
+                " because the computed capacity exceeded the collection's maximum"
+            }
+            TryReserveErrorKind::AllocError { .. } => {
+                " because the memory allocator returned an error"
+            }
+        };
+        f.write_str("memory allocation failed")?;
+        f.write_str(reason)
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for TryReserveError {}

src/map.rs

@@ -31,7 +31,7 @@ use std::collections::hash_map::RandomState;
 use self::core::IndexMapCore;
 use crate::equivalent::Equivalent;
 use crate::util::{third, try_simplify_range};
-use crate::{Bucket, Entries, HashValue};
+use crate::{Bucket, Entries, HashValue, TryReserveError};
 
 /// A hash table where the iteration order of the key-value pairs is independent
 /// of the hash values of the keys.
@@ -319,6 +319,37 @@ where
         self.core.reserve(additional);
     }
 
+    /// Reserve capacity for `additional` more key-value pairs, without over-allocating.
+    ///
+    /// Unlike `reserve`, this does not deliberately over-allocate the entry capacity to avoid
+    /// frequent re-allocations. However, the underlying data structures may still have internal
+    /// capacity requirements, and the allocator itself may give more space than requested, so this
+    /// cannot be relied upon to be precisely minimal.
+    ///
+    /// Computes in **O(n)** time.
+    pub fn reserve_exact(&mut self, additional: usize) {
+        self.core.reserve_exact(additional);
+    }
+
+    /// Try to reserve capacity for `additional` more key-value pairs.
+    ///
+    /// Computes in **O(n)** time.
+    pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
+        self.core.try_reserve(additional)
+    }
+
+    /// Try to reserve capacity for `additional` more key-value pairs, without over-allocating.
+    ///
+    /// Unlike `try_reserve`, this does not deliberately over-allocate the entry capacity to avoid
+    /// frequent re-allocations. However, the underlying data structures may still have internal
+    /// capacity requirements, and the allocator itself may give more space than requested, so this
+    /// cannot be relied upon to be precisely minimal.
+    ///
+    /// Computes in **O(n)** time.
+    pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> {
+        self.core.try_reserve_exact(additional)
+    }
+
     /// Shrink the capacity of the map as much as possible.
     ///
     /// Computes in **O(n)** time.

src/map/core.rs

@@ -12,9 +12,9 @@ mod raw;
 use hashbrown::raw::RawTable;
 
 use crate::vec::{Drain, Vec};
-use core::cmp;
+use crate::TryReserveError;
 use core::fmt;
-use core::mem::replace;
+use core::mem;
 use core::ops::RangeBounds;
 
 use crate::equivalent::Equivalent;
@@ -62,18 +62,18 @@ where
     V: Clone,
 {
     fn clone(&self) -> Self {
-        let indices = self.indices.clone();
-        let mut entries = Vec::with_capacity(indices.capacity());
-        entries.clone_from(&self.entries);
-        IndexMapCore { indices, entries }
+        let mut new = Self::new();
+        new.clone_from(self);
+        new
     }
 
     fn clone_from(&mut self, other: &Self) {
         let hasher = get_hash(&other.entries);
         self.indices.clone_from_with_hasher(&other.indices, hasher);
         if self.entries.capacity() < other.entries.len() {
-            // If we must resize, match the indices capacity
-            self.reserve_entries();
+            // If we must resize, match the indices capacity.
+            let additional = other.entries.len() - self.entries.len();
+            self.reserve_entries(additional);
         }
         self.entries.clone_from(&other.entries);
     }
@@ -120,6 +120,9 @@ impl<K, V> Entries for IndexMapCore<K, V> {
 }
 
 impl<K, V> IndexMapCore<K, V> {
+    /// The maximum capacity before the `entries` allocation would exceed `isize::MAX`.
+    const MAX_ENTRIES_CAPACITY: usize = (isize::MAX as usize) / mem::size_of::<Bucket<K, V>>();
+
     #[inline]
     pub(crate) const fn new() -> Self {
         IndexMapCore {
@@ -143,7 +146,7 @@ impl<K, V> IndexMapCore<K, V> {
 
     #[inline]
     pub(crate) fn capacity(&self) -> usize {
-        cmp::min(self.indices.capacity(), self.entries.capacity())
+        Ord::min(self.indices.capacity(), self.entries.capacity())
     }
 
     pub(crate) fn clear(&mut self) {
@@ -193,15 +196,67 @@ impl<K, V> IndexMapCore<K, V> {
     /// Reserve capacity for `additional` more key-value pairs.
     pub(crate) fn reserve(&mut self, additional: usize) {
         self.indices.reserve(additional, get_hash(&self.entries));
-        self.reserve_entries();
+        // Only grow entries if necessary, since we also round up capacity.
+        if additional > self.entries.capacity() - self.entries.len() {
+            self.reserve_entries(additional);
+        }
+    }
+
+    /// Reserve entries capacity, rounded up to match the indices
+    fn reserve_entries(&mut self, additional: usize) {
+        // Use a soft-limit on the maximum capacity, but if the caller explicitly
+        // requested more, do it and let them have the resulting panic.
+        let new_capacity = Ord::min(self.indices.capacity(), Self::MAX_ENTRIES_CAPACITY);
+        let try_add = new_capacity - self.entries.len();
+        if try_add > additional && self.entries.try_reserve_exact(try_add).is_ok() {
+            return;
+        }
+        self.entries.reserve_exact(additional);
     }
 
-    /// Reserve entries capacity to match the indices
-    fn reserve_entries(&mut self) {
-        let additional = self.indices.capacity() - self.entries.len();
+    /// Reserve capacity for `additional` more key-value pairs, without over-allocating.
+    pub(crate) fn reserve_exact(&mut self, additional: usize) {
+        self.indices.reserve(additional, get_hash(&self.entries));
         self.entries.reserve_exact(additional);
     }
 
+    /// Try to reserve capacity for `additional` more key-value pairs.
+    pub(crate) fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
+        self.indices
+            .try_reserve(additional, get_hash(&self.entries))
+            .map_err(TryReserveError::from_hashbrown)?;
+        // Only grow entries if necessary, since we also round up capacity.
+        if additional > self.entries.capacity() - self.entries.len() {
+            self.try_reserve_entries(additional)
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Try to reserve entries capacity, rounded up to match the indices
+    fn try_reserve_entries(&mut self, additional: usize) -> Result<(), TryReserveError> {
+        // Use a soft-limit on the maximum capacity, but if the caller explicitly
+        // requested more, do it and let them have the resulting error.
+        let new_capacity = Ord::min(self.indices.capacity(), Self::MAX_ENTRIES_CAPACITY);
+        let try_add = new_capacity - self.entries.len();
+        if try_add > additional && self.entries.try_reserve_exact(try_add).is_ok() {
+            return Ok(());
+        }
+        self.entries
+            .try_reserve_exact(additional)
+            .map_err(TryReserveError::from_alloc)
+    }
+
+    /// Try to reserve capacity for `additional` more key-value pairs, without over-allocating.
+    pub(crate) fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> {
+        self.indices
+            .try_reserve(additional, get_hash(&self.entries))
+            .map_err(TryReserveError::from_hashbrown)?;
+        self.entries
+            .try_reserve_exact(additional)
+            .map_err(TryReserveError::from_alloc)
+    }
+
     /// Shrink the capacity of the map with a lower bound
     pub(crate) fn shrink_to(&mut self, min_capacity: usize) {
         self.indices
@@ -228,7 +283,7 @@ impl<K, V> IndexMapCore<K, V> {
         if i == self.entries.capacity() {
             // Reserve our own capacity synced to the indices,
             // rather than letting `Vec::push` just double it.
-            self.reserve_entries();
+            self.reserve_entries(1);
         }
         self.entries.push(Bucket { hash, key, value });
         i
@@ -248,7 +303,7 @@ impl<K, V> IndexMapCore<K, V> {
         K: Eq,
     {
         match self.get_index_of(hash, &key) {
-            Some(i) => (i, Some(replace(&mut self.entries[i].value, value))),
+            Some(i) => (i, Some(mem::replace(&mut self.entries[i].value, value))),
             None => (self.push(hash, key, value), None),
         }
     }
@@ -601,7 +656,7 @@ pub use self::raw::OccupiedEntry;
 impl<K, V> OccupiedEntry<'_, K, V> {
     /// Sets the value of the entry to `value`, and returns the entry's old value.
     pub fn insert(&mut self, value: V) -> V {
-        replace(self.get_mut(), value)
+        mem::replace(self.get_mut(), value)
     }
 
     /// Remove the key, value pair stored in the map for this entry, and return the value.

src/map/tests.rs

@@ -158,6 +158,15 @@ fn reserve() {
     assert_eq!(map.get(&capacity), Some(&std::usize::MAX));
 }
 
+#[test]
+fn try_reserve() {
+    let mut map = IndexMap::<usize, usize>::new();
+    assert_eq!(map.capacity(), 0);
+    assert_eq!(map.try_reserve(100), Ok(()));
+    assert!(map.capacity() >= 100);
+    assert!(map.try_reserve(usize::MAX).is_err());
+}
+
 #[test]
 fn shrink_to_fit() {
     let mut map = IndexMap::<usize, usize>::new();

src/set.rs

@@ -11,6 +11,7 @@ pub use self::slice::Slice;
 
 #[cfg(feature = "rayon")]
 pub use crate::rayon::set as rayon;
+use crate::TryReserveError;
 
 #[cfg(feature = "std")]
 use std::collections::hash_map::RandomState;
@@ -266,6 +267,37 @@ where
         self.map.reserve(additional);
     }
 
+    /// Reserve capacity for `additional` more values, without over-allocating.
+    ///
+    /// Unlike `reserve`, this does not deliberately over-allocate the entry capacity to avoid
+    /// frequent re-allocations. However, the underlying data structures may still have internal
+    /// capacity requirements, and the allocator itself may give more space than requested, so this
+    /// cannot be relied upon to be precisely minimal.
+    ///
+    /// Computes in **O(n)** time.
+    pub fn reserve_exact(&mut self, additional: usize) {
+        self.map.reserve_exact(additional);
+    }
+
+    /// Try to reserve capacity for `additional` more values.
+    ///
+    /// Computes in **O(n)** time.
+    pub fn try_reserve(&mut self, additional: usize) -> Result<(), TryReserveError> {
+        self.map.try_reserve(additional)
+    }
+
+    /// Try to reserve capacity for `additional` more values, without over-allocating.
+    ///
+    /// Unlike `try_reserve`, this does not deliberately over-allocate the entry capacity to avoid
+    /// frequent re-allocations. However, the underlying data structures may still have internal
+    /// capacity requirements, and the allocator itself may give more space than requested, so this
+    /// cannot be relied upon to be precisely minimal.
+    ///
+    /// Computes in **O(n)** time.
+    pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), TryReserveError> {
+        self.map.try_reserve_exact(additional)
+    }
+
     /// Shrink the capacity of the set as much as possible.
     ///
     /// Computes in **O(n)** time.

src/set/tests.rs

@@ -282,6 +282,15 @@ fn reserve() {
     assert_eq!(set.get(&capacity), Some(&capacity));
 }
 
+#[test]
+fn try_reserve() {
+    let mut set = IndexSet::<usize>::new();
+    assert_eq!(set.capacity(), 0);
+    assert_eq!(set.try_reserve(100), Ok(()));
+    assert!(set.capacity() >= 100);
+    assert!(set.try_reserve(usize::MAX).is_err());
+}
+
 #[test]
 fn shrink_to_fit() {
     let mut set = IndexSet::<usize>::new();