BTreeMap::drain_filter should not touch the root during iteration

[ssomers]

Although Miri doesn't point it out, I believe there is undefined behaviour using drain_filter when draining the 11th-last element from a tree that was larger. When this happens, the last remaining child nodes are merged, the root becomes empty and is popped from the tree. That last step establishes a mutable reference to the node elected root and writes a pointer in node::Root, while iteration continues to visit the same node.

This is mostly code from #74437, slightly adapted.

[rust-highfive]

r? @Mark-Simulacrum

(rust_highfive has picked a reviewer for you, use r? to override)

[bors]

☔ The latest upstream changes (presumably #73265) made this pull request unmergeable. Please resolve the merge conflicts.

[Mark-Simulacrum]

I am still not certain that the UB is there, but it seems like the related code cleanups in this PR are good regardless.

@bors r+

[bors]

📌 Commit 99398dd has been approved by Mark-Simulacrum

[ssomers]

You're probably right. I assumed that pop_level would deference the new root node (in addition to the Root owner of that node) but it seems like it doesn't. It all works on raw pointers and NonNull. There's also the NodeRef::root raw pointer that stays around in the iterator's copy of the handle, that I think we couldn't safely deference a second time, but that never happened anyway because it's only needed to pop_level the empty root and that can only happen once in a drain filter's lifetime.

Note that the earlier commit simplified code a little more, but recent benchmarking said consistently this costs 5 to 10%, so I passed the emptied_internal_root information around instead. Not sure what I benchmarked earlier.

[ssomers]

Illustration of the alleged undefined behaviour under stacked borrows rules:

Say we have a 2-level map with 11 keys -5..=5. The root has key 0, a left child with keys -5..=-1, a right child with keys 1..=5. You won't get this 2nd level inserting 11 elements from scratch, but you can get there removing an element from a larger map.

We start drain_filter, leave -5 to -2 there, and drain -1. Removing key -1 comes down to simply decrementing the length of the left child. We see the left child is underfull, we see there is no left sibling so we turn to the right sibling, we see we can merge with the right sibling so we do that (appending the 0 and 1..=5 keys/values onto the left child, discarding the right child, shortening the parent). On the road to checking whether that has made the parent underfull instead, we see that the parent has become empty. So we pop the empty root off the tree: parent.into_root_mut().pop_level(). This dereferences the map's &'a mut Root<K, V> (which is unique), gets and later deallocates the raw Root::node of type BoxedNode AKA Unique (fine), creates a NodeRef around the same raw pointer, descends to the only remaining child (that descend uses and obtains only raw pointers) and writes its raw pointer to Root::node. Then it creates another NodeRef around the Unique to the new child, traverses that, and erases the former child's parent pointer.

(*self.as_mut().as_leaf_mut()).parent = ptr::null();

or virtually:

let x: *mut LeafNode<K, V> = …;
(*x).parent = ptr::null();

That's a raw pointer access (unless there's something like #73987 going on).

If we turn it into a "real" &mut LeafNode, Miri still does not complain, while I think the whole LeafNode would be invalidated, and that includes the raw pointer in the NodeRef that is returned and assigned to DrainFilterInner::cur_leaf_edgeand later used in the next iterator step. But I'm not sure, and in addition, the latter access is also done using raw pointers, despite the fact that the iterator has unique access, because code is shared with double-ended iterators that do not have unique access.