diff --git a/src/lib.rs b/src/lib.rs index e6bc8c7..07d9b53 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1232,7 +1232,16 @@ impl SmallVec { // 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 SmallVec { #[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 SmallVec { } } + #[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(); diff --git a/src/tests.rs b/src/tests.rs index 2082f96..d8098f0 100644 --- a/src/tests.rs +++ b/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::::from(&[1][..])[..], [1]); - assert_eq!( - &SmallVec::::from(&[1, 2, 3][..])[..], - [1, 2, 3] - ); + assert_eq!(&SmallVec::::from(&[1, 2, 3][..])[..], [1, 2, 3]); } #[test]