Skip to content

Commit 35c7567

Browse files
ssomersssomers
committed
BTreeMap: clean up implementation of clone
1 parent 5c56f40 commit 35c7567

File tree

3 files changed

+41
-62
lines changed

3 files changed

+41
-62
lines changed

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

+34-32
Original file line numberDiff line numberDiff line change
@@ -147,42 +147,59 @@ unsafe impl<#[may_dangle] K, #[may_dangle] V> Drop for BTreeMap<K, V> {
147147
#[stable(feature = "rust1", since = "1.0.0")]
148148
impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
149149
fn clone(&self) -> BTreeMap<K, V> {
150+
struct GuardedTree<K, V, Type>(Option<NodeRef<marker::Owned, K, V, Type>>);
151+
impl<K, V, Type> GuardedTree<K, V, Type> {
152+
fn new(root: NodeRef<marker::Owned, K, V, Type>) -> Self {
153+
Self(Some(root))
154+
}
155+
fn borrow_mut(&mut self) -> NodeRef<marker::Mut<'_>, K, V, Type> {
156+
self.0.as_mut().unwrap().borrow_mut()
157+
}
158+
fn leak(mut self) -> Root<K, V> {
159+
self.0.take().unwrap().forget_type()
160+
}
161+
}
162+
impl<K, V, Type> Drop for GuardedTree<K, V, Type> {
163+
fn drop(&mut self) {
164+
if let Some(root) = self.0.take() {
165+
let mut cur_edge = root.forget_type().into_dying().first_leaf_edge();
166+
while let Some((next_edge, _kv)) = unsafe { cur_edge.deallocating_next() } {
167+
cur_edge = next_edge;
168+
}
169+
}
170+
}
171+
}
172+
150173
fn clone_subtree<'a, K: Clone, V: Clone>(
151174
node: NodeRef<marker::Immut<'a>, K, V, marker::LeafOrInternal>,
152-
) -> BTreeMap<K, V>
175+
) -> Root<K, V>
153176
where
154177
K: 'a,
155178
V: 'a,
156179
{
157180
match node.force() {
158181
Leaf(leaf) => {
159-
let mut out_tree = BTreeMap { root: Some(Root::new()), length: 0 };
182+
let mut out_tree = GuardedTree::new(NodeRef::new_leaf());
160183

161184
{
162-
let root = out_tree.root.as_mut().unwrap(); // unwrap succeeds because we just wrapped
163-
let mut out_node = match root.borrow_mut().force() {
164-
Leaf(leaf) => leaf,
165-
Internal(_) => unreachable!(),
166-
};
167-
185+
let mut out_node = out_tree.borrow_mut();
168186
let mut in_edge = leaf.first_edge();
169187
while let Ok(kv) = in_edge.right_kv() {
170188
let (k, v) = kv.into_kv();
171189
in_edge = kv.right_edge();
172190

173191
out_node.push(k.clone(), v.clone());
174-
out_tree.length += 1;
175192
}
176193
}
177194

178-
out_tree
195+
out_tree.leak()
179196
}
180197
Internal(internal) => {
181-
let mut out_tree = clone_subtree(internal.first_edge().descend());
198+
let first_child = clone_subtree(internal.first_edge().descend());
199+
let mut out_tree = GuardedTree::new(NodeRef::new_internal(first_child));
182200

183201
{
184-
let out_root = BTreeMap::ensure_is_owned(&mut out_tree.root);
185-
let mut out_node = out_root.push_internal_level();
202+
let mut out_node = out_tree.borrow_mut();
186203
let mut in_edge = internal.first_edge();
187204
while let Ok(kv) = in_edge.right_kv() {
188205
let (k, v) = kv.into_kv();
@@ -192,32 +209,17 @@ impl<K: Clone, V: Clone> Clone for BTreeMap<K, V> {
192209
let v = (*v).clone();
193210
let subtree = clone_subtree(in_edge.descend());
194211

195-
// We can't destructure subtree directly
196-
// because BTreeMap implements Drop
197-
let (subroot, sublength) = unsafe {
198-
let subtree = ManuallyDrop::new(subtree);
199-
let root = ptr::read(&subtree.root);
200-
let length = subtree.length;
201-
(root, length)
202-
};
203-
204-
out_node.push(k, v, subroot.unwrap_or_else(Root::new));
205-
out_tree.length += 1 + sublength;
212+
out_node.push(k, v, subtree);
206213
}
207214
}
208215

209-
out_tree
216+
out_tree.leak()
210217
}
211218
}
212219
}
213220

214-
if self.is_empty() {
215-
// Ideally we'd call `BTreeMap::new` here, but that has the `K:
216-
// Ord` constraint, which this method lacks.
217-
BTreeMap { root: None, length: 0 }
218-
} else {
219-
clone_subtree(self.root.as_ref().unwrap().reborrow()) // unwrap succeeds because not empty
220-
}
221+
let cloned_root = self.root.as_ref().map(|root| clone_subtree(root.reborrow()));
222+
BTreeMap { root: cloned_root, length: self.length }
221223
}
222224
}
223225

‎library/alloc/src/collections/btree/navigate.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -301,7 +301,7 @@ impl<K, V> Handle<NodeRef<marker::Dying, K, V, marker::Leaf>, marker::Edge> {
301301
///
302302
/// # Safety
303303
/// The next KV must not have been previously returned by counterpart `deallocating_next_back`.
304-
unsafe fn deallocating_next(self) -> Option<(Self, (K, V))> {
304+
pub unsafe fn deallocating_next(self) -> Option<(Self, (K, V))> {
305305
let mut edge = self.forget_node_type();
306306
loop {
307307
edge = match edge.right_kv() {

‎library/alloc/src/collections/btree/node.rs

+6-29
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ impl<K, V> Root<K, V> {
132132
}
133133

134134
impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
135-
fn new_leaf() -> Self {
135+
pub fn new_leaf() -> Self {
136136
Self::from_new_leaf(Box::new(unsafe { LeafNode::new() }))
137137
}
138138

@@ -142,7 +142,7 @@ impl<K, V> NodeRef<marker::Owned, K, V, marker::Leaf> {
142142
}
143143

144144
impl<K, V> NodeRef<marker::Owned, K, V, marker::Internal> {
145-
fn new_internal(child: Root<K, V>) -> Self {
145+
pub fn new_internal(child: Root<K, V>) -> Self {
146146
let mut new_node = Box::new(unsafe { InternalNode::new() });
147147
new_node.edges[0].write(child.node);
148148
NodeRef::from_new_internal(new_node, child.height + 1)
@@ -1461,45 +1461,22 @@ impl<'a, K: 'a, V: 'a> BalancingContext<'a, K, V> {
14611461
}
14621462
}
14631463

1464-
impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Leaf> {
1465-
/// Removes any static information asserting that this node is a `Leaf` node.
1466-
pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
1467-
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
1468-
}
1469-
}
1470-
1471-
impl<BorrowType, K, V> NodeRef<BorrowType, K, V, marker::Internal> {
1472-
/// Removes any static information asserting that this node is an `Internal` node.
1464+
impl<BorrowType, K, V, Type> NodeRef<BorrowType, K, V, Type> {
1465+
/// Removes any static information asserting that this node is a `Leaf` or `Internal` node.
14731466
pub fn forget_type(self) -> NodeRef<BorrowType, K, V, marker::LeafOrInternal> {
14741467
NodeRef { height: self.height, node: self.node, _marker: PhantomData }
14751468
}
14761469
}
14771470

1478-
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::Edge> {
1479-
pub fn forget_node_type(
1480-
self,
1481-
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
1482-
unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
1483-
}
1484-
}
1485-
1486-
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::Edge> {
1471+
impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::Edge> {
14871472
pub fn forget_node_type(
14881473
self,
14891474
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::Edge> {
14901475
unsafe { Handle::new_edge(self.node.forget_type(), self.idx) }
14911476
}
14921477
}
14931478

1494-
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Leaf>, marker::KV> {
1495-
pub fn forget_node_type(
1496-
self,
1497-
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {
1498-
unsafe { Handle::new_kv(self.node.forget_type(), self.idx) }
1499-
}
1500-
}
1501-
1502-
impl<BorrowType, K, V> Handle<NodeRef<BorrowType, K, V, marker::Internal>, marker::KV> {
1479+
impl<BorrowType, K, V, NodeType> Handle<NodeRef<BorrowType, K, V, NodeType>, marker::KV> {
15031480
pub fn forget_node_type(
15041481
self,
15051482
) -> Handle<NodeRef<BorrowType, K, V, marker::LeafOrInternal>, marker::KV> {

0 commit comments

Comments
 (0)