Repository: fitzgen/bumpalo (rust)

Issue:
Add try_reserve/_exact methods for Vec?
I tried to add these try_* methods (since I need fallible allocation) to `Vec` by wrapping the `RawVec` methods of the same names. However, they still seem to be infallible due to the alloc impl always returning `Ok`? https://github.com/fitzgen/bumpalo/blob/master/src/lib.rs#L1076

Would it be possible to make them fallible?

Pull request description:
[WIP] Added try_alloc_layout, try_reserve, & try_reserve_exact
1) Made it a point to not make any breaking changes to the public API.
2) I implemented my own global allocator which always returns a null ptr to simulate OOM and it seemed like `Bump::new()` /  `Bump::with_capacity(0)` would always panic with `oom()`. Shouldn't it not try and reach for the allocator in that case? Any idea how to go about fixing that special case?
3) I named the new methods `try_reserve` and `try_reserve_exact` to match what raw/vec calls them, but I could rename them to `try_alloc_in_new_chunk_if_necessary` like you suggested if you wish.

Fixes #66 

Patch:
diff --git a/Cargo.lock b/Cargo.lock
index da21c72..11d35a5 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -52,6 +52,7 @@ version = "3.2.1"
 dependencies = [
  "criterion 0.3.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "quickcheck 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)",
+ "rand 0.7.2 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
diff --git a/Cargo.toml b/Cargo.toml
index 2ec3a28..8ca1506 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,6 +26,7 @@ required-features = ["collections"]
 [dev-dependencies]
 quickcheck = "0.9.0"
 criterion = "0.3.0"
+rand = "0.7"
 
 [features]
 default = []
diff --git a/src/collections/vec.rs b/src/collections/vec.rs
index 8459c78..2a6b414 100644
--- a/src/collections/vec.rs
+++ b/src/collections/vec.rs
@@ -85,6 +85,7 @@
 
 use super::raw_vec::RawVec;
 use crate::Bump;
+use crate::collections::CollectionAllocErr;
 use core::cmp::Ordering;
 use core::fmt;
 use core::hash::{self, Hash};
@@ -738,6 +739,63 @@ impl<'bump, T: 'bump> Vec<'bump, T> {
         self.buf.reserve_exact(self.len, additional);
     }
 
+    /// Attempts to reserve capacity for at least `additional` more elements to be inserted
+    /// in the given `Vec<'bump, T>`. The collection may reserve more space to avoid
+    /// frequent reallocations. After calling `try_reserve`, capacity will be
+    /// greater than or equal to `self.len() + additional`. Does nothing if
+    /// capacity is already sufficient.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the new capacity overflows `usize`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bumpalo::{Bump, collections::Vec};
+    ///
+    /// let b = Bump::new();
+    /// let mut vec = bumpalo::vec![in &b; 1];
+    /// vec.try_reserve(10).unwrap();
+    /// assert!(vec.capacity() >= 11);
+    /// ```
+    pub fn try_reserve(
+        &mut self,
+        additional: usize,
+    ) -> Result<(), CollectionAllocErr> {
+        self.buf.try_reserve(self.len, additional)
+    }
+
+    /// Attempts to reserve the minimum capacity for exactly `additional` more elements to
+    /// be inserted in the given `Vec<'bump, T>`. After calling `try_reserve_exact`,
+    /// capacity will be greater than or equal to `self.len() + additional`.
+    /// Does nothing if the capacity is already sufficient.
+    ///
+    /// Note that the allocator may give the collection more space than it
+    /// requests. Therefore capacity can not be relied upon to be precisely
+    /// minimal. Prefer `try_reserve` if future insertions are expected.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the new capacity overflows `usize`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bumpalo::{Bump, collections::Vec};
+    ///
+    /// let b = Bump::new();
+    /// let mut vec = bumpalo::vec![in &b; 1];
+    /// vec.try_reserve_exact(10).unwrap();
+    /// assert!(vec.capacity() >= 11);
+    /// ```
+    pub fn try_reserve_exact(
+        &mut self,
+        additional: usize,
+    ) -> Result<(), CollectionAllocErr> {
+        self.buf.try_reserve_exact(self.len, additional)
+    }
+
     /// Shrinks the capacity of the vector as much as possible.
     ///
     /// It will drop down as close as possible to the length but the allocator
diff --git a/src/lib.rs b/src/lib.rs
index 4ea60ea..fb72742 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -307,7 +307,19 @@ impl Bump {
         Self::with_capacity(0)
     }
 
-    /// Construct a new arena with the specified capacity to bump allocate into.
+    /// Attempt to construct a new arena to bump allocate into.
+    ///
+    /// ## Example
+    ///
+    /// ```
+    /// let bump = bumpalo::Bump::try_new();
+    /// # let _ = bump.unwrap();
+    /// ```
+    pub fn try_new() -> Result<Bump, alloc::AllocErr> {
+        Bump::try_with_capacity(0)
+    }
+
+    /// Construct a new arena with the specified byte capacity to bump allocate into.
     ///
     /// ## Example
     ///
@@ -316,14 +328,27 @@ impl Bump {
     /// # let _ = bump;
     /// ```
     pub fn with_capacity(capacity: usize) -> Bump {
+        Bump::try_with_capacity(capacity).unwrap_or_else(|_| oom())
+    }
+
+    /// Attempt to construct a new arena with the specified byte capacity to bump allocate into.
+    ///
+    /// ## Example
+    ///
+    /// ```
+    /// let bump = bumpalo::Bump::try_with_capacity(100);
+    /// # let _ = bump.unwrap();
+    /// ```
+    pub fn try_with_capacity(capacity: usize) -> Result<Self, alloc::AllocErr> {
         let chunk_footer = Self::new_chunk(
             None,
             Some(unsafe { layout_from_size_align(capacity, 1) }),
             None,
-        );
-        Bump {
+        ).ok_or(alloc::AllocErr {})?;
+
+        Ok(Bump {
             current_chunk_footer: Cell::new(chunk_footer),
-        }
+        })
     }
 
     /// Allocate a new chunk and return its initialized footer.
@@ -335,16 +360,14 @@ impl Bump {
         old_size_with_footer: Option<usize>,
         requested_layout: Option<Layout>,
         prev: Option<NonNull<ChunkFooter>>,
-    ) -> NonNull<ChunkFooter> {
+    ) -> Option<NonNull<ChunkFooter>> {
         unsafe {
             // As a sane default, we want our new allocation to be about twice as
             // big as the previous allocation
             let mut new_size_without_footer =
                 if let Some(old_size_with_footer) = old_size_with_footer {
                     let old_size_without_footer = old_size_with_footer - FOOTER_SIZE;
-                    old_size_without_footer
-                        .checked_mul(2)
-                        .unwrap_or_else(|| oom())
+                    old_size_without_footer.checked_mul(2)?
                 } else {
                     DEFAULT_CHUNK_SIZE_WITHOUT_FOOTER
                 };
@@ -372,8 +395,7 @@ impl Bump {
                 new_size_without_footer =
                     (new_size_without_footer + OVERHEAD).next_power_of_two() - OVERHEAD;
             } else {
-                new_size_without_footer = round_up_to(new_size_without_footer + OVERHEAD, 0x1000)
-                    .unwrap_or_else(|| oom())
+                new_size_without_footer = round_up_to(new_size_without_footer + OVERHEAD, 0x1000)?
                     - OVERHEAD;
             }
 
@@ -387,7 +409,7 @@ impl Bump {
             debug_assert!(size >= old_size_with_footer.unwrap_or(0) * 2);
 
             let data = alloc(layout);
-            let data = NonNull::new(data).unwrap_or_else(|| oom());
+            let data = NonNull::new(data)?;
 
             // The `ChunkFooter` is at the end of the chunk.
             let footer_ptr = data.as_ptr() as usize + new_size_without_footer;
@@ -409,7 +431,7 @@ impl Bump {
                 },
             );
 
-            NonNull::new_unchecked(footer_ptr)
+            Some(NonNull::new_unchecked(footer_ptr))
         }
     }
 
@@ -807,10 +829,21 @@ impl Bump {
     /// [`std::ptr::write`](https://doc.rust-lang.org/stable/std/ptr/fn.write.html).
     #[inline(always)]
     pub fn alloc_layout(&self, layout: Layout) -> NonNull<u8> {
+        self.try_alloc_layout(layout).unwrap_or_else(|_| oom())
+    }
+
+    /// Attempts to allocate space for an object with the given `Layout` or else returns
+    /// an `Err`.
+    ///
+    /// The returned pointer points at uninitialized memory, and should be
+    /// initialized with
+    /// [`std::ptr::write`](https://doc.rust-lang.org/stable/std/ptr/fn.write.html).
+    #[inline(always)]
+    pub fn try_alloc_layout(&self, layout: Layout) -> Result<NonNull<u8>, alloc::AllocErr> {
         if let Some(p) = self.try_alloc_layout_fast(layout) {
-            p
+            Ok(p)
         } else {
-            self.alloc_layout_slow(layout)
+            self.alloc_layout_slow(layout).ok_or(alloc::AllocErr {})
         }
     }
 
@@ -841,10 +874,18 @@ impl Bump {
         }
     }
 
-    // Slow path allocation for when we need to allocate a new chunk from the
-    // parent bump set because there isn't enough room in our current chunk.
+    /// Gets the remaining capacity in the current chunk.
+    pub fn chunk_capacity(&self) -> usize {
+        let current_footer = self.current_chunk_footer.get();
+        let current_footer = unsafe { current_footer.as_ref() };
+
+        current_footer as *const _ as usize - current_footer.data.as_ptr() as usize
+    }
+
+    /// Slow path allocation for when we need to allocate a new chunk from the
+    /// parent bump set because there isn't enough room in our current chunk.
     #[inline(never)]
-    fn alloc_layout_slow(&self, layout: Layout) -> NonNull<u8> {
+    fn alloc_layout_slow(&self, layout: Layout) -> Option<NonNull<u8>> {
         unsafe {
             let size = layout.size();
 
@@ -855,7 +896,7 @@ impl Bump {
                 Some(current_layout.size()),
                 Some(layout),
                 Some(current_footer),
-            );
+            )?;
             debug_assert_eq!(
                 new_footer.as_ref().data.as_ptr() as usize % layout.align(),
                 0
@@ -882,7 +923,7 @@ impl Bump {
             new_footer.ptr.set(ptr);
 
             // Return a pointer to the freshly allocated region in this chunk.
-            ptr
+            Some(ptr)
         }
     }
 
@@ -1066,7 +1107,7 @@ fn oom() -> ! {
 unsafe impl<'a> alloc::Alloc for &'a Bump {
     #[inline(always)]
     unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, alloc::AllocErr> {
-        Ok(self.alloc_layout(layout))
+        self.try_alloc_layout(layout)
     }
 
     #[inline]
@@ -1129,7 +1170,7 @@ unsafe impl<'a> alloc::Alloc for &'a Bump {
 
         // Fallback: do a fresh allocation and copy the existing data into it.
         let new_layout = layout_from_size_align(new_size, layout.align());
-        let new_ptr = self.alloc_layout(new_layout);
+        let new_ptr = self.try_alloc_layout(new_layout)?;
         ptr::copy_nonoverlapping(ptr.as_ptr(), new_ptr.as_ptr(), old_size);
         Ok(new_ptr)
     }


Test patch:
diff --git a/tests/try_alloc.rs b/tests/try_alloc.rs
new file mode 100644
index 0000000..14451fe
--- /dev/null
+++ b/tests/try_alloc.rs
@@ -0,0 +1,159 @@
+use bumpalo::Bump;
+#[cfg(feature = "collections")]
+use bumpalo::collections::Vec;
+use rand::Rng;
+
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::error::Error;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+struct Allocator(AtomicBool);
+
+impl Allocator {
+    fn is_returning_null(&self) -> bool {
+        self.0.load(Ordering::SeqCst)
+    }
+
+    fn toggle_state(&self) {
+        let current = self.0.load(Ordering::SeqCst);
+
+        self.0.store(!current, Ordering::SeqCst)
+    }
+
+    fn scoped_success<F>(&self, callback: F)
+    where
+        F: FnOnce(),
+    {
+        let returning_null = self.is_returning_null();
+
+        if returning_null {
+            self.toggle_state()
+        }
+
+        callback();
+
+        if returning_null {
+            self.toggle_state()
+        }
+    }
+}
+
+unsafe impl GlobalAlloc for Allocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        if self.is_returning_null() {
+            core::ptr::null_mut()
+        } else {
+            SYSTEM_ALLOCATOR.alloc(layout)
+        }
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        if !self.is_returning_null() {
+            SYSTEM_ALLOCATOR.dealloc(ptr, layout);
+        }
+    }
+
+    unsafe fn realloc(
+        &self,
+        ptr: *mut u8,
+        layout: Layout,
+        new_size: usize
+    ) -> *mut u8 {
+        if self.is_returning_null() {
+            core::ptr::null_mut()
+        } else {
+            SYSTEM_ALLOCATOR.realloc(ptr, layout, new_size)
+        }
+    }
+}
+
+#[global_allocator]
+static GLOBAL_ALLOCATOR: Allocator = Allocator(AtomicBool::new(false));
+static SYSTEM_ALLOCATOR: System = System;
+
+macro_rules! toggling_assert {
+    ($cond:expr $(, $args:tt)*) => {
+        let result = $cond;
+
+        // assert! may allocate on failure, so we must re-enable allocations
+        // prior to asserting
+        GLOBAL_ALLOCATOR.scoped_success(|| assert!(result $(, $args)*))
+    };
+}
+
+#[test]
+fn test_try_alloc_layout() -> Result<(), Box<dyn Error>> {
+    const NUM_TESTS: usize = 32;
+
+    if !GLOBAL_ALLOCATOR.is_returning_null() {
+        GLOBAL_ALLOCATOR.toggle_state();
+    }
+
+    toggling_assert!(Bump::try_new().is_err());
+
+    GLOBAL_ALLOCATOR.toggle_state();
+
+    let bump = Bump::try_new().unwrap();
+
+    // Bump preallocates space in the initial chunk, so we need to
+    // use up this block prior to the actual test
+    let layout = Layout::from_size_align(bump.chunk_capacity(), 1)?;
+
+    toggling_assert!(bump.try_alloc_layout(layout).is_ok());
+
+    let mut rng = rand::thread_rng();
+
+    for _ in 0..NUM_TESTS {
+        let layout = Layout::from_size_align(bump.chunk_capacity() + 1, 1)?;
+
+        if rng.gen() {
+            GLOBAL_ALLOCATOR.toggle_state();
+        } else if GLOBAL_ALLOCATOR.is_returning_null() {
+            toggling_assert!(bump.try_alloc_layout(layout).is_err());
+        } else {
+            toggling_assert!(bump.try_alloc_layout(layout).is_ok());
+        }
+    }
+
+    #[cfg(feature = "collections")]
+    {
+        if GLOBAL_ALLOCATOR.is_returning_null() {
+            GLOBAL_ALLOCATOR.toggle_state();
+        }
+
+        let bump = Bump::try_new().unwrap();
+
+        if !GLOBAL_ALLOCATOR.is_returning_null() {
+            GLOBAL_ALLOCATOR.toggle_state();
+        }
+
+        let mut vec = Vec::<u8>::new_in(&bump);
+
+        let chunk_cap = bump.chunk_capacity();
+
+        // Will always succeed since this size gets pre-allocated in Bump::try_new()
+        toggling_assert!(vec.try_reserve(chunk_cap).is_ok());
+        toggling_assert!(vec.try_reserve_exact(chunk_cap).is_ok());
+        // Fails to allocate futher since allocator returns null
+        toggling_assert!(vec.try_reserve(chunk_cap + 1).is_err());
+        toggling_assert!(vec.try_reserve_exact(chunk_cap + 1).is_err());
+
+        GLOBAL_ALLOCATOR.toggle_state();
+
+        let mut vec = Vec::<u8>::new_in(&bump);
+
+        // Will always succeed since this size gets pre-allocated in Bump::try_new()
+        toggling_assert!(vec.try_reserve(chunk_cap).is_ok());
+        toggling_assert!(vec.try_reserve_exact(chunk_cap).is_ok());
+        // Succeeds to allocate further
+        toggling_assert!(vec.try_reserve(chunk_cap + 1).is_ok());
+        toggling_assert!(vec.try_reserve_exact(chunk_cap + 1).is_ok());
+    }
+
+    // Reset the allocator for the test harness to use
+    if GLOBAL_ALLOCATOR.is_returning_null() {
+        GLOBAL_ALLOCATOR.toggle_state();
+    }
+
+    Ok(())
+}