Implement try_reserve

Author: Amanieu

src/lib.rs

@@ -80,3 +80,13 @@ pub mod hash_set {
 
 pub use map::HashMap;
 pub use set::HashSet;
+
+/// Augments `AllocErr` with a CapacityOverflow variant.
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub enum CollectionAllocErr {
+    /// Error due to the computed capacity exceeding the collection's maximum
+    /// (usually `isize::MAX` bytes).
+    CapacityOverflow,
+    /// Error due to the allocator (see the `AllocErr` type's docs).
+    AllocErr,
+}

src/map.rs

@@ -18,6 +18,7 @@ use core::marker::PhantomData;
 use core::mem;
 use core::ops::Index;
 use raw::{Bucket, RawDrain, RawIntoIter, RawIter, RawTable};
+use CollectionAllocErr;
 
 pub use fx::FxHashBuilder as DefaultHashBuilder;
 
@@ -361,6 +362,29 @@ where
             .reserve(additional, |x| make_hash(hash_builder, &x.0));
     }
 
+    /// Tries to reserve capacity for at least `additional` more elements to be inserted
+    /// in the given `HashMap<K,V>`. The collection may reserve more space to avoid
+    /// frequent reallocations.
+    ///
+    /// # Errors
+    ///
+    /// If the capacity overflows, or the allocator reports a failure, then an error
+    /// is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use hashbrown::HashMap;
+    /// let mut map: HashMap<&str, isize> = HashMap::new();
+    /// map.try_reserve(10).expect("why is the test harness OOMing on 10 bytes?");
+    /// ```
+    #[inline]
+    pub fn try_reserve(&mut self, additional: usize) -> Result<(), CollectionAllocErr> {
+        let hash_builder = &self.hash_builder;
+        self.table
+            .try_reserve(additional, |x| make_hash(hash_builder, &x.0))
+    }
+
     /// Shrinks the capacity of the map as much as possible. It will drop
     /// down as much as possible while maintaining the internal rules
     /// and possibly leaving some space in accordance with the resize policy.
@@ -2519,10 +2543,10 @@ mod test_map {
     use super::HashMap;
     use rand::{thread_rng, Rng};
     use std::cell::RefCell;
-    #[cfg(try_reserve_not_implemented)]
     use std::mem::size_of;
     use std::usize;
     use std::vec::Vec;
+    use CollectionAllocErr::*;
 
     #[test]
     fn test_zero_capacities() {
@@ -3397,21 +3421,20 @@ mod test_map {
         panic!("Adaptive early resize failed");
     }
 
-    #[cfg(try_reserve_not_implemented)]
     #[test]
     fn test_try_reserve() {
         let mut empty_bytes: HashMap<u8, u8> = HashMap::new();
 
         const MAX_USIZE: usize = usize::MAX;
 
         // HashMap and RawTables use complicated size calculations
-        // hashes_size is sizeof(HashUint) * capacity;
+        // hashes_size is sizeof(u8) * capacity;
         // pairs_size is sizeof((K. V)) * capacity;
         // alignment_hashes_size is 8
         // alignment_pairs size is 4
-        let size_of_multiplier = (size_of::<usize>() + size_of::<(u8, u8)>()).next_power_of_two();
+        let size_of_multiplier = (size_of::<u8>() + size_of::<(u8, u8)>()).next_power_of_two();
         // The following formula is used to calculate the new capacity
-        let max_no_ovf = ((MAX_USIZE / 11) * 10) / size_of_multiplier - 1;
+        let max_no_ovf = ((MAX_USIZE / 8) * 7) / size_of_multiplier - 1;
 
         if let Err(CapacityOverflow) = empty_bytes.try_reserve(MAX_USIZE) {
         } else {

src/raw/mod.rs

@@ -8,6 +8,7 @@ use core::mem::ManuallyDrop;
 use core::ops::Range;
 use core::ptr::NonNull;
 use scopeguard::guard;
+use CollectionAllocErr;
 
 // Branch prediction hint. This is currently only available on nightly but it
 // consistently improves performance by 10-15%.
@@ -61,6 +62,32 @@ mod bitmask;
 use self::bitmask::BitMask;
 use self::imp::Group;
 
+/// Whether memory allocation errors should return an error or abort.
+enum Fallibility {
+    Fallible,
+    Infallible,
+}
+
+impl Fallibility {
+    /// Error to return on capacity overflow.
+    #[inline]
+    fn capacity_overflow(&self) -> CollectionAllocErr {
+        match *self {
+            Fallibility::Fallible => CollectionAllocErr::CapacityOverflow,
+            Fallibility::Infallible => panic!("Hash table capacity overflow"),
+        }
+    }
+
+    /// Error to return on allocation error.
+    #[inline]
+    fn alloc_err(&self, layout: Layout) -> CollectionAllocErr {
+        match *self {
+            Fallibility::Fallible => CollectionAllocErr::AllocErr,
+            Fallibility::Infallible => handle_alloc_error(layout),
+        }
+    }
+}
+
 /// Control byte value for an empty bucket.
 const EMPTY: u8 = 0b11111111;
 
@@ -129,18 +156,23 @@ impl Iterator for ProbeSeq {
 
 /// Returns the number of buckets needed to hold the given number of items,
 /// taking the maximum load factor into account.
+///
+/// Returns `None` if an overflow occurs.
 #[inline]
-fn capacity_to_buckets(cap: usize) -> usize {
+fn capacity_to_buckets(cap: usize) -> Option<usize> {
     let adjusted_cap = if cap < 8 {
         // Need at least 1 free bucket on small tables
         cap + 1
     } else {
         // Otherwise require 1/8 buckets to be empty (87.5% load)
-        cap.checked_mul(8).expect("Hash table capacity overflow") / 7
+        //
+        // Be careful when modifying this, calculate_layout relies on the
+        // overflow check here.
+        cap.checked_mul(8)? / 7
     };
 
     // Any overflows will have been caught by the checked_mul.
-    adjusted_cap.next_power_of_two()
+    Some(adjusted_cap.next_power_of_two())
 }
 
 /// Returns the maximum effective capacity for the given bucket mask, taking
@@ -156,6 +188,8 @@ fn bucket_mask_to_capacity(bucket_mask: usize) -> usize {
 
 // Returns a Layout which describes the allocation required for a hash table,
 // and the offset of the buckets in the allocation.
+///
+/// Returns `None` if an overflow occurs.
 #[inline]
 #[cfg(feature = "nightly")]
 fn calculate_layout<T>(buckets: usize) -> Option<(Layout, usize)> {
@@ -169,10 +203,10 @@ fn calculate_layout<T>(buckets: usize) -> Option<(Layout, usize)> {
     // We add `Group::WIDTH` control bytes at the end of the array which
     // replicate the bytes at the start of the array and thus avoids the need to
     // perform bounds-checking while probing.
-    let ctrl = Layout::array::<u8>(buckets + Group::WIDTH)
-        .ok()?
-        .align_to(Group::WIDTH)
-        .ok()?;
+    //
+    // There is no possible overflow here since buckets is a power of two and
+    // Group::WIDTH is a small number.
+    let ctrl = unsafe { Layout::from_size_align_unchecked(buckets + Group::WIDTH, Group::WIDTH) };
 
     ctrl.extend(data).ok()
 }
@@ -188,12 +222,11 @@ fn calculate_layout<T>(buckets: usize) -> Option<(Layout, usize)> {
     let data_align = usize::max(mem::align_of::<T>(), Group::WIDTH);
     let data_offset = (buckets + Group::WIDTH).checked_add(data_align - 1)? & !(data_align - 1);
     let len = data_offset.checked_add(mem::size_of::<T>().checked_mul(buckets)?)?;
-    unsafe {
-        Some((
-            Layout::from_size_align_unchecked(len, data_align),
-            data_offset,
-        ))
-    }
+
+    Some((
+        unsafe { Layout::from_size_align_unchecked(len, data_align) },
+        data_offset,
+    ))
 }
 
 /// A reference to a hash table bucket containing a `T`.
@@ -271,28 +304,36 @@ impl<T> RawTable<T> {
     ///
     /// The control bytes are left uninitialized.
     #[inline]
-    unsafe fn new_uninitialized(buckets: usize) -> RawTable<T> {
+    unsafe fn new_uninitialized(
+        buckets: usize,
+        fallability: Fallibility,
+    ) -> Result<RawTable<T>, CollectionAllocErr> {
         let (layout, data_offset) =
-            calculate_layout::<T>(buckets).expect("Hash table capacity overflow");
-        let ctrl = NonNull::new(alloc(layout)).unwrap_or_else(|| handle_alloc_error(layout));
+            calculate_layout::<T>(buckets).ok_or_else(|| fallability.capacity_overflow())?;
+        let ctrl = NonNull::new(alloc(layout)).ok_or_else(|| fallability.alloc_err(layout))?;
         let data = NonNull::new_unchecked(ctrl.as_ptr().add(data_offset) as *mut T);
-        RawTable {
+        Ok(RawTable {
             data,
             ctrl,
             bucket_mask: buckets - 1,
             items: 0,
             growth_left: bucket_mask_to_capacity(buckets - 1),
-        }
+        })
     }
 
-    /// Allocates a new hash table with at least enough capacity for inserting
-    /// the given number of elements without reallocating.
-    pub fn with_capacity(capacity: usize) -> RawTable<T> {
+    /// Attempts to allocate a new hash table with at least enough capacity
+    /// for inserting the given number of elements without reallocating.
+    fn try_with_capacity(
+        capacity: usize,
+        fallability: Fallibility,
+    ) -> Result<RawTable<T>, CollectionAllocErr> {
         if capacity == 0 {
-            RawTable::new()
+            Ok(RawTable::new())
         } else {
             unsafe {
-                let result = RawTable::new_uninitialized(capacity_to_buckets(capacity));
+                let buckets =
+                    capacity_to_buckets(capacity).ok_or_else(|| fallability.capacity_overflow())?;
+                let result = RawTable::new_uninitialized(buckets, fallability)?;
                 result
                     .ctrl(0)
                     .write_bytes(EMPTY, result.buckets() + Group::WIDTH);
@@ -306,11 +347,18 @@ impl<T> RawTable<T> {
                         .write_bytes(DELETED, Group::WIDTH - result.buckets());
                 }
 
-                result
+                Ok(result)
             }
         }
     }
 
+    /// Allocates a new hash table with at least enough capacity for inserting
+    /// the given number of elements without reallocating.
+    pub fn with_capacity(capacity: usize) -> RawTable<T> {
+        RawTable::try_with_capacity(capacity, Fallibility::Infallible)
+            .unwrap_or_else(|_| unsafe { hint::unreachable_unchecked() })
+    }
+
     /// Deallocates the table without dropping any entries.
     #[inline]
     unsafe fn free_buckets(&mut self) {
@@ -473,8 +521,9 @@ impl<T> RawTable<T> {
     #[inline]
     pub fn shrink_to(&mut self, min_size: usize, hasher: impl Fn(&T) -> u64) {
         let min_size = usize::max(self.items, min_size);
-        if bucket_mask_to_capacity(self.bucket_mask) >= min_size * 2 {
-            self.resize(min_size, hasher);
+        if self.bucket_mask != 0 && bucket_mask_to_capacity(self.bucket_mask) >= min_size * 2 {
+            self.resize(min_size, hasher, Fallibility::Infallible)
+                .unwrap_or_else(|_| unsafe { hint::unreachable_unchecked() });
         }
     }
 
@@ -483,25 +532,47 @@ impl<T> RawTable<T> {
     #[inline]
     pub fn reserve(&mut self, additional: usize, hasher: impl Fn(&T) -> u64) {
         if additional > self.growth_left {
-            self.reserve_rehash(additional, hasher);
+            self.reserve_rehash(additional, hasher, Fallibility::Infallible)
+                .unwrap_or_else(|_| unsafe { hint::unreachable_unchecked() });
+        }
+    }
+
+    /// Tries to ensure that at least `additional` items can be inserted into
+    /// the table without reallocation.
+    #[inline]
+    pub fn try_reserve(
+        &mut self,
+        additional: usize,
+        hasher: impl Fn(&T) -> u64,
+    ) -> Result<(), CollectionAllocErr> {
+        if additional > self.growth_left {
+            self.reserve_rehash(additional, hasher, Fallibility::Fallible)
+        } else {
+            Ok(())
         }
     }
 
-    /// Out-of-line slow path for `reserve`.
+    /// Out-of-line slow path for `reserve` and `try_reserve`.
     #[cold]
     #[inline(never)]
-    fn reserve_rehash(&mut self, additional: usize, hasher: impl Fn(&T) -> u64) {
+    fn reserve_rehash(
+        &mut self,
+        additional: usize,
+        hasher: impl Fn(&T) -> u64,
+        fallability: Fallibility,
+    ) -> Result<(), CollectionAllocErr> {
         let new_items = self
             .items
             .checked_add(additional)
-            .expect("Hash table capacity overflow");
+            .ok_or_else(|| fallability.capacity_overflow())?;
 
         // Rehash in-place without re-allocating if we have plenty of spare
         // capacity that is locked up due to DELETED entries.
         if new_items < bucket_mask_to_capacity(self.bucket_mask) / 2 {
             self.rehash_in_place(hasher);
+            Ok(())
         } else {
-            self.resize(new_items, hasher);
+            self.resize(new_items, hasher, fallability)
         }
     }
 
@@ -607,12 +678,17 @@ impl<T> RawTable<T> {
 
     /// Allocates a new table of a different size and moves the contents of the
     /// current table into it.
-    fn resize(&mut self, capacity: usize, hasher: impl Fn(&T) -> u64) {
+    fn resize(
+        &mut self,
+        capacity: usize,
+        hasher: impl Fn(&T) -> u64,
+        fallability: Fallibility,
+    ) -> Result<(), CollectionAllocErr> {
         unsafe {
             debug_assert!(self.items <= capacity);
 
             // Allocate and initialize the new table.
-            let mut new_table = RawTable::with_capacity(capacity);
+            let mut new_table = RawTable::try_with_capacity(capacity, fallability)?;
             new_table.growth_left -= self.items;
             new_table.items = self.items;
 
@@ -642,6 +718,8 @@ impl<T> RawTable<T> {
             // the items will not be dropped (since they have been moved into the
             // new table).
             mem::swap(self, &mut new_table);
+
+            Ok(())
         }
     }
 
@@ -762,7 +840,10 @@ impl<T: Clone> Clone for RawTable<T> {
             Self::new()
         } else {
             unsafe {
-                let mut new_table = ManuallyDrop::new(Self::new_uninitialized(self.buckets()));
+                let mut new_table = ManuallyDrop::new(
+                    Self::new_uninitialized(self.buckets(), Fallibility::Infallible)
+                        .unwrap_or_else(|_| hint::unreachable_unchecked()),
+                );
 
                 // Copy the control bytes unchanged. We do this in a single pass
                 self.ctrl(0)