Introduce a Vec<T> collection that forces fallible allocation (#12440)

* Introduce a `Vec<T>` collection that forces fallible allocation

Part of #12069

* review feedback

Author: fitzgen

crates/environ/src/collections.rs

@@ -2,7 +2,9 @@
 
 mod entity_set;
 mod primary_map;
+mod vec;
 
 pub use entity_set::EntitySet;
 pub use primary_map::PrimaryMap;
+pub use vec::Vec;
 pub use wasmtime_core::alloc::{TryNew, try_new};

crates/environ/src/collections/vec.rs

@@ -0,0 +1,133 @@
+use crate::error::OutOfMemory;
+use alloc::vec::Vec as StdVec;
+use core::{
+    fmt,
+    ops::{Deref, DerefMut, Index, IndexMut},
+};
+
+/// Like `std::vec::Vec` but all methods that allocate force handling allocation
+/// failure.
+#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct Vec<T> {
+    inner: StdVec<T>,
+}
+
+impl<T> Default for Vec<T> {
+    fn default() -> Self {
+        Self {
+            inner: Default::default(),
+        }
+    }
+}
+
+impl<T: fmt::Debug> fmt::Debug for Vec<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Debug::fmt(&self.inner, f)
+    }
+}
+
+impl<T> Vec<T> {
+    /// Same as [`std::vec::Vec::new`].
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    /// Same as [`std::vec::Vec::with_capacity`] but returns an error on
+    /// allocation failure.
+    pub fn with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
+        let mut v = Self::new();
+        v.reserve(capacity)?;
+        Ok(v)
+    }
+
+    /// Same as [`std::vec::Vec::reserve`] but returns an error on allocation
+    /// failure.
+    pub fn reserve(&mut self, additional: usize) -> Result<(), OutOfMemory> {
+        self.inner.try_reserve(additional).map_err(|_| {
+            OutOfMemory::new(
+                self.len()
+                    .saturating_add(additional)
+                    .saturating_mul(core::mem::size_of::<T>()),
+            )
+        })
+    }
+
+    /// Same as [`std::vec::Vec::reserve_exact`] but returns an error on allocation
+    /// failure.
+    pub fn reserve_exact(&mut self, additional: usize) -> Result<(), OutOfMemory> {
+        self.inner
+            .try_reserve_exact(additional)
+            .map_err(|_| OutOfMemory::new(self.len().saturating_add(additional)))
+    }
+
+    /// Same as [`std::vec::Vec::len`].
+    pub fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    /// Same as [`std::vec::Vec::capacity`].
+    pub fn capacity(&self) -> usize {
+        self.inner.capacity()
+    }
+
+    /// Same as [`std::vec::Vec::is_empty`].
+    pub fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
+
+    /// Same as [`std::vec::Vec::push`] but returns an error on allocation
+    /// failure.
+    pub fn push(&mut self, value: T) -> Result<(), OutOfMemory> {
+        self.reserve(1)?;
+        self.inner.push(value);
+        Ok(())
+    }
+}
+
+impl<T> Deref for Vec<T> {
+    type Target = [T];
+
+    fn deref(&self) -> &Self::Target {
+        &self.inner
+    }
+}
+
+impl<T> DerefMut for Vec<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.inner
+    }
+}
+
+impl<T> Index<usize> for Vec<T> {
+    type Output = T;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.inner[index]
+    }
+}
+
+impl<T> IndexMut<usize> for Vec<T> {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        &mut self.inner[index]
+    }
+}
+
+impl<'a, T> IntoIterator for &'a Vec<T> {
+    type Item = &'a T;
+
+    type IntoIter = core::slice::Iter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        (**self).iter()
+    }
+}
+
+impl<'a, T> IntoIterator for &'a mut Vec<T> {
+    type Item = &'a mut T;
+
+    type IntoIter = core::slice::IterMut<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        (**self).iter_mut()
+    }
+}

crates/fuzzing/tests/oom.rs

@@ -106,6 +106,41 @@ fn primary_map_try_reserve_exact() -> Result<()> {
     })
 }
 
+#[test]
+fn vec_with_capacity() -> Result<()> {
+    OomTest::new().test(|| {
+        let _v = wasmtime_environ::collections::Vec::<usize>::with_capacity(100)?;
+        Ok(())
+    })
+}
+
+#[test]
+fn vec_reserve() -> Result<()> {
+    OomTest::new().test(|| {
+        let mut v = wasmtime_environ::collections::Vec::<usize>::new();
+        v.reserve(10)?;
+        Ok(())
+    })
+}
+
+#[test]
+fn vec_reserve_exact() -> Result<()> {
+    OomTest::new().test(|| {
+        let mut v = wasmtime_environ::collections::Vec::<usize>::new();
+        v.reserve_exact(3)?;
+        Ok(())
+    })
+}
+
+#[test]
+fn vec_push() -> Result<()> {
+    OomTest::new().test(|| {
+        let mut v = wasmtime_environ::collections::Vec::new();
+        v.push(42)?;
+        Ok(())
+    })
+}
+
 #[test]
 fn config_new() -> Result<()> {
     OomTest::new().test(|| {