diff --git a/src/lib.rs b/src/lib.rs index 22bddd3..2a4d678 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1726,12 +1726,7 @@ impl SmallVec { where F: FnMut() -> T { let old_len = self.len(); if old_len < new_len { - let mut f = f; - let additional = new_len - old_len; - self.reserve(additional); - for _ in 0..additional { - self.push(f()); - } + self.extend(core::iter::repeat_with(f).take(new_len - old_len)); } else if old_len > new_len { self.truncate(new_len); } @@ -2449,11 +2444,55 @@ impl SmallVec { fn extend_fallback(&mut self, iter: I) where I: IntoIterator { - let iter = iter.into_iter(); - let (size, _) = iter.size_hint(); - self.reserve(size); - for x in iter { - self.push(x); + struct SetLenOnDrop<'a, T, const N: usize> { + vec: &'a mut SmallVec, + len: usize + } + impl Drop for SetLenOnDrop<'_, T, N> { + #[inline(always)] + fn drop(&mut self) { + // SAFETY: restores `len` only len of initialized items. + unsafe { self.vec.set_len(self.len) }; + } + } + + let mut iter = iter.into_iter(); + let (lower, _) = iter.size_hint(); + self.reserve(lower); + + // gaurd will auto drop with return statements inside loop + let mut guard = SetLenOnDrop { + len: self.len(), + vec: self + }; + + loop { + let capacity = guard.vec.capacity(); + // ptr stays valid until we reserve(memory relocation) + let ptr = guard.vec.as_mut_ptr(); + + // fast path for when size_hint provides exact length + while guard.len < capacity { + let Some(value) = iter.next() else { return }; + // SAFETY: `guard.len < capacity` and slot is uninitialized. + unsafe { ptr.add(guard.len).write(value) }; + guard.len += 1; + } + // exit happens mostly in the while body `else { return }`, + // if we are here iterator is not exact size. + let Some(value) = iter.next() else { return }; + let (lower, _) = iter.size_hint(); + // restore vector length before requesting reserve + // SAFETY: elements up to `guard.len` are initialized. + unsafe { guard.vec.set_len(guard.len) }; + // UNSAFE TO USE `ptr` AFTER THIS POINT. + // reserve pessimistic lower+1. + guard.vec.reserve(lower.saturating_add(1)); + // write the currently read value + // SAFETY: `reserve` made room for us. + unsafe { guard.vec.as_mut_ptr().add(guard.len).write(value) }; + guard.len += 1; + // retry in the fast path } } @@ -2503,9 +2542,7 @@ impl SmallVec { where I: Iterator { let (size, _) = iter.size_hint(); let mut v = Self::with_capacity(size); - for x in iter { - v.push(x); - } + v.extend_fallback(iter); v }