BTreeMap: clean up implementation of clone

library/alloc/src/collections/btree/map.rs

@@ -138,51 +138,68 @@ pub struct BTreeMap<K, V> {
 #[stable(feature = "btree_drop", since = "1.7.0")]
 unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap<K, V> {
  fn drop(&mut self) {
- unsafe {
- drop(ptr::read(self).into_iter());
+ if let Some(root) = self.root.take() {
+ Dropper { front: root.into_dying().first_leaf_edge(), remaining_length: self.length };
  }
  }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
  fn clone(&self) -> BTreeMap<K, V> {
+ struct GuardedTree<K, V, Type>(Option<NodeRef<marker::Owned, K, V, Type>>);
+ impl<K, V, Type> GuardedTree<K, V, Type> {
+ fn new(root: NodeRef<marker::Owned, K, V, Type>) -> Self {
+ Self(Some(root))
+ }
+ fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> {
+ self.0.as_mut().unwrap().borrow_mut()
+ }
+ fn leak(mut self) -> Root<K, V> {
+ self.0.take().unwrap().forget_type()
+ }
+ }
+ impl<K, V, Type> Drop for GuardedTree<K, V, Type> {
+ fn drop(&mut self) {
+ if let Some(root) = self.0.take() {
+ let mut cur_edge = root.forget_type().into_dying().first_leaf_edge();
+ while let Some((next_edge, _kv)) = unsafe { cur_edge.deallocating_next() } {
+ cur_edge = next_edge;
+ }
+ }
+ }
+ }
+
  fn clone_subtree<'a, K: Clone, V: Clone>(
  node: NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>,
- ) -> BTreeMap<K, V>
+ ) -> Root<K, V>
  where
  K: 'a,
  V: 'a,
  {
  match node.force() {
  Leaf(leaf) => {
- let mut out_tree = BTreeMap { root: Some(Root::new()), length: 0 };
+ let mut out_tree = GuardedTree::new(NodeRef::new_leaf());
 
  {
- let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped
- let mut out_node = match root.borrow_mut().force() {
- Leaf(leaf) => leaf,
- Internal(_) => unreachable!(),
- };
-
+ let mut out_node = out_tree.borrow_mut();
  let mut in_edge = leaf.first_edge();
  while let Ok(kv) = in_edge.right_kv() {
  let (k, v) = kv.into_kv();
  in_edge = kv.right_edge();
 
  out_node.push(k.clone(), v.clone());
- out_tree.length += 1;
  }
  }
 
- out_tree
+ out_tree.leak()
  }
  Internal(internal) => {
- let mut out_tree = clone_subtree(internal.first_edge().descend());
+ let first_child = clone_subtree(internal.first_edge().descend());
+ let mut out_tree = GuardedTree::new(NodeRef::new_internal(first_child));
 
  {
- let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
- let mut out_node = out_root.push_internal_level();
+ let mut out_node = out_tree.borrow_mut();
  let mut in_edge = internal.first_edge();
  while let Ok(kv) = in_edge.right_kv() {
  let (k, v) = kv.into_kv();
@@ -192,32 +209,17 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
  let v = (*v).clone();
  let subtree = clone_subtree(in_edge.descend());
 
- // We can't destructure subtree directly
- // because BTreeMap implements Drop
- let (subroot, sublength) = unsafe {
- let subtree = ManuallyDrop::new(subtree);
- let root = ptr::read(&subtree.root);
- let length = subtree.length;
- (root, length)
- };
-
- out_node.push(k, v, subroot.unwrap_or_else(Root::new));
- out_tree.length += 1 + sublength;
+ out_node.push(k, v, subtree);
  }
  }
 
- out_tree
+ out_tree.leak()
  }
  }
  }
 
- if self.is_empty() {
- // Ideally we'd call `BTreeMap::new` here, but that has the `K:
- // Ord` constraint, which this method lacks.
- BTreeMap { root: None, length: 0 }
- } else {
- clone_subtree(self.root.as_ref().unwrap().reborrow()) // unwrap succeeds because not empty
- }
+ let cloned_root = self.root.as_ref().map(|root| clone_subtree(root.reborrow()));
+ BTreeMap { root: cloned_root, length: self.length }
  }
 }
 
@@ -325,6 +327,14 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IntoIter<K, V> {
  }
 }
 
+/// A simplified version of `IntoIter` that is not double-ended and has only one
+/// purpose: to drop the remainder of an `IntoIter`. Therefore it also serves to
+/// drop an entire tree without the need to look up a `back` leaf edge.
+struct Dropper<K, V> {
+ front: Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge>,
+ remaining_length: usize,
+}
+
 /// An iterator over the keys of a `BTreeMap`.
 ///
 /// This `struct` is created by the [`keys`] method on [`BTreeMap`]. See its
@@ -1373,42 +1383,42 @@ impl<K, V> IntoIterator for BTreeMap<K, V> {
  }
 }
 
-#[stable(feature = "btree_drop", since = "1.7.0")]
-impl<K, V> Drop for IntoIter<K, V> {
+impl<K, V> Drop for Dropper<K, V> {
  fn drop(&mut self) {
- struct DropGuard<'a, K, V>(&'a mut IntoIter<K, V>);
+ // Similar to advancing a non-fusing iterator.
+ fn next_or_end<K, V>(this: &mut Dropper<K, V>) -> Option<(K, V)> {
+ if this.remaining_length == 0 {
+ unsafe { ptr::read(&this.front).deallocating_end() }
+ None
+ } else {
+ this.remaining_length -= 1;
+ Some(unsafe { this.front.next_unchecked() })
+ }
+ }
+
+ struct DropGuard<'a, K, V>(&'a mut Dropper<K, V>);
 
  impl<'a, K, V> Drop for DropGuard<'a, K, V> {
  fn drop(&mut self) {
  // Continue the same loop we perform below. This only runs when unwinding, so we
  // don't have to care about panics this time (they'll abort).
- while let Some(_) = self.0.next() {}
-
- unsafe {
- let mut node =
- unwrap_unchecked(ptr::read(&self.0.front)).into_node().forget_type();
- while let Some(parent) = node.deallocate_and_ascend() {
- node = parent.into_node().forget_type();
- }
- }
+ while let Some(_pair) = next_or_end(&mut self.0) {}
  }
  }
 
- while let Some(pair) = self.next() {
+ while let Some(pair) = next_or_end(self) {
  let guard = DropGuard(self);
  drop(pair);
  mem::forget(guard);
  }
+ }
+}
 
- unsafe {
- if let Some(front) = ptr::read(&self.front) {
- let mut node = front.into_node().forget_type();
- // Most of the nodes have been deallocated while traversing
- // but one pile from a leaf up to the root is left standing.
- while let Some(parent) = node.deallocate_and_ascend() {
- node = parent.into_node().forget_type();
- }
- }
+#[stable(feature = "btree_drop", since = "1.7.0")]
+impl<K, V> Drop for IntoIter<K, V> {
+ fn drop(&mut self) {
+ if let Some(front) = self.front.take() {
+ Dropper { front, remaining_length: self.length };
  }
  }
 }