From cc1a9a0ed67c194f231fe9478069b57fd0c52761 Mon Sep 17 00:00:00 2001
From: Arthur Silva <arthurprs@gmail.com>
Date: Sun, 3 Mar 2024 18:13:52 +0100
Subject: [PATCH] Implement SmallVec::split_off

---
 src/lib.rs   | 42 ++++++++++++++++++++++++++++++++++++++++++
 src/tests.rs | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 752a0ab..3c07bee 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -880,6 +880,48 @@ impl<T, const N: usize> SmallVec<T, N> {
         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<R>(&mut self, range: R) -> Drain<'_, T, N>
     where
         R: core::ops::RangeBounds<usize>,
diff --git a/src/tests.rs b/src/tests.rs
index b1a4f45..57e0277 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<u32, 4> = 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::<u32, 4>::with_capacity(1000);
+    vec.extend([1, 2, 3, 4, 5, 6]);
+    let orig_ptr = vec.as_ptr();
+    let orig_capacity: usize = 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<u8, 8> = SmallVec::new();