diff --git a/Cargo.toml b/Cargo.toml index a838593..7c4e83c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,9 +5,7 @@ edition = "2024" authors = ["Kyle Thapar "] description = "Stack-allocated collections for Rust" license = "MIT OR Apache-2.0" -readme = "README.md" repository = "https://github.com/OverwrittenCode/stack_collections" -homepage = "https://github.com/OverwrittenCode/stack_collections" documentation = "https://docs.rs/stack_collections" keywords = ["stack", "collections", "no-heap"] categories = ["data-structures", "no-std", "no-std::no-alloc"] diff --git a/src/stack_vec.rs b/src/stack_vec.rs index 03be4ea..64bc838 100644 --- a/src/stack_vec.rs +++ b/src/stack_vec.rs @@ -461,34 +461,80 @@ impl StackVec { where F: FnMut(&mut T) -> bool, { - let base_ptr = self.as_mut_ptr(); - let mut kept = 0; + struct BackshiftOnDrop<'vec, T, const CAP: usize> { + vec: &'vec mut StackVec, + processed: usize, + deleted: usize, + original_len: usize, + } + + impl Drop for BackshiftOnDrop<'_, T, CAP> { + fn drop(&mut self) { + if self.deleted > 0 { + let base = self.vec.as_mut_ptr(); + // SAFETY: processed and processed - deleted are both + // within bounds of the buffer. + let src = unsafe { base.add(self.processed) }; + // SAFETY: see above. + let dst = unsafe { base.add(self.processed - self.deleted) }; + // SAFETY: the unprocessed tail is still initialized and + // the destination range is inside the buffer. + // src/dst derive from a single reborrow, so + // neither invalidates the other. + unsafe { + ptr::copy(src, dst, self.original_len - self.processed); + } + } + self.vec.len = self.original_len - self.deleted; + } + } + + // Both the predicate and a removed element's `Drop` are user code and + // may unwind. Retire the whole buffer up front and let a guard restore + // a correct length on every exit path, so a duplicated or + // already-dropped slot can never stay inside `0..len`. + let original_len = self.len; + self.len = 0; + + let mut g = BackshiftOnDrop { + vec: self, + processed: 0, + deleted: 0, + original_len, + }; + + while g.processed < original_len { + let base = g.vec.as_mut_ptr(); - for i in 0..self.len { - // SAFETY: i < self.len, so this is a valid and initialized element - let ptr = unsafe { base_ptr.add(i) }; - // SAFETY: ptr points to valid initialized element + // SAFETY: processed < original_len, so this element is initialized. + let ptr = unsafe { base.add(g.processed) }; + // SAFETY: ptr points to a valid initialized element. let elem = unsafe { &mut *ptr }; if f(elem) { - if kept != i { - // SAFETY: kept < i < self.len, ranges don't overlap - let dst = unsafe { base_ptr.add(kept) }; - // SAFETY: Copying single element from valid source to valid dest + if g.deleted > 0 { + // SAFETY: dst is an earlier, already-vacated slot; ptr and dst + // both derive from the same reborrow, so neither + // invalidates the other. + let dst = unsafe { base.add(g.processed - g.deleted) }; + // SAFETY: copying a single initialized element to a vacated slot. unsafe { ptr::copy_nonoverlapping(ptr, dst, 1); } } - kept += 1; + g.processed += 1; } else { - // SAFETY: Dropping initialized element that won't be kept + // Count the slot out before destroying it, so an unwinding + // `Drop` leaves the guard with a consistent view. + g.processed += 1; + g.deleted += 1; + // SAFETY: dropping an initialized element that won't be kept. unsafe { - ptr::drop_in_place(elem); + ptr::drop_in_place(ptr); } } } - - self.len = kept; + // `g`'s `Drop` commits the final length. } // index @@ -1837,4 +1883,84 @@ mod tests { drop(iter); assert_eq!(counter.load(Ordering::SeqCst), 3); } + + /// A panic in `retain` must not leave a duplicated or already-destroyed + /// slot inside `0..len`. + #[cfg(feature = "std")] + #[test] + fn retain_panicking_predicate() { + extern crate std; + use core::panic::AssertUnwindSafe; + use std::panic::catch_unwind; + + struct DropCounter(i32, Arc); + impl Drop for DropCounter { + fn drop(&mut self) { + self.1.fetch_add(1, Ordering::SeqCst); + } + } + + let counter = Arc::new(AtomicUsize::new(0)); + let mut vec = StackVec::::new(); + for i in 0..4_i32 { + vec.push(DropCounter(i, Arc::clone(&counter))); + } + + // Reject element 0 so the next kept one is compacted forward, then + // unwind on element 2 before the length is committed. + let r = catch_unwind(AssertUnwindSafe(|| { + vec.retain(|e| { + assert_ne!(e.0, 2_i32, "predicate panics"); + e.0 != 0_i32 + }); + })); + assert!(r.is_err(), "the predicate should have panicked"); + + drop(vec); + + // Four elements exist. Fewer drops mean a leak, which is sound; more + // mean a value was destroyed twice. + let drops = counter.load(Ordering::SeqCst); + assert!(drops <= 4, "{drops} drops for 4 elements"); + } + + /// The same, with the panic coming from a removed element's `Drop`. + #[cfg(feature = "std")] + #[test] + fn retain_panicking_element_drop() { + extern crate std; + use core::cell::Cell; + use core::panic::AssertUnwindSafe; + use std::panic::catch_unwind; + + std::thread_local! { + static ARMED: Cell = const { Cell::new(false) }; + } + + struct Boom(i32, Arc); + impl Drop for Boom { + fn drop(&mut self) { + self.1.fetch_add(1, Ordering::SeqCst); + assert!(!ARMED.replace(false), "element Drop panics"); + } + } + + let counter = Arc::new(AtomicUsize::new(0)); + let mut vec = StackVec::::new(); + for i in 0..4_i32 { + vec.push(Boom(i, Arc::clone(&counter))); + } + + ARMED.set(true); + let r = catch_unwind(AssertUnwindSafe(|| { + vec.retain(|e| e.0 != 1_i32); + })); + ARMED.set(false); + assert!(r.is_err(), "the armed Drop should have panicked"); + + drop(vec); + + let drops = counter.load(Ordering::SeqCst); + assert!(drops <= 4, "{drops} drops for 4 elements"); + } }