Skip to content

Commit bb6521e

Browse files
committed
Auto merge of #458 - JustForFun88:simplify-clone, r=Amanieu
Simplify `Clone` by removing redundant guards These extra guards only complicate the code where it is not needed anyway. The `Drop` function checks the number of `items` in a table before dropping elements (see below), so that dropping the uninitialize table as well as table with no actual data (but with `FULL` control bytes) is safe . Added tests to check that the current behavior of `Drop` won't change in the future. ```rust impl<T, A: Allocator + Clone> Drop for RawTable<T, A> { fn drop(&mut self) { if !self.table.is_empty_singleton() { unsafe { self.drop_elements(); self.free_buckets(); } } } } impl<T, A: Allocator + Clone> RawTable<T, A> { unsafe fn drop_elements(&mut self) { if Self::DATA_NEEDS_DROP && !self.is_empty() { for item in self.iter() { item.drop(); } } } } ```
2 parents 5e578e7 + 50a9e7b commit bb6521e

File tree

3 files changed

+227
-57
lines changed

3 files changed

+227
-57
lines changed

src/map.rs

Lines changed: 120 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -6658,7 +6658,7 @@ mod test_map {
66586658
use allocator_api2::alloc::{AllocError, Allocator, Global};
66596659
use core::alloc::Layout;
66606660
use core::ptr::NonNull;
6661-
use core::sync::atomic::{AtomicBool, Ordering};
6661+
use core::sync::atomic::{AtomicI8, Ordering};
66626662
use rand::{rngs::SmallRng, Rng, SeedableRng};
66636663
use std::borrow::ToOwned;
66646664
use std::cell::RefCell;
@@ -8510,47 +8510,44 @@ mod test_map {
85108510
let _map2 = map1.clone();
85118511
}
85128512

8513-
#[test]
8514-
fn test_hashmap_into_iter_bug() {
8515-
let dropped: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
8513+
struct MyAllocInner {
8514+
drop_count: Arc<AtomicI8>,
8515+
}
85168516

8517-
{
8518-
struct MyAllocInner {
8519-
drop_flag: Arc<AtomicBool>,
8520-
}
8517+
#[derive(Clone)]
8518+
struct MyAlloc {
8519+
_inner: Arc<MyAllocInner>,
8520+
}
85218521

8522-
#[derive(Clone)]
8523-
struct MyAlloc {
8524-
_inner: Arc<MyAllocInner>,
8525-
}
8522+
impl Drop for MyAllocInner {
8523+
fn drop(&mut self) {
8524+
println!("MyAlloc freed.");
8525+
self.drop_count.fetch_sub(1, Ordering::SeqCst);
8526+
}
8527+
}
85268528

8527-
impl Drop for MyAllocInner {
8528-
fn drop(&mut self) {
8529-
println!("MyAlloc freed.");
8530-
self.drop_flag.store(true, Ordering::SeqCst);
8531-
}
8532-
}
8529+
unsafe impl Allocator for MyAlloc {
8530+
fn allocate(&self, layout: Layout) -> std::result::Result<NonNull<[u8]>, AllocError> {
8531+
let g = Global;
8532+
g.allocate(layout)
8533+
}
85338534

8534-
unsafe impl Allocator for MyAlloc {
8535-
fn allocate(
8536-
&self,
8537-
layout: Layout,
8538-
) -> std::result::Result<NonNull<[u8]>, AllocError> {
8539-
let g = Global;
8540-
g.allocate(layout)
8541-
}
8535+
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
8536+
let g = Global;
8537+
g.deallocate(ptr, layout)
8538+
}
8539+
}
85428540

8543-
unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
8544-
let g = Global;
8545-
g.deallocate(ptr, layout)
8546-
}
8547-
}
8541+
#[test]
8542+
fn test_hashmap_into_iter_bug() {
8543+
let dropped: Arc<AtomicI8> = Arc::new(AtomicI8::new(1));
85488544

8545+
{
85498546
let mut map = crate::HashMap::with_capacity_in(
85508547
10,
85518548
MyAlloc {
85528549
_inner: Arc::new(MyAllocInner {
8553-
drop_flag: dropped.clone(),
8550+
drop_count: dropped.clone(),
85548551
}),
85558552
},
85568553
);
@@ -8563,6 +8560,96 @@ mod test_map {
85638560
}
85648561
}
85658562

8566-
assert!(dropped.load(Ordering::SeqCst));
8563+
// All allocator clones should already be dropped.
8564+
assert_eq!(dropped.load(Ordering::SeqCst), 0);
8565+
}
8566+
8567+
#[test]
8568+
#[should_panic = "panic in clone"]
8569+
fn test_clone_memory_leaks_and_double_drop() {
8570+
let dropped: Arc<AtomicI8> = Arc::new(AtomicI8::new(2));
8571+
8572+
{
8573+
let mut map = crate::HashMap::with_capacity_in(
8574+
10,
8575+
MyAlloc {
8576+
_inner: Arc::new(MyAllocInner {
8577+
drop_count: dropped.clone(),
8578+
}),
8579+
},
8580+
);
8581+
8582+
const DISARMED: bool = false;
8583+
const ARMED: bool = true;
8584+
8585+
struct CheckedCloneDrop {
8586+
panic_in_clone: bool,
8587+
dropped: bool,
8588+
need_drop: Vec<i32>,
8589+
}
8590+
8591+
impl Clone for CheckedCloneDrop {
8592+
fn clone(&self) -> Self {
8593+
if self.panic_in_clone {
8594+
panic!("panic in clone")
8595+
}
8596+
Self {
8597+
panic_in_clone: self.panic_in_clone,
8598+
dropped: self.dropped,
8599+
need_drop: self.need_drop.clone(),
8600+
}
8601+
}
8602+
}
8603+
8604+
impl Drop for CheckedCloneDrop {
8605+
fn drop(&mut self) {
8606+
if self.dropped {
8607+
panic!("double drop");
8608+
}
8609+
self.dropped = true;
8610+
}
8611+
}
8612+
8613+
let armed_flags = [
8614+
DISARMED, DISARMED, ARMED, DISARMED, DISARMED, DISARMED, DISARMED,
8615+
];
8616+
8617+
// Hash and Key must be equal to each other for controlling the elements placement
8618+
// so that we can be sure that we first clone elements that don't panic during cloning.
8619+
for (idx, &panic_in_clone) in armed_flags.iter().enumerate() {
8620+
let idx = idx as u64;
8621+
map.table.insert(
8622+
idx,
8623+
(
8624+
idx,
8625+
CheckedCloneDrop {
8626+
panic_in_clone,
8627+
dropped: false,
8628+
need_drop: vec![0, 1, 2, 3],
8629+
},
8630+
),
8631+
|(k, _)| *k,
8632+
);
8633+
}
8634+
8635+
let mut count = 0;
8636+
// Let's check that all elements are located as we wanted
8637+
//
8638+
// SAFETY: We know for sure that `RawTable` will outlive
8639+
// the returned `RawIter` iterator.
8640+
for ((key, value), panic_in_clone) in map.iter().zip(armed_flags) {
8641+
assert_eq!(*key, count);
8642+
assert_eq!(value.panic_in_clone, panic_in_clone);
8643+
count += 1;
8644+
}
8645+
assert_eq!(map.len(), 7);
8646+
assert_eq!(count, 7);
8647+
8648+
// Clone should normally clone a few elements, and then (when the
8649+
// clone function panics), deallocate both its own memory, memory
8650+
// of `dropped: Arc<AtomicI8>` and the memory of already cloned
8651+
// elements (Vec<i32> memory inside CheckedCloneDrop).
8652+
let _table2 = map.clone();
8653+
}
85678654
}
85688655
}

src/raw/mod.rs

Lines changed: 106 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use crate::alloc::alloc::{handle_alloc_error, Layout};
2-
use crate::scopeguard::{guard, ScopeGuard};
2+
use crate::scopeguard::guard;
33
use crate::TryReserveError;
44
use core::iter::FusedIterator;
55
use core::marker::PhantomData;
@@ -2494,7 +2494,11 @@ impl<T: Clone, A: Allocator + Clone> Clone for RawTable<T, A> {
24942494
} else {
24952495
unsafe {
24962496
// Avoid `Result::ok_or_else` because it bloats LLVM IR.
2497-
let new_table = match Self::new_uninitialized(
2497+
//
2498+
// SAFETY: This is safe as we are taking the size of an already allocated table
2499+
// and therefore сapacity overflow cannot occur, `self.table.buckets()` is power
2500+
// of two and all allocator errors will be caught inside `RawTableInner::new_uninitialized`.
2501+
let mut new_table = match Self::new_uninitialized(
24982502
self.table.alloc.clone(),
24992503
self.table.buckets(),
25002504
Fallibility::Infallible,
@@ -2503,29 +2507,29 @@ impl<T: Clone, A: Allocator + Clone> Clone for RawTable<T, A> {
25032507
Err(_) => hint::unreachable_unchecked(),
25042508
};
25052509

2506-
// If cloning fails then we need to free the allocation for the
2507-
// new table. However we don't run its drop since its control
2508-
// bytes are not initialized yet.
2509-
let mut guard = guard(ManuallyDrop::new(new_table), |new_table| {
2510-
new_table.free_buckets();
2511-
});
2512-
2513-
guard.clone_from_spec(self);
2514-
2515-
// Disarm the scope guard and return the newly created table.
2516-
ManuallyDrop::into_inner(ScopeGuard::into_inner(guard))
2510+
// Cloning elements may fail (the clone function may panic). But we don't
2511+
// need to worry about uninitialized control bits, since:
2512+
// 1. The number of items (elements) in the table is zero, which means that
2513+
// the control bits will not be readed by Drop function.
2514+
// 2. The `clone_from_spec` method will first copy all control bits from
2515+
// `self` (thus initializing them). But this will not affect the `Drop`
2516+
// function, since the `clone_from_spec` function sets `items` only after
2517+
// successfully clonning all elements.
2518+
new_table.clone_from_spec(self);
2519+
new_table
25172520
}
25182521
}
25192522
}
25202523

25212524
fn clone_from(&mut self, source: &Self) {
25222525
if source.table.is_empty_singleton() {
2526+
// Dereference drops old `self` table
25232527
*self = Self::new_in(self.table.alloc.clone());
25242528
} else {
25252529
unsafe {
25262530
// Make sure that if any panics occurs, we clear the table and
25272531
// leave it in an empty state.
2528-
let mut self_ = guard(self, |self_| {
2532+
let mut guard = guard(&mut *self, |self_| {
25292533
self_.clear_no_drop();
25302534
});
25312535

@@ -2535,18 +2539,32 @@ impl<T: Clone, A: Allocator + Clone> Clone for RawTable<T, A> {
25352539
//
25362540
// This leak is unavoidable: we can't try dropping more elements
25372541
// since this could lead to another panic and abort the process.
2538-
self_.drop_elements();
2542+
//
2543+
// SAFETY: We clear our table right after dropping the elements,
2544+
// so there is no double drop, since `items` will be equal to zero.
2545+
guard.drop_elements();
2546+
2547+
// Okay, we've successfully dropped all elements, so we'll just set
2548+
// `items` to zero (so that the `Drop` of `RawTable` doesn't try to
2549+
// drop all elements twice) and just forget about the guard.
2550+
guard.table.items = 0;
2551+
mem::forget(guard);
25392552

25402553
// If necessary, resize our table to match the source.
2541-
if self_.buckets() != source.buckets() {
2554+
if self.buckets() != source.buckets() {
25422555
// Skip our drop by using ptr::write.
2543-
if !self_.table.is_empty_singleton() {
2544-
self_.free_buckets();
2556+
if !self.table.is_empty_singleton() {
2557+
// SAFETY: We have verified that the table is allocated.
2558+
self.free_buckets();
25452559
}
2546-
(&mut **self_ as *mut Self).write(
2560+
(self as *mut Self).write(
25472561
// Avoid `Result::unwrap_or_else` because it bloats LLVM IR.
2562+
//
2563+
// SAFETY: This is safe as we are taking the size of an already allocated table
2564+
// and therefore сapacity overflow cannot occur, `self.table.buckets()` is power
2565+
// of two and all allocator errors will be caught inside `RawTableInner::new_uninitialized`.
25482566
match Self::new_uninitialized(
2549-
self_.table.alloc.clone(),
2567+
self.table.alloc.clone(),
25502568
source.buckets(),
25512569
Fallibility::Infallible,
25522570
) {
@@ -2556,10 +2574,11 @@ impl<T: Clone, A: Allocator + Clone> Clone for RawTable<T, A> {
25562574
);
25572575
}
25582576

2559-
self_.clone_from_spec(source);
2560-
2561-
// Disarm the scope guard if cloning was successful.
2562-
ScopeGuard::into_inner(self_);
2577+
// Cloning elements may fail (the clone function may panic), but the `ScopeGuard`
2578+
// inside the `clone_from_impl` function will take care of that, dropping all
2579+
// cloned elements if necessary. The `Drop` of `RawTable` takes care of the rest
2580+
// by freeing up the allocated memory.
2581+
self.clone_from_spec(source);
25632582
}
25642583
}
25652584
}
@@ -3373,4 +3392,67 @@ mod test_map {
33733392
assert!(table.find(i + 100, |x| *x == i + 100).is_none());
33743393
}
33753394
}
3395+
3396+
/// CHECKING THAT WE ARE NOT TRYING TO READ THE MEMORY OF
3397+
/// AN UNINITIALIZED TABLE DURING THE DROP
3398+
#[test]
3399+
fn test_drop_uninitialized() {
3400+
use ::alloc::vec::Vec;
3401+
3402+
let table = unsafe {
3403+
// SAFETY: The `buckets` is power of two and we're not
3404+
// trying to actually use the returned RawTable.
3405+
RawTable::<(u64, Vec<i32>)>::new_uninitialized(Global, 8, Fallibility::Infallible)
3406+
.unwrap()
3407+
};
3408+
drop(table);
3409+
}
3410+
3411+
/// CHECKING THAT WE DON'T TRY TO DROP DATA IF THE `ITEMS`
3412+
/// ARE ZERO, EVEN IF WE HAVE `FULL` CONTROL BYTES.
3413+
#[test]
3414+
fn test_drop_zero_items() {
3415+
use ::alloc::vec::Vec;
3416+
unsafe {
3417+
// SAFETY: The `buckets` is power of two and we're not
3418+
// trying to actually use the returned RawTable.
3419+
let table =
3420+
RawTable::<(u64, Vec<i32>)>::new_uninitialized(Global, 8, Fallibility::Infallible)
3421+
.unwrap();
3422+
3423+
// WE SIMULATE, AS IT WERE, A FULL TABLE.
3424+
3425+
// SAFETY: We checked that the table is allocated and therefore the table already has
3426+
// `self.bucket_mask + 1 + Group::WIDTH` number of control bytes (see TableLayout::calculate_layout_for)
3427+
// so writing `table.table.num_ctrl_bytes() == bucket_mask + 1 + Group::WIDTH` bytes is safe.
3428+
table
3429+
.table
3430+
.ctrl(0)
3431+
.write_bytes(EMPTY, table.table.num_ctrl_bytes());
3432+
3433+
// SAFETY: table.capacity() is guaranteed to be smaller than table.buckets()
3434+
table.table.ctrl(0).write_bytes(0, table.capacity());
3435+
3436+
// Fix up the trailing control bytes. See the comments in set_ctrl
3437+
// for the handling of tables smaller than the group width.
3438+
if table.buckets() < Group::WIDTH {
3439+
// SAFETY: We have `self.bucket_mask + 1 + Group::WIDTH` number of control bytes,
3440+
// so copying `self.buckets() == self.bucket_mask + 1` bytes with offset equal to
3441+
// `Group::WIDTH` is safe
3442+
table
3443+
.table
3444+
.ctrl(0)
3445+
.copy_to(table.table.ctrl(Group::WIDTH), table.table.buckets());
3446+
} else {
3447+
// SAFETY: We have `self.bucket_mask + 1 + Group::WIDTH` number of
3448+
// control bytes,so copying `Group::WIDTH` bytes with offset equal
3449+
// to `self.buckets() == self.bucket_mask + 1` is safe
3450+
table
3451+
.table
3452+
.ctrl(0)
3453+
.copy_to(table.table.ctrl(table.table.buckets()), Group::WIDTH);
3454+
}
3455+
drop(table);
3456+
}
3457+
}
33763458
}

src/scopeguard.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ impl<T, F> ScopeGuard<T, F>
2525
where
2626
F: FnMut(&mut T),
2727
{
28+
#[allow(dead_code)]
2829
#[inline]
2930
pub fn into_inner(guard: Self) -> T {
3031
// Cannot move out of Drop-implementing types, so

0 commit comments

Comments
 (0)