From fe272c28a81e5ab31e06215c1a627845d727c627 Mon Sep 17 00:00:00 2001
From: Adam Cimarosti <cimarosti@gmail.com>
Date: Fri, 27 Sep 2024 16:47:36 +0100
Subject: [PATCH] resize_with

---
 src/vec.rs | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/src/vec.rs b/src/vec.rs
index 663ffd5..591c80b 100644
--- a/src/vec.rs
+++ b/src/vec.rs
@@ -118,6 +118,26 @@ impl<T, A: Allocator> Vec<T, A> {
     pub fn truncate(&mut self, new_len: usize) {
         self.inner.truncate(new_len);
     }
+
+    #[inline]
+    pub fn resize_with<F: FnMut() -> T>(
+        &mut self,
+        new_len: usize,
+        mut f: F,
+    ) -> Result<(), TryReserveError> {
+        let len = self.len();
+        if new_len > len {
+            self.reserve(new_len - len)?;
+            for _ in 0..new_len - len {
+                unsafe {
+                    self.unsafe_push(f());
+                }
+            }
+        } else {
+            self.truncate(new_len);
+        }
+        Ok(())
+    }
 }
 
 impl<T: Clone, A: Allocator> Vec<T, A> {
@@ -546,4 +566,16 @@ mod tests {
         vec.resize(2, 3).unwrap();
         assert_eq!(vec.inner.as_slice(), &[1, 1]);
     }
+
+    #[test]
+    fn test_resize_with() {
+        let wma = WatermarkAllocator::new(64);
+        let mut vec = Vec::new_in(wma);
+        vec.resize_with(3, || 1).unwrap();
+        assert_eq!(vec.inner.as_slice(), &[1, 1, 1]);
+        vec.resize_with(5, || 2).unwrap();
+        assert_eq!(vec.inner.as_slice(), &[1, 1, 1, 2, 2]);
+        vec.resize_with(2, || 3).unwrap();
+        assert_eq!(vec.inner.as_slice(), &[1, 1]);
+    }
 }