Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

BTree: refactor clone, mostly for readability #89513

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 54 additions & 0 deletions library/alloc/src/collections/btree/bare_tree.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
use super::node::{marker, NodeRef};

/// Simple tree around a raw `NodeRef`, that provides a destructor.
/// Unlike smarter sister `BTreeMap`, it does not keep track of element count,
/// does not allow replacing the root, therefore can be statically tyoed by the
/// node type of the root, and cannot represent a tree without root node
/// (its `Option` exists only to disable the `drop` method when needed).
pub struct BareTree<K, V, RootType>(Option<NodeRef<marker::Owned, K, V, RootType>>);

impl<K, V, RootType> Drop for BareTree<K, V, RootType> {
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() } {
unsafe { kv.drop_key_val() };
cur_edge = next_edge;
}
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like this doesn't replicate the continue-dropping-if-we-see-a-panic code, presumably because the currently only user only drops BareTrees if we're already panicking. But that seems like a footgun should we start using it somewhere that does need to keep dropping.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

presumably because

Actually because I didn't realize it belongs there.

}
}

impl<K, V> BareTree<K, V, marker::Leaf> {
/// Returns a new tree consisting of a single leaf that is initially empty.
pub fn new_leaf() -> Self {
Self(Some(NodeRef::new_leaf()))
}
}

impl<K, V> BareTree<K, V, marker::Internal> {
/// Returns a new tree with an internal root, that initially has no elements
/// and one child.
pub fn new_internal(child: NodeRef<marker::Owned, K, V, marker::LeafOrInternal>) -> Self {
Self(Some(NodeRef::new_internal(child)))
}
}

impl<K, V, RootType> BareTree<K, V, RootType> {
/// Mutably borrows the owned root node.
pub fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, RootType> {
self.0.as_mut().unwrap().borrow_mut()
}

/// Removes any static information asserting that the root is a `Leaf` or
/// `Internal` node.
pub fn forget_root_type(mut self) -> BareTree<K, V, marker::LeafOrInternal> {
BareTree(self.0.take().map(NodeRef::forget_type))
}

/// Consumes the tree, returning a `NodeRef` that still enforces ownership
/// but lacks a destructor.
pub fn into_inner(mut self) -> NodeRef<marker::Owned, K, V, RootType> {
self.0.take().unwrap()
}
}
55 changes: 14 additions & 41 deletions library/alloc/src/collections/btree/map.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use core::mem::{self, ManuallyDrop};
use core::ops::{Index, RangeBounds};
use core::ptr;

use super::bare_tree::BareTree;
use super::borrow::DormantMutRef;
use super::dedup_sorted_iter::DedupSortedIter;
use super::navigate::{LazyLeafRange, LeafRange};
Expand Down Expand Up @@ -171,75 +172,47 @@ 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> {
fn clone(&self) -> BTreeMap<K, V> {
fn clone_subtree<'a, K: Clone, V: Clone>(
fn clone_subtree<'a, K: Clone + 'a, V: Clone + 'a>(
node: NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>,
) -> BTreeMap<K, V>
where
K: 'a,
V: 'a,
{
) -> BareTree<K, V, marker::LeafOrInternal> {
match node.force() {
Leaf(leaf) => {
let mut out_tree = BTreeMap { root: Some(Root::new()), length: 0 };

let mut out_tree = BareTree::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.forget_root_type()
}
Internal(internal) => {
let mut out_tree = clone_subtree(internal.first_edge().descend());

let mut in_edge = internal.first_edge();
let first_child = clone_subtree(in_edge.descend());
let mut out_tree = BareTree::new_internal(first_child.into_inner());
{
let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
let mut out_node = out_root.push_internal_level();
let mut in_edge = internal.first_edge();
let mut out_node = out_tree.borrow_mut();
while let Ok(kv) = in_edge.right_kv() {
let (k, v) = kv.into_kv();
in_edge = kv.right_edge();

let k = (*k).clone();
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.into_inner());
}
}

out_tree
out_tree.forget_root_type()
}
}
}

if self.is_empty() {
BTreeMap::new()
} else {
clone_subtree(self.root.as_ref().unwrap().reborrow()) // unwrap succeeds because not empty
}
let cloned_root = self.root.as_ref().map(|r| clone_subtree(r.reborrow()).into_inner());
BTreeMap { root: cloned_root, length: self.length }
}
}

Expand Down
1 change: 1 addition & 0 deletions library/alloc/src/collections/btree/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
mod append;
mod bare_tree;
mod borrow;
mod dedup_sorted_iter;
mod fix;
Expand Down
2 changes: 1 addition & 1 deletion library/alloc/src/collections/btree/navigate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -441,7 +441,7 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
/// `deallocating_next_back`.
/// - The returned KV handle is only valid to access the key and value,
/// and only valid until the next call to a `deallocating_` method.
unsafe fn deallocating_next(
pub unsafe fn deallocating_next(
self,
) -> Option<(Self, Handle<NodeRef<marker::Dying, K, V, marker::LeafOrInternal>, marker::KV>)>
{
Expand Down
37 changes: 8 additions & 29 deletions library/alloc/src/collections/btree/node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,8 @@ unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Owned, K, V, Type>
unsafe impl<K: Send, V: Send, Type> Send for NodeRef<marker::Dying, K, V, Type> {}

impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
fn new_leaf() -> Self {
/// Returns a new owned leaf node, that is initially empty.
pub fn new_leaf() -> Self {
Self::from_new_leaf(LeafNode::new())
}

Expand All @@ -223,7 +224,8 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
}

impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
fn new_internal(child: Root<K, V>) -> Self {
/// Returns a new owned internal node, that initially has no elements and one child.
pub fn new_internal(child: Root<K, V>) -> Self {
let mut new_node = unsafe { InternalNode::new() };
new_node.edges[0].write(child.node);
unsafe { NodeRef::from_new_internal(new_node, child.height + 1) }
Expand Down Expand Up @@ -651,15 +653,8 @@ impl<'a, K: 'a, V: 'a> NodeRef<marker::Mut<'a>, K, V, marker::Internal> {
}
}

impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Leaf> {
/// Removes any static information asserting that this node is a `Leaf` node.
pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}
}

impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
/// Removes any static information asserting that this node is an `Internal` node.
impl<BorrowType, K, V, NodeType> NodeRef<BorrowType, K, V, NodeType> {
/// Removes any static information asserting that this node is a `Leaf` or `Internal` node.
pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
}
Expand Down Expand Up @@ -1505,31 +1500,15 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
pub fn forget_node_type(
self,
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
pub fn forget_node_type(
self,
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::KV> {
pub fn forget_node_type(
self,
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
unsafe { Handle::new_kv(self.node.forget_type(), self.idx) }
}
}

impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV> {
impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV> {
pub fn forget_node_type(
self,
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
Expand Down