diff --git a/src/lib.rs b/src/lib.rs index 752a0ab..6ad42a5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -880,6 +880,48 @@ impl SmallVec { self.len.on_heap(Self::is_zst()) } + /// Splits the collection into two at the given index. + /// + /// Returns a newly allocated vector containing the elements in the range + /// `[at, len)`. After the call, the original vector will be left containing + /// the elements `[0, at)` with its previous capacity unchanged. + /// + /// - If you want to take ownership of the entire contents and capacity of + /// the vector, see [`mem::take`] or [`mem::replace`]. + /// - If you don't need the returned vector at all, see [`SmallVec::truncate`]. + /// - If you want to take ownership of an arbitrary subslice, or you don't + /// necessarily want to store the removed items in a vector, see [`SmallVec::drain`]. + /// + /// # Panics + /// + /// Panics if `at > len`. + /// + /// # Examples + /// + /// ``` + /// let mut vec = vec![1, 2, 3]; + /// let vec2 = vec.split_off(1); + /// assert_eq!(vec, [1]); + /// assert_eq!(vec2, [2, 3]); + /// ``` + #[inline] + pub fn split_off(&mut self, at: usize) -> Self { + let len = self.len(); + assert!(at <= len); + + let other_len = len - at; + let mut other = Self::with_capacity(other_len); + + // Unsafely `set_len` and copy items to `other`. + unsafe { + self.set_len(at); + other.set_len(other_len); + + core::ptr::copy_nonoverlapping(self.as_ptr().add(at), other.as_mut_ptr(), other.len()); + } + other + } + pub fn drain(&mut self, range: R) -> Drain<'_, T, N> where R: core::ops::RangeBounds, diff --git a/src/tests.rs b/src/tests.rs index b1a4f45..c242c28 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -280,6 +280,40 @@ fn test_truncate_references() { ); } +#[test] +fn test_split_off() { + let mut vec: SmallVec = smallvec![1, 2, 3, 4, 5, 6]; + let orig_ptr = vec.as_ptr(); + let orig_capacity = vec.capacity(); + + let split_off = vec.split_off(4); + assert_eq!(&vec[..], &[1, 2, 3, 4]); + assert_eq!(&split_off[..], &[5, 6]); + assert_eq!(vec.capacity(), orig_capacity); + assert_eq!(vec.as_ptr(), orig_ptr); +} + +#[test] +fn test_split_off_take_all() { + // Allocate enough capacity that we can tell whether the split-off vector's + // capacity is based on its size, or (incorrectly) on the original capacity. + let mut vec = SmallVec::::with_capacity(1000); + vec.extend([1, 2, 3, 4, 5, 6]); + let orig_ptr = vec.as_ptr(); + let orig_capacity = vec.capacity(); + + let split_off = vec.split_off(0); + assert_eq!(&vec[..], &[]); + assert_eq!(&split_off[..], &[1, 2, 3, 4, 5, 6]); + assert_eq!(vec.capacity(), orig_capacity); + assert_eq!(vec.as_ptr(), orig_ptr); + + // The split-off vector should be newly-allocated, and should not have + // stolen the original vector's allocation. + assert!(split_off.capacity() < orig_capacity); + assert_ne!(split_off.as_ptr(), orig_ptr); +} + #[test] fn test_insert_many() { let mut v: SmallVec = SmallVec::new();