Auto merge of #439 - Amanieu:fix_alloc_leak, r=Amanieu

Fix leaking of allocator in `RawIntoIter` and `RawIntoParIter`

The allocator was not getting properly dropped when it was moved into the iterator.

Fixes #438

Author: bors

src/external_trait_impls/rayon/raw.rs

@@ -1,7 +1,6 @@
 use crate::raw::Bucket;
 use crate::raw::{Allocator, Global, RawIter, RawIterRange, RawTable};
 use crate::scopeguard::guard;
-use alloc::alloc::dealloc;
 use core::marker::PhantomData;
 use core::mem;
 use core::ptr::NonNull;
@@ -97,9 +96,9 @@ impl<T: Send, A: Allocator + Clone + Send> ParallelIterator for RawIntoParIter<T
     {
         let iter = unsafe { self.table.iter().iter };
         let _guard = guard(self.table.into_allocation(), |alloc| {
-            if let Some((ptr, layout)) = *alloc {
+            if let Some((ptr, layout, ref alloc)) = *alloc {
                 unsafe {
-                    dealloc(ptr.as_ptr(), layout);
+                    alloc.deallocate(ptr, layout);
                 }
             }
         });

src/map.rs

@@ -6653,6 +6653,12 @@ mod test_map {
     use super::Entry::{Occupied, Vacant};
     use super::EntryRef;
     use super::{HashMap, RawEntryMut};
+    use alloc::string::ToString;
+    use alloc::sync::Arc;
+    use allocator_api2::alloc::{AllocError, Allocator, Global};
+    use core::alloc::Layout;
+    use core::ptr::NonNull;
+    use core::sync::atomic::{AtomicBool, Ordering};
     use rand::{rngs::SmallRng, Rng, SeedableRng};
     use std::borrow::ToOwned;
     use std::cell::RefCell;
@@ -8503,4 +8509,60 @@ mod test_map {
         );
         let _map2 = map1.clone();
     }
+
+    #[test]
+    fn test_hashmap_into_iter_bug() {
+        let dropped: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
+
+        {
+            struct MyAllocInner {
+                drop_flag: Arc<AtomicBool>,
+            }
+
+            #[derive(Clone)]
+            struct MyAlloc {
+                _inner: Arc<MyAllocInner>,
+            }
+
+            impl Drop for MyAllocInner {
+                fn drop(&mut self) {
+                    println!("MyAlloc freed.");
+                    self.drop_flag.store(true, Ordering::SeqCst);
+                }
+            }
+
+            unsafe impl Allocator for MyAlloc {
+                fn allocate(
+                    &self,
+                    layout: Layout,
+                ) -> std::result::Result<NonNull<[u8]>, AllocError> {
+                    let g = Global;
+                    g.allocate(layout)
+                }
+
+                unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+                    let g = Global;
+                    g.deallocate(ptr, layout)
+                }
+            }
+
+            let mut map = crate::HashMap::with_capacity_in(
+                10,
+                MyAlloc {
+                    _inner: Arc::new(MyAllocInner {
+                        drop_flag: dropped.clone(),
+                    }),
+                },
+            );
+            for i in 0..10 {
+                map.entry(i).or_insert_with(|| "i".to_string());
+            }
+
+            for (k, v) in map {
+                println!("{}, {}", k, v);
+            }
+        }
+
+        assert!(dropped.load(Ordering::SeqCst));
+    }
 }

src/raw/mod.rs

@@ -1482,20 +1482,18 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
     pub unsafe fn into_iter_from(self, iter: RawIter<T>) -> RawIntoIter<T, A> {
         debug_assert_eq!(iter.len(), self.len());
 
-        let alloc = self.table.alloc.clone();
         let allocation = self.into_allocation();
         RawIntoIter {
             iter,
             allocation,
             marker: PhantomData,
-            alloc,
         }
     }
 
     /// Converts the table into a raw allocation. The contents of the table
     /// should be dropped using a `RawIter` before freeing the allocation.
     #[cfg_attr(feature = "inline-more", inline)]
-    pub(crate) fn into_allocation(self) -> Option<(NonNull<u8>, Layout)> {
+    pub(crate) fn into_allocation(self) -> Option<(NonNull<u8>, Layout, A)> {
         let alloc = if self.table.is_empty_singleton() {
             None
         } else {
@@ -1508,6 +1506,7 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
             Some((
                 unsafe { NonNull::new_unchecked(self.table.ctrl.as_ptr().sub(ctrl_offset)) },
                 layout,
+                unsafe { ptr::read(&self.table.alloc) },
             ))
         };
         mem::forget(self);
@@ -3069,9 +3068,8 @@ impl<T> FusedIterator for RawIter<T> {}
 /// Iterator which consumes a table and returns elements.
 pub struct RawIntoIter<T, A: Allocator + Clone = Global> {
     iter: RawIter<T>,
-    allocation: Option<(NonNull<u8>, Layout)>,
+    allocation: Option<(NonNull<u8>, Layout, A)>,
     marker: PhantomData<T>,
-    alloc: A,
 }
 
 impl<T, A: Allocator + Clone> RawIntoIter<T, A> {
@@ -3103,8 +3101,8 @@ unsafe impl<#[may_dangle] T, A: Allocator + Clone> Drop for RawIntoIter<T, A> {
             self.iter.drop_elements();
 
             // Free the table
-            if let Some((ptr, layout)) = self.allocation {
-                self.alloc.deallocate(ptr, layout);
+            if let Some((ptr, layout, ref alloc)) = self.allocation {
+                alloc.deallocate(ptr, layout);
             }
         }
     }
@@ -3118,8 +3116,8 @@ impl<T, A: Allocator + Clone> Drop for RawIntoIter<T, A> {
             self.iter.drop_elements();
 
             // Free the table
-            if let Some((ptr, layout)) = self.allocation {
-                self.alloc.deallocate(ptr, layout);
+            if let Some((ptr, layout, ref alloc)) = self.allocation {
+                alloc.deallocate(ptr, layout);
             }
         }
     }