From d6dec1ebe384f29cc99435412f581c77ad131708 Mon Sep 17 00:00:00 2001 From: oxalica Date: Mon, 18 Jan 2021 00:51:37 +0800 Subject: [PATCH 1/4] Optimize Vec::retain --- library/alloc/src/vec/mod.rs | 79 +++++++++++++++++++++++++++++++----- library/alloc/tests/vec.rs | 72 ++++++++++++++++++++++++++++++++ 2 files changed, 140 insertions(+), 11 deletions(-) diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index ccc4f03a1e505..2c510b8b2ae61 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -1371,21 +1371,78 @@ impl Vec { F: FnMut(&T) -> bool, { let len = self.len(); - let mut del = 0; - { - let v = &mut **self; - - for i in 0..len { - if !f(&v[i]) { - del += 1; - } else if del > 0 { - v.swap(i - del, i); + // Avoid double drop if the drop guard is not executed, + // since we may make some holes during the process. + unsafe { self.set_len(0) }; + + // Vec: [Kept, Kept, Hole, Hole, Hole, Hole, Unchecked, Unchecked] + // |<- processed len ->| ^- next to check + // |<- deleted cnt ->| + // |<- original_len ->| + // Kept: Elements which predicate returns true on. + // Hole: Moved or dropped element slot. + // Unchecked: Unchecked valid elements. + // + // This drop guard will be invoked when predicate or `drop` of element panicked. + // It shifts unchecked elements to cover holes and `set_len` to the correct length. + // In cases when predicate and `drop` never panick, it will be optimized out. + struct BackshiftOnDrop<'a, T, A: Allocator> { + v: &'a mut Vec, + processed_len: usize, + deleted_cnt: usize, + original_len: usize, + } + + impl Drop for BackshiftOnDrop<'_, T, A> { + fn drop(&mut self) { + if self.deleted_cnt > 0 { + // SAFETY: Fill the hole of dropped or moved + unsafe { + ptr::copy( + self.v.as_ptr().offset(self.processed_len as isize), + self.v + .as_mut_ptr() + .offset(self.processed_len as isize - self.deleted_cnt as isize), + self.original_len - self.processed_len, + ); + self.v.set_len(self.original_len - self.deleted_cnt); + } } } } - if del > 0 { - self.truncate(len - del); + + let mut guard = BackshiftOnDrop { + v: self, + processed_len: 0, + deleted_cnt: 0, + original_len: len, + }; + + let mut del = 0usize; + for i in 0..len { + // SAFETY: Unchecked element must be valid. + let cur = unsafe { &mut *guard.v.as_mut_ptr().offset(i as isize) }; + if !f(cur) { + del += 1; + // Advance early to avoid double drop if `drop_in_place` panicked. + guard.processed_len = i + 1; + guard.deleted_cnt = del; + // SAFETY: We never touch this element again after dropped. + unsafe { ptr::drop_in_place(cur) }; + } else if del > 0 { + // SAFETY: `del` > 0 so the hole slot must not overlap with current element. + // We use copy for move, and never touch this element again. + unsafe { + let hole_slot = guard.v.as_mut_ptr().offset(i as isize - del as isize); + ptr::copy_nonoverlapping(cur, hole_slot, 1); + } + guard.processed_len = i + 1; + } } + + // All holes are at the end now. Simply cut them out. + unsafe { guard.v.set_len(len - del) }; + mem::forget(guard); } /// Removes all but the first of consecutive elements in the vector that resolve to the same diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index e19406d7a0697..856efb1d3a98e 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -287,6 +287,78 @@ fn test_retain() { assert_eq!(vec, [2, 4]); } +#[test] +fn test_retain_pred_panic() { + use std::sync::atomic::{AtomicU64, Ordering}; + + struct Wrap<'a>(&'a AtomicU64, u64, bool); + + impl Drop for Wrap<'_> { + fn drop(&mut self) { + self.0.fetch_or(self.1, Ordering::SeqCst); + } + } + + let dropped = AtomicU64::new(0); + + let ret = std::panic::catch_unwind(|| { + let mut v = vec![ + Wrap(&dropped, 1, false), + Wrap(&dropped, 2, false), + Wrap(&dropped, 4, false), + Wrap(&dropped, 8, false), + Wrap(&dropped, 16, false), + ]; + v.retain(|w| match w.1 { + 1 => true, + 2 => false, + 4 => true, + _ => panic!(), + }); + }); + assert!(ret.is_err()); + // Everything is dropped when predicate panicked. + assert_eq!(dropped.load(Ordering::SeqCst), 1 | 2 | 4 | 8 | 16); +} + +#[test] +fn test_retain_drop_panic() { + use std::sync::atomic::{AtomicU64, Ordering}; + + struct Wrap<'a>(&'a AtomicU64, u64); + + impl Drop for Wrap<'_> { + fn drop(&mut self) { + if self.1 == 8 { + panic!(); + } + self.0.fetch_or(self.1, Ordering::SeqCst); + } + } + + let dropped = AtomicU64::new(0); + + let ret = std::panic::catch_unwind(|| { + let mut v = vec![ + Wrap(&dropped, 1), + Wrap(&dropped, 2), + Wrap(&dropped, 4), + Wrap(&dropped, 8), + Wrap(&dropped, 16), + ]; + v.retain(|w| match w.1 { + 1 => true, + 2 => false, + 4 => true, + 8 => false, + _ => true, + }); + }); + assert!(ret.is_err()); + // Other elements are dropped when `drop` of one element panicked. + assert_eq!(dropped.load(Ordering::SeqCst), 1 | 2 | 4 | 16); +} + #[test] fn test_dedup() { fn case(a: Vec, b: Vec) { From 16deaec7f95b7478c9d9538f45ef79a0e4535b0f Mon Sep 17 00:00:00 2001 From: oxalica Date: Mon, 18 Jan 2021 02:11:01 +0800 Subject: [PATCH 2/4] Format code --- library/alloc/src/vec/mod.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 2c510b8b2ae61..9fbf06c8e3059 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -1411,12 +1411,8 @@ impl Vec { } } - let mut guard = BackshiftOnDrop { - v: self, - processed_len: 0, - deleted_cnt: 0, - original_len: len, - }; + let mut guard = + BackshiftOnDrop { v: self, processed_len: 0, deleted_cnt: 0, original_len: len }; let mut del = 0usize; for i in 0..len { From 969e5523554a6d53a385be35bab449d9b69a82d4 Mon Sep 17 00:00:00 2001 From: oxalica Date: Sat, 23 Jan 2021 04:00:55 +0800 Subject: [PATCH 3/4] Simplify and fix tests --- library/alloc/tests/vec.rs | 91 +++++++++++++++++--------------------- 1 file changed, 40 insertions(+), 51 deletions(-) diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index 856efb1d3a98e..2e56e06b4d0b1 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -288,75 +288,64 @@ fn test_retain() { } #[test] -fn test_retain_pred_panic() { - use std::sync::atomic::{AtomicU64, Ordering}; - - struct Wrap<'a>(&'a AtomicU64, u64, bool); - - impl Drop for Wrap<'_> { - fn drop(&mut self) { - self.0.fetch_or(self.1, Ordering::SeqCst); - } - } - - let dropped = AtomicU64::new(0); +fn test_retain_pred_panic_with_hole() { + let v = (0..5).map(Rc::new).collect::>(); + catch_unwind(AssertUnwindSafe(|| { + let mut v = v.clone(); + v.retain(|r| match **r { + 0 => true, + 1 => false, + 2 => true, + _ => panic!(), + }); + })) + .unwrap_err(); + // Everything is dropped when predicate panicked. + assert!(v.iter().all(|r| Rc::strong_count(r) == 1)); +} - let ret = std::panic::catch_unwind(|| { - let mut v = vec![ - Wrap(&dropped, 1, false), - Wrap(&dropped, 2, false), - Wrap(&dropped, 4, false), - Wrap(&dropped, 8, false), - Wrap(&dropped, 16, false), - ]; - v.retain(|w| match w.1 { - 1 => true, - 2 => false, - 4 => true, +#[test] +fn test_retain_pred_panic_no_hole() { + let v = (0..5).map(Rc::new).collect::>(); + catch_unwind(AssertUnwindSafe(|| { + let mut v = v.clone(); + v.retain(|r| match **r { + 0 | 1 | 2 => true, _ => panic!(), }); - }); - assert!(ret.is_err()); + })) + .unwrap_err(); // Everything is dropped when predicate panicked. - assert_eq!(dropped.load(Ordering::SeqCst), 1 | 2 | 4 | 8 | 16); + assert!(v.iter().all(|r| Rc::strong_count(r) == 1)); } #[test] fn test_retain_drop_panic() { - use std::sync::atomic::{AtomicU64, Ordering}; - - struct Wrap<'a>(&'a AtomicU64, u64); + struct Wrap(Rc); - impl Drop for Wrap<'_> { + impl Drop for Wrap { fn drop(&mut self) { - if self.1 == 8 { + if *self.0 == 3 { panic!(); } - self.0.fetch_or(self.1, Ordering::SeqCst); } } - let dropped = AtomicU64::new(0); - - let ret = std::panic::catch_unwind(|| { - let mut v = vec![ - Wrap(&dropped, 1), - Wrap(&dropped, 2), - Wrap(&dropped, 4), - Wrap(&dropped, 8), - Wrap(&dropped, 16), - ]; - v.retain(|w| match w.1 { - 1 => true, - 2 => false, - 4 => true, - 8 => false, + let v = (0..5).map(|x| Rc::new(x)).collect::>(); + catch_unwind(AssertUnwindSafe(|| { + let mut v = v.iter().map(|r| Wrap(r.clone())).collect::>(); + v.retain(|w| match *w.0 { + 0 => true, + 1 => false, + 2 => true, + 3 => false, // Drop panic. _ => true, }); - }); - assert!(ret.is_err()); + })) + .unwrap_err(); // Other elements are dropped when `drop` of one element panicked. - assert_eq!(dropped.load(Ordering::SeqCst), 1 | 2 | 4 | 16); + // The panicked wrapper also has its Rc dropped. + assert!(v.iter().all(|r| Rc::strong_count(r) == 1)); } #[test] From 2a11c57fb04d68c261dc14f6cf1aa3f8f14f0a2e Mon Sep 17 00:00:00 2001 From: oxalica Date: Sun, 24 Jan 2021 00:11:04 +0800 Subject: [PATCH 4/4] Fix and simplify --- library/alloc/src/vec/mod.rs | 44 ++++++++++++++++++------------------ 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 9fbf06c8e3059..aa6c0d1f06dd0 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -1370,7 +1370,7 @@ impl Vec { where F: FnMut(&T) -> bool, { - let len = self.len(); + let original_len = self.len(); // Avoid double drop if the drop guard is not executed, // since we may make some holes during the process. unsafe { self.set_len(0) }; @@ -1396,49 +1396,49 @@ impl Vec { impl Drop for BackshiftOnDrop<'_, T, A> { fn drop(&mut self) { if self.deleted_cnt > 0 { - // SAFETY: Fill the hole of dropped or moved + // SAFETY: Trailing unchecked items must be valid since we never touch them. unsafe { ptr::copy( - self.v.as_ptr().offset(self.processed_len as isize), - self.v - .as_mut_ptr() - .offset(self.processed_len as isize - self.deleted_cnt as isize), + self.v.as_ptr().add(self.processed_len), + self.v.as_mut_ptr().add(self.processed_len - self.deleted_cnt), self.original_len - self.processed_len, ); - self.v.set_len(self.original_len - self.deleted_cnt); } } + // SAFETY: After filling holes, all items are in contiguous memory. + unsafe { + self.v.set_len(self.original_len - self.deleted_cnt); + } } } - let mut guard = - BackshiftOnDrop { v: self, processed_len: 0, deleted_cnt: 0, original_len: len }; + let mut g = BackshiftOnDrop { v: self, processed_len: 0, deleted_cnt: 0, original_len }; - let mut del = 0usize; - for i in 0..len { + while g.processed_len < original_len { // SAFETY: Unchecked element must be valid. - let cur = unsafe { &mut *guard.v.as_mut_ptr().offset(i as isize) }; + let cur = unsafe { &mut *g.v.as_mut_ptr().add(g.processed_len) }; if !f(cur) { - del += 1; // Advance early to avoid double drop if `drop_in_place` panicked. - guard.processed_len = i + 1; - guard.deleted_cnt = del; + g.processed_len += 1; + g.deleted_cnt += 1; // SAFETY: We never touch this element again after dropped. unsafe { ptr::drop_in_place(cur) }; - } else if del > 0 { - // SAFETY: `del` > 0 so the hole slot must not overlap with current element. + // We already advanced the counter. + continue; + } + if g.deleted_cnt > 0 { + // SAFETY: `deleted_cnt` > 0, so the hole slot must not overlap with current element. // We use copy for move, and never touch this element again. unsafe { - let hole_slot = guard.v.as_mut_ptr().offset(i as isize - del as isize); + let hole_slot = g.v.as_mut_ptr().add(g.processed_len - g.deleted_cnt); ptr::copy_nonoverlapping(cur, hole_slot, 1); } - guard.processed_len = i + 1; } + g.processed_len += 1; } - // All holes are at the end now. Simply cut them out. - unsafe { guard.v.set_len(len - del) }; - mem::forget(guard); + // All item are processed. This can be optimized to `set_len` by LLVM. + drop(g); } /// Removes all but the first of consecutive elements in the vector that resolve to the same