Add truncate and resize methods to our OOM-handling Vec (#12583)

Author: fitzgen

crates/core/src/alloc/vec.rs

@@ -1,5 +1,6 @@
 use crate::alloc::{TryClone, try_realloc};
 use crate::error::OutOfMemory;
+use core::cmp::Ordering;
 use core::marker::PhantomData;
 use core::{
     fmt, mem,
@@ -108,6 +109,34 @@ impl<T> Vec<T> {
         self.inner.pop()
     }
 
+    /// Same as [`std::vec::Vec::truncate`].
+    pub fn truncate(&mut self, len: usize) {
+        self.inner.truncate(len);
+    }
+
+    /// Same as [`std::vec::Vec::resize`] but returns an error on allocation
+    /// failure.
+    pub fn resize(&mut self, new_len: usize, value: T) -> Result<(), OutOfMemory>
+    where
+        T: TryClone,
+    {
+        match new_len.cmp(&self.len()) {
+            Ordering::Less => self.truncate(new_len),
+            Ordering::Equal => {}
+            Ordering::Greater => {
+                let delta = new_len - self.len();
+                self.reserve(delta)?;
+                // Minimize `try_clone` calls by always pushing `value` directly
+                // as the last element.
+                for _ in 0..delta - 1 {
+                    self.push(value.try_clone()?)?;
+                }
+                self.push(value)?;
+            }
+        }
+        Ok(())
+    }
+
     /// Same as [`std::vec::Vec::into_raw_parts`].
     pub fn into_raw_parts(mut self) -> (*mut T, usize, usize) {
         // NB: Can't use `Vec::into_raw_parts` until our MSRV is >= 1.93.

crates/fuzzing/tests/oom.rs

@@ -708,6 +708,19 @@ fn vec_shrink_to_fit() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn vec_resize() -> Result<()> {
+    OomTest::new().test(|| {
+        let mut v = Vec::new();
+        v.resize(10, 'a')?; // Grow.
+        v.resize(1, 'b')?; // Truncate.
+        v.resize(1, 'c')?; // Same length.
+        v.resize(3, 'd')?; // Grow again.
+        assert_eq!(&*v, &['a', 'd', 'd']);
+        Ok(())
+    })
+}
+
 #[test]
 fn vec_try_collect() -> Result<()> {
     OomTest::new().test(|| {