Skip to content

Commit

Permalink
Implement spare_capacity_mut (#364)
Browse files Browse the repository at this point in the history
Mirrors the implementation from `Vec::spare_capacity_mut`.

This is useful to write directly into the buffer after `SmallVec::with_capacity`,
removing the overhead of the `SmallVec` model entirely when doing so.
  • Loading branch information
DaniPopes authored Dec 10, 2024
1 parent f6ec5b2 commit 9eacf68
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 1 deletion.
16 changes: 16 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1413,6 +1413,22 @@ impl<T, const N: usize> SmallVec<T, N> {
}
}

/// Returns the remaining spare capacity of the vector as a slice of
/// `MaybeUninit<T>`.
///
/// The returned slice can be used to fill the vector with data (e.g. by
/// reading from a file) before marking the data as initialized using the
/// [`set_len`](Self::set_len) method.
#[inline]
pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
unsafe {
core::slice::from_raw_parts_mut(
self.as_mut_ptr().add(self.len()) as *mut MaybeUninit<T>,
self.capacity() - self.len(),
)
}
}

/// Creates a `SmallVec` directly from the raw components of another `SmallVec`.
///
/// # Safety
Expand Down
31 changes: 30 additions & 1 deletion src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1136,7 +1136,9 @@ fn collect_from_iter() {

impl<I: Iterator> Iterator for IterNoHint<I> {
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> { self.0.next() }
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}

// no implementation of size_hint means it returns (0, None) - which forces from_iter to
// grow the allocated space iteratively.
Expand All @@ -1148,3 +1150,30 @@ fn collect_from_iter() {

let _y: SmallVec<u8, 1> = SmallVec::from_iter(iter);
}

#[test]
fn test_spare_capacity_mut() {
let mut v: SmallVec<u8, 2> = SmallVec::new();
assert!(!v.spilled());
let spare = v.spare_capacity_mut();
assert_eq!(spare.len(), 2);
assert_eq!(spare.as_ptr().cast::<u8>(), v.as_ptr());

v.push(1);
assert!(!v.spilled());
let spare = v.spare_capacity_mut();
assert_eq!(spare.len(), 1);
assert_eq!(spare.as_ptr().cast::<u8>(), unsafe { v.as_ptr().add(1) });

v.push(2);
assert!(!v.spilled());
let spare = v.spare_capacity_mut();
assert_eq!(spare.len(), 0);
assert_eq!(spare.as_ptr().cast::<u8>(), unsafe { v.as_ptr().add(2) });

v.push(3);
assert!(v.spilled());
let spare = v.spare_capacity_mut();
assert!(spare.len() >= 1);
assert_eq!(spare.as_ptr().cast::<u8>(), unsafe { v.as_ptr().add(3) });
}

0 comments on commit 9eacf68

Please sign in to comment.