Implement clone_from for HashMap

Author: Amanieu

src/map.rs

@@ -188,12 +188,29 @@ pub enum DefaultHashBuilder {}
 ///  .iter().cloned().collect();
 /// // use the values stored in map
 /// ```
-#[derive(Clone)]
 pub struct HashMap<K, V, S = DefaultHashBuilder> {
     pub(crate) hash_builder: S,
     pub(crate) table: RawTable<(K, V)>,
 }
 
+impl<K: Clone, V: Clone, S: Clone> Clone for HashMap<K, V, S> {
+    fn clone(&self) -> Self {
+        HashMap {
+            hash_builder: self.hash_builder.clone(),
+            table: self.table.clone(),
+        }
+    }
+
+    fn clone_from(&mut self, source: &Self) {
+        // We clone the hash_builder first since this might panic and we don't
+        // want the table to have elements hashed with the wrong hash_builder.
+        let hash_builder = source.hash_builder.clone();
+
+        self.table.clone_from(&source.table);
+        self.hash_builder = hash_builder;
+    }
+}
+
 #[cfg_attr(feature = "inline-more", inline)]
 pub(crate) fn make_hash<K: Hash + ?Sized>(hash_builder: &impl BuildHasher, val: &K) -> u64 {
     let mut state = hash_builder.build_hasher();
@@ -2820,6 +2837,21 @@ mod test_map {
         assert_eq!(m2.len(), 2);
     }
 
+    #[test]
+    fn test_clone_from() {
+        let mut m = HashMap::new();
+        let mut m2 = HashMap::new();
+        assert_eq!(m.len(), 0);
+        assert!(m.insert(1, 2).is_none());
+        assert_eq!(m.len(), 1);
+        assert!(m.insert(2, 4).is_none());
+        assert_eq!(m.len(), 2);
+        m2.clone_from(&m);
+        assert_eq!(*m2.get(&1).unwrap(), 2);
+        assert_eq!(*m2.get(&2).unwrap(), 4);
+        assert_eq!(m2.len(), 2);
+    }
+
     thread_local! { static DROP_VECTOR: RefCell<Vec<i32>> = RefCell::new(Vec::new()) }
 
     #[derive(Hash, PartialEq, Eq)]

src/raw/mod.rs

@@ -987,47 +987,92 @@ impl<T: Clone> Clone for RawTable<T> {
                         .unwrap_or_else(|_| hint::unreachable_unchecked()),
                 );
 
-                // Copy the control bytes unchanged. We do this in a single pass
-                self.ctrl(0)
-                    .copy_to_nonoverlapping(new_table.ctrl(0), self.num_ctrl_bytes());
-
-                {
-                    // The cloning of elements may panic, in which case we need
-                    // to make sure we drop only the elements that have been
-                    // cloned so far.
-                    let mut guard = guard((0, &mut new_table), |(index, new_table)| {
-                        if mem::needs_drop::<T>() {
-                            for i in 0..=*index {
-                                if is_full(*new_table.ctrl(i)) {
-                                    new_table.bucket(i).drop();
-                                }
-                            }
-                        }
-                        new_table.free_buckets();
-                    });
+                new_table.clone_from_impl(self, |new_table| {
+                    // We need to free the memory allocated for the new table.
+                    new_table.free_buckets();
+                });
 
-                    for from in self.iter() {
-                        let index = self.bucket_index(&from);
-                        let to = guard.1.bucket(index);
-                        to.write(from.as_ref().clone());
+                // Return the newly created table.
+                ManuallyDrop::into_inner(new_table)
+            }
+        }
+    }
 
-                        // Update the index in case we need to unwind.
-                        guard.0 = index;
+    fn clone_from(&mut self, source: &Self) {
+        if source.is_empty_singleton() {
+            *self = Self::new();
+        } else {
+            unsafe {
+                // First, drop all our elements without clearing the control bytes.
+                if mem::needs_drop::<T>() {
+                    for item in self.iter() {
+                        item.drop();
                     }
+                }
 
-                    // Successfully cloned all items, no need to clean up.
-                    mem::forget(guard);
+                // If necessary, resize our table to match the source.
+                if self.buckets() != source.buckets() {
+                    // Skip our drop by using ptr::write.
+                    self.free_buckets();
+                    (self as *mut Self).write(
+                        Self::new_uninitialized(source.buckets(), Fallibility::Infallible)
+                            .unwrap_or_else(|_| hint::unreachable_unchecked()),
+                    );
                 }
 
-                // Return the newly created table.
-                new_table.items = self.items;
-                new_table.growth_left = self.growth_left;
-                ManuallyDrop::into_inner(new_table)
+                self.clone_from_impl(source, |self_| {
+                    // We need to leave the table in an empty state.
+                    self_.clear_no_drop()
+                });
             }
         }
     }
 }
 
+impl<T: Clone> RawTable<T> {
+    /// Common code for clone and clone_from. Assumes `self.buckets() == source.buckets()`.
+    #[cfg_attr(feature = "inline-more", inline)]
+    unsafe fn clone_from_impl(&mut self, source: &Self, mut on_panic: impl FnMut(&mut Self)) {
+        // Copy the control bytes unchanged. We do this in a single pass
+        source
+            .ctrl(0)
+            .copy_to_nonoverlapping(self.ctrl(0), self.num_ctrl_bytes());
+
+        // The cloning of elements may panic, in which case we need
+        // to make sure we drop only the elements that have been
+        // cloned so far.
+        let mut guard = guard((0, &mut *self), |(index, self_)| {
+            if mem::needs_drop::<T>() {
+                for i in 0..=*index {
+                    if is_full(*self_.ctrl(i)) {
+                        self_.bucket(i).drop();
+                    }
+                }
+            }
+
+            // Depending on whether we were called from clone or clone_from, we
+            // either need to free the memory for the destination table or just
+            // clear the control bytes.
+            on_panic(self_);
+        });
+
+        for from in source.iter() {
+            let index = source.bucket_index(&from);
+            let to = guard.1.bucket(index);
+            to.write(from.as_ref().clone());
+
+            // Update the index in case we need to unwind.
+            guard.0 = index;
+        }
+
+        // Successfully cloned all items, no need to clean up.
+        mem::forget(guard);
+
+        self.items = source.items;
+        self.growth_left = source.growth_left;
+    }
+}
+
 #[cfg(feature = "nightly")]
 unsafe impl<#[may_dangle] T> Drop for RawTable<T> {
     #[cfg_attr(feature = "inline-more", inline)]

src/set.rs

@@ -111,11 +111,22 @@ use super::map::{self, DefaultHashBuilder, HashMap, Keys};
 /// [`HashMap`]: struct.HashMap.html
 /// [`PartialEq`]: https://doc.rust-lang.org/std/cmp/trait.PartialEq.html
 /// [`RefCell`]: https://doc.rust-lang.org/std/cell/struct.RefCell.html
-#[derive(Clone)]
 pub struct HashSet<T, S = DefaultHashBuilder> {
     pub(crate) map: HashMap<T, (), S>,
 }
 
+impl<T: Clone, S: Clone> Clone for HashSet<T, S> {
+    fn clone(&self) -> Self {
+        HashSet {
+            map: self.map.clone(),
+        }
+    }
+
+    fn clone_from(&mut self, source: &Self) {
+        self.map.clone_from(&source.map);
+    }
+}
+
 #[cfg(feature = "ahash")]
 impl<T: Hash + Eq> HashSet<T, DefaultHashBuilder> {
     /// Creates an empty `HashSet`.