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 Vec { pub fn truncate(&mut self, new_len: usize) { self.inner.truncate(new_len); } + + #[inline] + pub fn resize_with 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 Vec { @@ -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]); + } }