Add fallible try_collect to collections (#12505)

* Add fallible `try_collect` to collections

This adds a new `TryCollect` extension trait to `wasmtime-core` which
enables `iter.try_collect()` to fallibly collect into a `Vec<T>` or
`Box<[T]>` while handling OOM. This is intended to serve as a
replacement for the `new_boxed_slice_*` helpers in appropriate
situations since it more closely matches what Wasmtime likely already
does today where we throw `.collect()` on iterators.

* Extend oom tests a bit more

* Use a TryExtend trait

* Fix oom tests

Author: alexcrichton

cranelift/bitset/src/compound.rs

@@ -3,6 +3,7 @@
 use crate::scalar::{self, ScalarBitSet, ScalarBitSetStorage};
 use alloc::boxed::Box;
 use core::{cmp, iter, mem};
+use wasmtime_core::alloc::{TryExtend, Vec};
 use wasmtime_core::error::OutOfMemory;
 
 /// A large bit set backed by dynamically-sized storage.
@@ -294,20 +295,11 @@ impl<T: ScalarBitSetStorage> CompoundBitSet<T> {
         // Don't make ridiculously small allocations.
         let to_grow = cmp::max(to_grow, 4);
 
-        let new_len = self.elems.len() + to_grow;
-        match wasmtime_core::alloc::new_boxed_slice_from_iter_with_len(
-            new_len,
-            self.elems
-                .iter()
-                .copied()
-                .chain(iter::repeat(ScalarBitSet::new())),
-        ) {
-            Ok(new_elems) => {
-                self.elems = new_elems;
-                Ok(())
-            }
-            Err(e) => Err(e.unwrap_oom()),
-        }
+        let mut new_elems = Vec::from(mem::take(&mut self.elems));
+        new_elems.reserve_exact(to_grow)?;
+        new_elems.try_extend(iter::repeat(ScalarBitSet::new()).take(to_grow))?;
+        self.elems = new_elems.into_boxed_slice()?;
+        Ok(())
     }
 
     /// Insert `i` into this bitset.

crates/core/src/alloc.rs

@@ -2,6 +2,7 @@
 
 mod arc;
 mod boxed;
+mod try_collect;
 mod try_new;
 mod vec;
 
@@ -10,6 +11,7 @@ pub use boxed::{
     new_boxed_slice_from_fallible_iter, new_boxed_slice_from_iter,
     new_boxed_slice_from_iter_with_len, new_uninit_boxed_slice,
 };
+pub use try_collect::{TryCollect, TryExtend, TryFromIterator};
 pub use try_new::{TryNew, try_new};
 pub use vec::Vec;
 
@@ -35,6 +37,32 @@ unsafe fn try_alloc(layout: Layout) -> Result<NonNull<u8>, OutOfMemory> {
     }
 }
 
+/// Tries to reallocate a block of memory, returning `OutOfMemory` on failure.
+///
+/// Analogue of [`alloc::alloc::realloc`].
+///
+/// # Safety
+///
+/// Same as `alloc::alloc::realloc`: `ptr` must be allocated with `layout`,
+/// `layout` must be nonzero in size, and `new_size` must be nonzero and valid.
+#[inline]
+unsafe fn try_realloc(
+    ptr: *mut u8,
+    layout: Layout,
+    new_size: usize,
+) -> Result<NonNull<u8>, OutOfMemory> {
+    // Safety: same as our safety conditions.
+    debug_assert!(layout.size() > 0);
+    debug_assert!(new_size > 0);
+    let ptr = unsafe { std_alloc::alloc::realloc(ptr, layout, new_size) };
+
+    if let Some(ptr) = NonNull::new(ptr) {
+        Ok(ptr)
+    } else {
+        Err(OutOfMemory::new(new_size))
+    }
+}
+
 /// An extension trait for ignoring `OutOfMemory` errors.
 ///
 /// Use this to unwrap a `Result<T, OutOfMemory>` into its inner `T` or

crates/core/src/alloc/try_collect.rs

@@ -0,0 +1,179 @@
+use crate::alloc::Vec;
+use crate::error::OutOfMemory;
+use std_alloc::boxed::Box;
+
+/// Extension trait for an `Iterator` to fallibly collect into a container.
+pub trait TryCollect: Iterator {
+    /// Attempts to collect the iterator `self` into `B`.
+    ///
+    /// Same as [`Iterator::collect`] except returns OOM instead of aborting.
+    fn try_collect<B, E>(self) -> Result<B, E>
+    where
+        B: TryFromIterator<Self::Item, E>,
+        Self: Sized,
+    {
+        B::try_from_iter(self)
+    }
+}
+
+impl<I: Iterator> TryCollect for I {}
+
+/// Analogue of [`FromIterator`] in the standard library, but used with
+/// [`TryCollect::try_collect`] instead.
+pub trait TryFromIterator<T, E>: Sized {
+    /// Creates an intance of this collection from the `iter` provided.
+    ///
+    /// Does not abort on OOM but insteads return an error.
+    fn try_from_iter<I>(iter: I) -> Result<Self, E>
+    where
+        I: Iterator<Item = T>;
+}
+
+impl<T> TryFromIterator<T, OutOfMemory> for Vec<T> {
+    fn try_from_iter<I>(iter: I) -> Result<Self, OutOfMemory>
+    where
+        I: Iterator<Item = T>,
+    {
+        let mut result = Vec::with_capacity(iter.size_hint().0)?;
+        for item in iter {
+            result.push(item)?;
+        }
+        Ok(result)
+    }
+}
+
+impl<T> TryFromIterator<T, OutOfMemory> for Box<[T]> {
+    fn try_from_iter<I>(iter: I) -> Result<Self, OutOfMemory>
+    where
+        I: Iterator<Item = T>,
+    {
+        let vec = Vec::try_from_iter(iter)?;
+        vec.into_boxed_slice()
+    }
+}
+
+impl<T, E> TryFromIterator<Result<T, E>, E> for Vec<T>
+where
+    E: From<OutOfMemory>,
+{
+    fn try_from_iter<I>(iter: I) -> Result<Self, E>
+    where
+        I: Iterator<Item = Result<T, E>>,
+    {
+        let mut result = Vec::with_capacity(iter.size_hint().0)?;
+        for item in iter {
+            result.push(item?)?;
+        }
+        Ok(result)
+    }
+}
+
+impl<T, E> TryFromIterator<Result<T, E>, E> for Box<[T]>
+where
+    E: From<OutOfMemory>,
+{
+    fn try_from_iter<I>(iter: I) -> Result<Self, E>
+    where
+        I: Iterator<Item = Result<T, E>>,
+    {
+        let vec = iter.try_collect::<Vec<_>, E>()?;
+        Ok(vec.into_boxed_slice()?)
+    }
+}
+
+/// Analogue of [`Extend`] except handles OOM conditions.
+pub trait TryExtend<T> {
+    /// Extends `self` with the items from `iter`.
+    ///
+    /// Returns an error if allocation fails while adding items to `self`. If an
+    /// OOM happens then some items from `iter` may have been added to `self`
+    /// already. On OOM no further items from the iterator will be consumed.
+    fn try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory>
+    where
+        I: IntoIterator<Item = T>;
+}
+
+impl<T> TryExtend<T> for Vec<T> {
+    fn try_extend<I>(&mut self, iter: I) -> Result<(), OutOfMemory>
+    where
+        I: IntoIterator<Item = T>,
+    {
+        let iter = iter.into_iter();
+        self.reserve(iter.size_hint().0)?;
+        for item in iter {
+            self.push(item)?;
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{Box, TryCollect, TryExtend, Vec};
+    use crate::error::{OutOfMemory, Result};
+
+    #[test]
+    fn test_vec_collect() -> Result<(), OutOfMemory> {
+        let v: Vec<i32> = (0..10).try_collect()?;
+        assert_eq!(&*v, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
+        Ok(())
+    }
+
+    #[test]
+    fn test_box_collect() -> Result<(), OutOfMemory> {
+        let v: Box<[i32]> = (0..10).try_collect()?;
+        assert_eq!(&*v, &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
+        Ok(())
+    }
+
+    #[test]
+    fn test_vec_result_collect() -> Result<()> {
+        let v: Result<Vec<i32>> = [].into_iter().try_collect();
+        assert!(v?.is_empty());
+
+        let v: Result<Vec<i32>> = [Ok(1), Ok(2)].into_iter().try_collect();
+        assert_eq!(&*v?, &[1, 2]);
+
+        let v: Result<Vec<i32>> = [Ok(1), Err(crate::format_err!("hi"))]
+            .into_iter()
+            .try_collect();
+        assert!(v.is_err());
+
+        let v: Result<Vec<i32>> = [Err(crate::format_err!("hi")), Ok(1)]
+            .into_iter()
+            .try_collect();
+        assert!(v.is_err());
+        Ok(())
+    }
+
+    #[test]
+    fn test_box_result_collect() -> Result<()> {
+        let v: Result<Box<[i32]>> = [].into_iter().try_collect();
+        assert!(v?.is_empty());
+
+        let v: Result<Box<[i32]>> = [Ok(1), Ok(2)].into_iter().try_collect();
+        assert_eq!(&*v?, &[1, 2]);
+
+        let v: Result<Box<[i32]>> = [Ok(1), Err(crate::format_err!("hi"))]
+            .into_iter()
+            .try_collect();
+        assert!(v.is_err());
+
+        let v: Result<Box<[i32]>> = [Err(crate::format_err!("hi")), Ok(1)]
+            .into_iter()
+            .try_collect();
+        assert!(v.is_err());
+        Ok(())
+    }
+
+    #[test]
+    fn test_try_extend() -> Result<(), OutOfMemory> {
+        let mut vec = Vec::new();
+        vec.try_extend([1, 2, 3].iter().cloned())?;
+        assert_eq!(&*vec, &[1, 2, 3]);
+
+        vec.try_extend([])?;
+        assert_eq!(&*vec, &[1, 2, 3]);
+        Ok(())
+    }
+}