From 9eacf68356ea1582962f6f61f24a8e0fff155cf3 Mon Sep 17 00:00:00 2001
From: DaniPopes <57450786+DaniPopes@users.noreply.github.com>
Date: Tue, 10 Dec 2024 23:42:02 +0100
Subject: [PATCH] Implement spare_capacity_mut (#364)

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.
---
 src/lib.rs   | 16 ++++++++++++++++
 src/tests.rs | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index 803e50e..1fbdc7b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -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
diff --git a/src/tests.rs b/src/tests.rs
index 3186fb5..2f0a1f2 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -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.
@@ -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) });
+}