Added SmallVec::push_mut andSmallVec::insert_mut

src/lib.rs

@@ -1232,7 +1232,16 @@ impl<T, const N: usize> SmallVec<T, N> {
         // 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) }
+        unsafe { self.set_len(len + 1) };
+    }
+
+    #[inline]
+    #[must_use]
+    pub fn push_mut(&mut self, value: T) -> &mut T {
+        let idx = self.len();
+        let () = self.push(value);
+        let ptr = unsafe { self.as_mut_ptr().add(idx) };
+        unsafe { &mut *ptr }
     }
 
     #[inline]
@@ -1498,7 +1507,10 @@ impl<T, const N: usize> SmallVec<T, N> {
     #[inline]
     pub fn insert(&mut self, index: usize, value: T) {
         let len = self.len();
-        assert!(index <= len, "insertion index (is {index}) should be <= len (is {len})");
+        assert!(
+            index <= len,
+            "insertion index (is {index}) should be <= len (is {len})"
+        );
         self.reserve(1);
         let ptr = self.as_mut_ptr();
         unsafe {
@@ -1514,6 +1526,14 @@ impl<T, const N: usize> SmallVec<T, N> {
         }
     }
 
+    #[inline]
+    #[must_use]
+    pub fn insert_mut(&mut self, index: usize, value: T) -> &mut T {
+        let () = self.insert(index, value);
+        let ptr = unsafe { self.as_mut_ptr().add(index) };
+        unsafe { &mut *ptr }
+    }
+
     #[inline]
     pub const fn as_slice(&self) -> &[T] {
         let len = self.len();

src/tests.rs

@@ -19,6 +19,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();
@@ -587,10 +619,7 @@ fn test_from() {
 #[test]
 fn test_from_slice() {
     assert_eq!(&SmallVec::<u32, 2>::from(&[1][..])[..], [1]);
-    assert_eq!(
-        &SmallVec::<u32, 2>::from(&[1, 2, 3][..])[..],
-        [1, 2, 3]
-    );
+    assert_eq!(&SmallVec::<u32, 2>::from(&[1, 2, 3][..])[..], [1, 2, 3]);
 }
 
 #[test]