Implement clone_from for BTree collections

2 files changed, 66 insertions, 1 deletions

diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs
index e25a5e0773e..12174ffcbfa 100644
--- a/src/liballoc/collections/btree/map.rs
+++ b/src/liballoc/collections/btree/map.rs
@@ -207,6 +207,60 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
             clone_subtree(self.root.as_ref())
         }
     }
+
+    fn clone_from(&mut self, other: &Self) {
+        BTreeClone::clone_from(self, other);
+    }
+}
+
+trait BTreeClone {
+    fn clone_from(&mut self, other: &Self);
+}
+
+impl<K: Clone, V: Clone> BTreeClone for BTreeMap<K, V> {
+    default fn clone_from(&mut self, other: &Self) {
+        *self = other.clone();
+    }
+}
+
+impl<K: Clone + Ord, V: Clone> BTreeClone for BTreeMap<K, V> {
+    fn clone_from(&mut self, other: &Self) {
+        // This truncates `self` to `other.len()` by calling `split_off` on
+        // the first key after `other.len()` elements if it exists
+        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();
+        // After truncation, `self` is at most as long as `other` so this loop
+        // replaces every key-value pair in `self`. Since `oiter` is in sorted
+        // order and the structure of the `BTreeMap` stays the same,
+        // the BTree invariants are maintained at the end of the loop
+        while siter.front != siter.back {
+            if let Some((ok, ov)) = oiter.next() {
+                // This is safe because the `siter.front != siter.back` check
+                // ensures that `siter` is nonempty
+                let (sk, sv) = unsafe { siter.next_unchecked() };
+                sk.clone_from(ok);
+                sv.clone_from(ov);
+            } else {
+                break;
+            }
+        }
+        // If `other` is longer than `self`, the remaining elements are inserted
+        self.extend(oiter.map(|(k, v)| ((*k).clone(), (*v).clone())));
+    }
 }
 
 impl<K, Q: ?Sized> super::Recover<Q> for BTreeMap<K, ()>
diff --git a/src/liballoc/collections/btree/set.rs b/src/liballoc/collections/btree/set.rs
index 282d163141b..5bdefe5cecf 100644
--- a/src/liballoc/collections/btree/set.rs
+++ b/src/liballoc/collections/btree/set.rs
@@ -56,12 +56,23 @@ use crate::collections::btree_map::{self, BTreeMap, Keys};
 ///     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> 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`].