Add fallible slice allocation functions (#262)

This change adds slice allocation functions  that can fail, to match the
scalar allocation functions that can fail.

Author: 414owen

src/lib.rs

@@ -1351,6 +1351,58 @@ impl Bump {
         }
     }
 
+    /// Allocates a new slice of size `len` into this `Bump` and returns an
+    /// exclusive reference to the copy, failing if the closure return an Err.
+    ///
+    /// The elements of the slice are initialized using the supplied closure.
+    /// The closure argument is the position in the slice.
+    ///
+    /// ## Panics
+    ///
+    /// Panics if reserving space for the slice fails.
+    ///
+    /// ## Example
+    ///
+    /// ```
+    /// let bump = bumpalo::Bump::new();
+    /// let x: Result<&mut [usize], ()> = bump.alloc_slice_try_fill_with(5, |i| Ok(5 * i));
+    /// assert_eq!(x, Ok(bump.alloc_slice_copy(&[0, 5, 10, 15, 20])));
+    /// ```
+    ///
+    /// ```
+    /// let bump = bumpalo::Bump::new();
+    /// let x: Result<&mut [usize], ()> = bump.alloc_slice_try_fill_with(
+    ///    5,
+    ///    |n| if n == 2 { Err(()) } else { Ok(n) }
+    /// );
+    /// assert_eq!(x, Err(()));
+    /// ```
+    #[inline(always)]
+    pub fn alloc_slice_try_fill_with<T, F, E>(&self, len: usize, mut f: F) -> Result<&mut [T], E>
+    where
+        F: FnMut(usize) -> Result<T, E>,
+    {
+        let layout = Layout::array::<T>(len).unwrap_or_else(|_| oom());
+        let base_ptr = self.alloc_layout(layout);
+        let dst = base_ptr.cast::<T>();
+
+        unsafe {
+            for i in 0..len {
+                match f(i) {
+                    Ok(el) => ptr::write(dst.as_ptr().add(i), el),
+                    Err(e) => {
+                        self.dealloc(base_ptr, layout);
+                        return Err(e);
+                    }
+                }
+            }
+
+            let result = slice::from_raw_parts_mut(dst.as_ptr(), len);
+            debug_assert_eq!(Layout::for_value(result), layout);
+            Ok(result)
+        }
+    }
+
     /// Allocates a new slice of size `len` into this `Bump` and returns an
     /// exclusive reference to the copy.
     ///
@@ -1487,6 +1539,45 @@ impl Bump {
         })
     }
 
+    /// Allocates a new slice of size `len` slice into this `Bump` and return an
+    /// exclusive reference to the copy, failing if the iterator returns an Err.
+    ///
+    /// The elements are initialized using the supplied iterator.
+    ///
+    /// ## Panics
+    ///
+    /// Panics if reserving space for the slice fails, or if the supplied
+    /// iterator returns fewer elements than it promised.
+    ///
+    /// ## Examples
+    ///
+    /// ```
+    /// let bump = bumpalo::Bump::new();
+    /// let x: Result<&mut [i32], ()> = bump.alloc_slice_try_fill_iter(
+    ///    [2, 3, 5].iter().cloned().map(|i| Ok(i * i))
+    /// );
+    /// assert_eq!(x, Ok(bump.alloc_slice_copy(&[4, 9, 25])));
+    /// ```
+    ///
+    /// ```
+    /// let bump = bumpalo::Bump::new();
+    /// let x: Result<&mut [i32], ()> = bump.alloc_slice_try_fill_iter(
+    ///    [Ok(2), Err(()), Ok(5)].iter().cloned()
+    /// );
+    /// assert_eq!(x, Err(()));
+    /// ```
+    #[inline(always)]
+    pub fn alloc_slice_try_fill_iter<T, I, E>(&self, iter: I) -> Result<&mut [T], E>
+    where
+        I: IntoIterator<Item = Result<T, E>>,
+        I::IntoIter: ExactSizeIterator,
+    {
+        let mut iter = iter.into_iter();
+        self.alloc_slice_try_fill_with(iter.len(), |_| {
+            iter.next().expect("Iterator supplied too few elements")
+        })
+    }
+
     /// Allocates a new slice of size `iter.len()` slice into this `Bump` and return an
     /// exclusive reference to the copy. Does not panic on failure.
     ///

tests/all/alloc_fill.rs

@@ -1,6 +1,7 @@
 use bumpalo::Bump;
 use std::alloc::Layout;
 use std::cmp;
+use std::iter::repeat;
 use std::mem;
 
 #[test]
@@ -29,6 +30,35 @@ fn alloc_slice_fill_zero() {
     assert_eq!(ptr2 as *mut _ as usize, ptr3.as_ptr() as usize + 1);
 }
 
+#[test]
+fn alloc_slice_try_fill_with_succeeds() {
+    let b = Bump::new();
+    let res: Result<&mut [usize], ()> = b.alloc_slice_try_fill_with(100, |n| Ok(n));
+    assert_eq!(res.map(|arr| arr[50]), Ok(50));
+}
+
+#[test]
+fn alloc_slice_try_fill_with_fails() {
+    let b = Bump::new();
+    let res: Result<&mut [u16], ()> =
+        b.alloc_slice_try_fill_with(1000, |n| if n == 100 { Err(()) } else { Ok(42) });
+    assert_eq!(res, Err(()));
+}
+
+#[test]
+fn alloc_slice_try_fill_iter_succeeds() {
+    let b = Bump::new();
+    let res: Result<&mut [u16], ()> = b.alloc_slice_try_fill_iter(repeat(42).take(10).map(Ok));
+    assert_eq!(res.map(|arr| arr[5]), Ok(42));
+}
+
+#[test]
+fn alloc_slice_try_fill_iter_fails() {
+    let b = Bump::new();
+    let res: Result<&mut [u16], ()> = b.alloc_slice_try_fill_iter(repeat(()).take(10).map(Err));
+    assert_eq!(res, Err(()));
+}
+
 #[test]
 #[should_panic(expected = "out of memory")]
 fn alloc_slice_overflow() {

tests/all/alloc_try_with.rs

@@ -4,6 +4,8 @@
 // We only run them when debug_assertions are not set, as we expect them to fail outside release
 // mode.
 
+use std::iter::repeat;
+
 use bumpalo::Bump;
 
 #[test]
@@ -116,3 +118,23 @@ fn alloc_try_with_large_enum_err() {
         .alloc_try_with(|| Result::<LargeEnum, _>::Err(()))
         .is_err());
 }
+
+#[test]
+#[cfg_attr(debug_assertions, ignore)]
+fn alloc_slice_try_fill_with_large_length() {
+    let b = Bump::new();
+
+    assert!(b
+        .alloc_slice_try_fill_with(10_000_000, |_| Err::<u8, _>(()))
+        .is_err());
+}
+
+#[test]
+#[cfg_attr(debug_assertions, ignore)]
+fn alloc_slice_try_fill_iter_large_length() {
+    let b = Bump::new();
+
+    assert!(b
+        .alloc_slice_try_fill_iter(repeat(Err::<u8, _>(())).take(10_000_000))
+        .is_err());
+}