Auto merge of #66648 - crgl:btree-clone-from, r=<try>

Implement clone_from for BTreeMap and BTreeSet

See #28481. This results in up to 90% speedups on simple data types when `self` and `other` are the same size, and is generally comparable or faster. Some concerns:

1. This implementation requires an `Ord` bound on the `Clone` implementation for `BTreeMap` and `BTreeSet`. Since these structs can only be created externally for keys with `Ord` implemented, this should be fine? If not, there's certainly a less safe way to do this.
2. Changing `next_unchecked` on `RangeMut` to return mutable key references allows for replacing the entire overlapping portion of both maps without changing the external interface in any way. However, if `clone_from` fails it can leave the `BTreeMap` in an invalid state, which might be unacceptable.

This probably needs an FCP since it changes a trait bound, but (as far as I know?) that change cannot break any external code.

Author: bors

src/liballoc/collections/btree/map.rs

@@ -135,7 +135,7 @@ unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap<K, V> {
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
-impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
+impl<K: Clone + Ord, V: Clone> Clone for BTreeMap<K, V> {
     fn clone(&self) -> BTreeMap<K, V> {
         fn clone_subtree<'a, K: Clone, V: Clone>(
             node: node::NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>
@@ -201,16 +201,40 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
         }
 
         if self.is_empty() {
-            // Ideally we'd call `BTreeMap::new` here, but that has the `K:
-            // Ord` constraint, which this method lacks.
-            BTreeMap {
-                root: node::Root::shared_empty_root(),
-                length: 0,
-            }
+            BTreeMap::new()
         } else {
             clone_subtree(self.root.as_ref())
         }
     }
+
+    fn clone_from(&mut self, other: &Self) {
+        if let Some(key) = {
+            if self.len() > other.len() {
+                let diff = self.len() - other.len();
+                if diff <= other.len() {
+                    self.iter().nth_back(diff - 1).map(|pair| (*pair.0).clone())
+                } else {
+                    self.iter().nth(other.len()).map(|pair| (*pair.0).clone())
+                }
+            } else {
+                None
+            }
+        } {
+            self.split_off(&key);
+        }
+        let mut siter = self.range_mut(..);
+        let mut oiter = other.iter();
+        while siter.front != siter.back {
+            if let Some((ok, ov)) = oiter.next() {
+                let (sk, sv) = unsafe { siter.next_unchecked() };
+                sk.clone_from(ok);
+                sv.clone_from(ov);
+            } else {
+                break ;
+            }
+        }
+        self.extend(oiter.map(|(k, v)| ((*k).clone(), (*v).clone())));
+    }
 }
 
 impl<K, Q: ?Sized> super::Recover<Q> for BTreeMap<K, ()>
@@ -1364,7 +1388,10 @@ impl<'a, K: 'a, V: 'a> Iterator for IterMut<'a, K, V> {
             None
         } else {
             self.length -= 1;
-            unsafe { Some(self.range.next_unchecked()) }
+            unsafe {
+                let (k, v) = self.range.next_unchecked();
+                Some((k, v)) // coerce k from `&mut K` to `&K`
+            }
         }
     }
 
@@ -1761,7 +1788,10 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
         if self.front == self.back {
             None
         } else {
-            unsafe { Some(self.next_unchecked()) }
+            unsafe {
+                let (k, v) = self.next_unchecked();
+                Some((k, v)) // coerce k from `&mut K` to `&K`
+            }
         }
     }
 
@@ -1771,16 +1801,15 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> {
 }
 
 impl<'a, K, V> RangeMut<'a, K, V> {
-    unsafe fn next_unchecked(&mut self) -> (&'a K, &'a mut V) {
+    unsafe fn next_unchecked(&mut self) -> (&'a mut K, &'a mut V) {
         let handle = ptr::read(&self.front);
 
         let mut cur_handle = match handle.right_kv() {
             Ok(kv) => {
                 self.front = ptr::read(&kv).right_edge();
                 // Doing the descend invalidates the references returned by `into_kv_mut`,
                 // so we have to do this last.
-                let (k, v) = kv.into_kv_mut();
-                return (k, v); // coerce k from `&mut K` to `&K`
+                return kv.into_kv_mut();
             }
             Err(last_edge) => {
                 let next_level = last_edge.into_node().ascend().ok();
@@ -1794,8 +1823,7 @@ impl<'a, K, V> RangeMut<'a, K, V> {
                     self.front = first_leaf_edge(ptr::read(&kv).right_edge().descend());
                     // Doing the descend invalidates the references returned by `into_kv_mut`,
                     // so we have to do this last.
-                    let (k, v) = kv.into_kv_mut();
-                    return (k, v); // coerce k from `&mut K` to `&K`
+                    return kv.into_kv_mut();
                 }
                 Err(last_edge) => {
                     let next_level = last_edge.into_node().ascend().ok();

src/liballoc/collections/btree/set.rs

@@ -56,12 +56,23 @@ use super::Recover;
 ///     println!("{}", book);
 /// }
 /// ```
-#[derive(Clone, Hash, PartialEq, Eq, Ord, PartialOrd)]
+#[derive(Hash, PartialEq, Eq, Ord, PartialOrd)]
 #[stable(feature = "rust1", since = "1.0.0")]
 pub struct BTreeSet<T> {
     map: BTreeMap<T, ()>,
 }
 
+#[stable(feature = "rust1", since = "1.0.0")]
+impl<T: Clone + Ord> Clone for BTreeSet<T> {
+    fn clone(&self) -> Self {
+        BTreeSet { map: self.map.clone() }
+    }
+
+    fn clone_from(&mut self, other: &Self) {
+        self.map.clone_from(&other.map);
+    }
+}
+
 /// An iterator over the items of a `BTreeSet`.
 ///
 /// This `struct` is created by the [`iter`] method on [`BTreeSet`].

src/liballoc/tests/btree/map.rs

@@ -552,6 +552,26 @@ fn test_clone() {
     }
 }
 
+#[test]
+fn test_clone_from() {
+    let mut map1 = BTreeMap::new();
+    let size = 30;
+
+    for i in 0..size {
+        map1.insert(i, 10 * i);
+        let mut map2 = BTreeMap::new();
+        for j in 0..i {
+            map2.insert(100 * j + 1, 2 * j + 1);
+            let mut map1_copy = map2.clone();
+            map1_copy.clone_from(&map1);
+            assert_eq!(map1_copy, map1);
+            let mut map2_copy = map1.clone();
+            map2_copy.clone_from(&map2);
+            assert_eq!(map2_copy, map2);
+        }
+    }
+}
+
 #[test]
 #[allow(dead_code)]
 fn test_variance() {