Add reserve_exact, try_reserve, try_reserve_exact

Author: cuviper

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,6 +12,7 @@ mod raw;
 use hashbrown::raw::RawTable;
 
 use crate::vec::{Drain, Vec};
+use crate::TryReserveError;
 use core::cmp;
 use core::fmt;
 use core::mem::replace;
@@ -202,6 +203,38 @@ impl<K, V> IndexMapCore<K, V> {
         self.entries.reserve_exact(additional);
     }
 
+    /// 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)?;
+        self.try_reserve_entries()
+    }
+
+    /// Try to reserve entries capacity to match the indices
+    fn try_reserve_entries(&mut self) -> Result<(), TryReserveError> {
+        let additional = self.indices.capacity() - self.entries.len();
+        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

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();