diff --git a/src/liballoc/collections/btree/map.rs b/src/liballoc/collections/btree/map.rs index 4dc004864fd46..11c81ceb19430 100644 --- a/src/liballoc/collections/btree/map.rs +++ b/src/liballoc/collections/btree/map.rs @@ -207,6 +207,60 @@ impl Clone for BTreeMap { 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 BTreeClone for BTreeMap { + default fn clone_from(&mut self, other: &Self) { + *self = other.clone(); + } +} + +impl BTreeClone for BTreeMap { + 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 + let split_off_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 + }; + if let Some(key) = split_off_key { + 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.is_empty() { + if let Some((ok, ov)) = oiter.next() { + // SAFETY: 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 super::Recover for BTreeMap @@ -1357,7 +1411,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` + } } } @@ -1736,7 +1793,14 @@ impl<'a, K, V> Iterator for RangeMut<'a, K, V> { type Item = (&'a K, &'a mut V); fn next(&mut self) -> Option<(&'a K, &'a mut V)> { - if self.front == self.back { None } else { unsafe { Some(self.next_unchecked()) } } + if self.is_empty() { + None + } else { + unsafe { + let (k, v) = self.next_unchecked(); + Some((k, v)) // coerce k from `&mut K` to `&K` + } + } } fn last(mut self) -> Option<(&'a K, &'a mut V)> { @@ -1745,7 +1809,11 @@ 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) { + fn is_empty(&self) -> bool { + self.front == self.back + } + + 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() { @@ -1753,8 +1821,7 @@ impl<'a, K, V> RangeMut<'a, K, V> { 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(); @@ -1768,8 +1835,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(); @@ -1783,7 +1849,7 @@ impl<'a, K, V> RangeMut<'a, K, V> { #[stable(feature = "btree_range", since = "1.17.0")] impl<'a, K, V> DoubleEndedIterator for RangeMut<'a, K, V> { fn next_back(&mut self) -> Option<(&'a K, &'a mut V)> { - if self.front == self.back { None } else { unsafe { Some(self.next_back_unchecked()) } } + if self.is_empty() { None } else { unsafe { Some(self.next_back_unchecked()) } } } } diff --git a/src/liballoc/collections/btree/set.rs b/src/liballoc/collections/btree/set.rs index f5487426814a5..b100ce754caad 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 { map: BTreeMap, } +#[stable(feature = "rust1", since = "1.0.0")] +impl Clone for BTreeSet { + 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`]. diff --git a/src/liballoc/tests/btree/map.rs b/src/liballoc/tests/btree/map.rs index f5be72c39b20c..0d009507fc7aa 100644 --- a/src/liballoc/tests/btree/map.rs +++ b/src/liballoc/tests/btree/map.rs @@ -785,6 +785,26 @@ fn test_clone() { } } +#[test] +fn test_clone_from() { + let mut map1 = BTreeMap::new(); + let size = 30; + + for i in 0..size { + let mut map2 = BTreeMap::new(); + for j in 0..i { + 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); + map2.insert(100 * j + 1, 2 * j + 1); + } + map1.insert(i, 10 * i); + } +} + #[test] #[allow(dead_code)] fn test_variance() {