Added SmallVec::push_mut andSmallVec::insert_mut

src/lib.rs

@@ -1240,16 +1240,39 @@ impl<T, const N: usize> SmallVec<T, N> {
 
     #[inline]
     pub fn push(&mut self, value: T) {
+        _ = self.push_mut(value);
+    }
+
+    #[inline]
+    #[must_use]
+    pub fn push_mut(&mut self, value: T) -> &mut T {
         let len = self.len();
         if len == self.capacity() {
             self.reserve(1);
         }
-        // SAFETY: both the input and output are within the allocation
+
+        // SAFETY: `len < capacity` after the reserve,
+        //         so the offset stays in bounds of the allocation.
         let ptr = unsafe { self.as_mut_ptr().add(len) };
         // SAFETY: we allocated enough space in case it wasn't enough, so the address is
         // valid for writes.
         unsafe { ptr.write(value) };
-        unsafe { self.set_len(len + 1) }
+
+        // LEGAL: all elements in `0..len + 1` are initialized.
+        {
+            // This block is an exact copy of `self.set_len`.
+            // We have to do this so that Miri doesn't report a "Stacked Borrows"
+            // rule violation. See PR/406
+
+            let new_len = len + 1;
+            debug_assert!(new_len <= self.capacity());
+            let on_heap = self.len.on_heap();
+            self.len = TaggedLen::new(new_len, on_heap);
+        }
+
+        // SAFETY: `ptr` is aligned, non-null and points to the element initialized
+        //         above; the borrow is tied to `&mut self`, so it is exclusive.
+        unsafe { &mut *ptr }
     }
 
     #[inline]
@@ -1529,24 +1552,49 @@ impl<T, const N: usize> SmallVec<T, N> {
 
     #[inline]
     pub fn insert(&mut self, index: usize, value: T) {
+        _ = self.insert_mut(index, value);
+    }
+
+    #[inline]
+    #[must_use]
+    pub fn insert_mut(&mut self, index: usize, value: T) -> &mut T {
         let len = self.len();
         assert!(
             index <= len,
             "insertion index (is {index}) should be <= len (is {len})"
         );
         self.reserve(1);
-        let ptr = self.as_mut_ptr();
-        unsafe {
-            // the elements at `index + 1..len + 1` are now initialized
-            if index < len {
-                copy(ptr.add(index), ptr.add(index + 1), len - index);
-            }
-            // the element at `index` is now initialized
-            ptr.add(index).write(value);
 
-            // SAFETY: all the elements are initialized
-            self.set_len(len + 1);
+        // SAFETY: `index <= len <= capacity`,
+        //         so the offset stays in bounds of the allocation.
+        let ptr = unsafe { self.as_mut_ptr().add(index) };
+
+        if index < len {
+            // SAFETY: `reserve(1)` guarantees capacity for `len + 1` elements,
+            //         so shifting `len - index` elements one slot up stays in bounds.
+            //         Source and destination overlap, hence `copy` instead of
+            //         `copy_nonoverlapping`.
+            unsafe { copy(ptr, ptr.add(1), len - index) };
         }
+
+        // SAFETY: the slot at `index` is free and properly aligned for `T`.
+        unsafe { ptr.write(value) };
+
+        // LEGAL: all elements in `0..len + 1` are initialized.
+        {
+            // This block is an exact copy of `self.set_len`.
+            // We have to do this so that Miri doesn't report a "Stacked Borrows"
+            // rule violation. See PR/406
+
+            let new_len = len + 1;
+            debug_assert!(new_len <= self.capacity());
+            let on_heap = self.len.on_heap();
+            self.len = TaggedLen::new(new_len, on_heap);
+        }
+
+        // SAFETY: `ptr` is aligned, non-null and points to the element initialized
+        //         above; the borrow is tied to `&mut self`, so it is exclusive.
+        unsafe { &mut *ptr }
     }
 
     #[inline]

src/tests.rs

@@ -18,6 +18,38 @@ pub fn test_zero() {
 // We heap allocate all these strings so that double frees will show up under
 // valgrind.
 
+#[test]
+pub fn test_push_mut() {
+    let mut v = SmallVec::<_, 16>::new();
+
+    let first_elem = v.push_mut("hello".to_owned());
+    assert_eq!(&*first_elem, &"hello".to_owned());
+
+    *first_elem = "hi".to_owned();
+    assert_eq!(&*first_elem, &"hi".to_owned());
+
+    v.push("there".to_owned());
+    assert_eq!(&*v, &["hi".to_owned(), "there".to_owned(),][..]);
+}
+
+#[test]
+pub fn test_insert_mut() {
+    let mut v = SmallVec::<_, 16>::new();
+    v.push("hello".to_owned());
+    v.push("there".to_owned());
+
+    let second_elem = v.insert_mut(1, ",".to_owned());
+    assert_eq!(&*second_elem, &",".to_owned());
+
+    *second_elem = ";".to_owned();
+    assert_eq!(&*second_elem, &";".to_owned());
+
+    assert_eq!(
+        &*v,
+        &["hello".to_owned(), ";".to_owned(), "there".to_owned(),][..]
+    );
+}
+
 #[test]
 pub fn test_inline() {
     let mut v = SmallVec::<_, 16>::new();