diff --git a/src/lib.rs b/src/lib.rs index 3889463..3995c8e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1240,16 +1240,39 @@ impl SmallVec { #[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 SmallVec { #[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] diff --git a/src/tests.rs b/src/tests.rs index 74f732f..d344de7 100644 --- a/src/tests.rs +++ b/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();