Add a vec! macro for our OOM-handling Vec (#12581)

Author: fitzgen

crates/core/src/alloc/vec.rs

@@ -2,6 +2,7 @@ use crate::alloc::{TryClone, try_realloc};
 use crate::error::OutOfMemory;
 use core::cmp::Ordering;
 use core::marker::PhantomData;
+use core::num::NonZeroUsize;
 use core::{
     fmt, mem,
     ops::{Deref, DerefMut, Index, IndexMut},
@@ -11,6 +12,38 @@ use std_alloc::alloc::Layout;
 use std_alloc::boxed::Box;
 use std_alloc::vec::Vec as StdVec;
 
+/// Same as the [`std::vec!`] macro but returns an error on allocation failure.
+#[macro_export]
+macro_rules! vec {
+    ( $( $elem:expr ),* ) => {{
+        let len = $crate::private_len!( $( $elem ),* );
+        $crate::alloc::Vec::with_capacity(len).and_then(|mut v| {
+            $( v.push($elem)?; )*
+            let _ = &mut v;
+            Ok(v)
+        })
+    }};
+
+    ( $elem:expr; $len:expr ) => {{
+        let len: usize = $len;
+        if let Some(len) = ::core::num::NonZeroUsize::new(len) {
+            let elem = $elem;
+            $crate::alloc::Vec::from_elem(elem, len)
+        } else {
+            Ok($crate::alloc::Vec::new())
+        }
+    }};
+
+}
+
+// Only for use by the `vec!` macro.
+#[doc(hidden)]
+#[macro_export]
+macro_rules! private_len {
+    ( ) => { 0 };
+    ( $e:expr $( , $es:expr )* ) => { 1 + $crate::private_len!( $( $es ),* ) };
+}
+
 /// Like `std::vec::Vec` but all methods that allocate force handling allocation
 /// failure.
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -61,6 +94,25 @@ impl<T> Vec<T> {
         Ok(v)
     }
 
+    // For use with the `vec!` macro.
+    #[doc(hidden)]
+    #[inline]
+    pub fn from_elem(elem: T, len: NonZeroUsize) -> Result<Self, OutOfMemory>
+    where
+        T: TryClone,
+    {
+        let mut v = Self::with_capacity(len.get())?;
+
+        // Minimize calls to `TryClone` by always pushing `elem` itself as the
+        // last element.
+        for _ in 0..len.get() - 1 {
+            v.push(elem.try_clone()?)?;
+        }
+        v.push(elem)?;
+
+        Ok(v)
+    }
+
     /// Same as [`std::vec::Vec::reserve`] but returns an error on allocation
     /// failure.
     pub fn reserve(&mut self, additional: usize) -> Result<(), OutOfMemory> {

crates/environ/src/collections.rs

@@ -13,7 +13,10 @@ pub use hash_set::HashSet;
 pub use index_map::IndexMap;
 pub use primary_map::PrimaryMap;
 pub use secondary_map::SecondaryMap;
-pub use wasmtime_core::alloc::{String, TryClone, TryNew, Vec, try_new};
+pub use wasmtime_core::{
+    alloc::{String, TryClone, TryNew, Vec, try_new},
+    vec,
+};
 
 /// Collections which abort on OOM.
 //

crates/fuzzing/tests/oom.rs

@@ -8,7 +8,10 @@ use std::{
 };
 use wasmtime::{error::OutOfMemory, *};
 use wasmtime_core::alloc::TryCollect;
-use wasmtime_environ::{PrimaryMap, SecondaryMap, collections::*};
+use wasmtime_environ::{
+    PrimaryMap, SecondaryMap,
+    collections::{self, *},
+};
 use wasmtime_fuzzing::oom::{OomTest, OomTestAllocator};
 
 #[global_allocator]
@@ -749,6 +752,24 @@ fn vec_extend() -> Result<()> {
     })
 }
 
+#[test]
+fn vec_macro_elems() -> Result<()> {
+    OomTest::new().test(|| {
+        let v = collections::vec![100, 200, 300, 400]?;
+        assert_eq!(&*v, &[100, 200, 300, 400]);
+        Ok(())
+    })
+}
+
+#[test]
+fn vec_macro_elem_len() -> Result<()> {
+    OomTest::new().test(|| {
+        let v = collections::vec![100; 3]?;
+        assert_eq!(&*v, &[100, 100, 100]);
+        Ok(())
+    })
+}
+
 #[test]
 fn index_map_try_clone() -> Result<()> {
     OomTest::new()