Pre-RFC: `{BTreeMap, BTreeSet}::drain` with a range

Summary

Implement BTreeMap::drain(range) and BTreeSet::drain(range) following rfc#1257-drain-range-2.

Motivation

As BTreeMap and BTreeSet are ordered collections, it is reasonable to have a drain method which takes a range parameter to allow range removal.

Currently, there are .range(range) and .drain_filter(pred) methods available, and the .drain(range) function can be achieved by .drain_filter(|k| range.contains(k)). It requires full traverse of the tree, in O(n log(n)), but for a B-Trees, only O(m log(n)) is needed, where m is the length of the range. Having .drain(range) method can undoubtedly benefit the situations where it needs to remove a range from a BTreeMap or BTreeSet.

There have be previous discussions on the furum such as BTree{Map,Set} range removal and BTreeMap::drain with a range.

Guide-level explanation

The signature of BTreeMap::<K, V>::drain method will be

pub fn drain<T: ?Sized, R>(&mut self, range: R) -> DrainRange<'_, K, V>
where 
    T: Ord,
    K: Borrow<T> + Ord,
    R: RangeBounds<T>,
{ /* ... */ }

Like .range, It returns an iterator that drains out all the nodes bounded in the given range, and like .drain_filter, the drained items will be definitely removed from the tree when the returned iterator is dropped.

Here is a short example:

use std::collections::BTreeSet;

let mut set: BTreeSet<i32> = (0..=10).collect();
let drained: Vec<i32> = set.drain(5..=8).collect();
assert_eq!(set.into_iter().collect::<Vec<_>>(), [0, 1, 2, 3, 4, 9, 10]);
assert_eq!(drained, [5, 6, 7, 8]);

Reference-level explanation

Public API

pub mod alloc {
    pub mod collections {
        pub mod btree_map {
            impl<K, V> BTreeMap<K, V> {
                pub fn drain<T: ?Sized, R>(&mut self, range: R) -> DrainRange<'_, K, V>
                where 
                    T: Ord,
                    K: Borrow<T> + Ord,
                    R: RangeBounds<T>,
                { }
            }

            pub struct DrainRange<'a, K: 'a, V: 'a> {}

            impl<K, V> Iterator for DrainRange<'_, K, V> {
                type Item = (K, V);
                fn next(&mut self) -> Option<(K, V)> {}
                fn min(self) -> Option<(K, V)> {}
                fn max(self) -> Option<(K, V)> {}
            }

            impl<K, V> FuseIterator for DrainRange<'_, K, V> {}

            impl<K, V> DoubleEndedIterator for DrainRange<'_, K, V> {
                fn next_back(&mut self) -> Option<(K, V)> {}
            }

            impl<K, V> Drop for DrainRange<'_, K, V> {
                fn drop(&mut self) {}
            }

            impl<K, V> Debug for DrainRange<'_, K, V> {
                fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {}
            }
        }

        pub mod btree_set {
            impl<T> BTreeSet<T> {
                pub fn drain<K: ?Sized, R>(&mut self, range: R) -> DrainRange<'_, T>
                where 
                    K: Ord,
                    T: Borrow<K> + Ord,
                    R: RangeBounds<K>,
                { }
            }

            pub struct DrainRange<'a, T: 'a> {}

            impl<T> Iterator for DrainRange<'_, T> {
                type Item = T;
                fn next(&mut self) -> Option<T> {}
                fn min(self) -> Option<T> {}
                fn max(self) -> Option<T> {}
            }

            impl<T> FuseIterator for DrainRange<'_, T> {}

            impl<T> DoubleEndedIterator for DrainRange<'_, T> {
                fn next_back(&mut self) -> Option<T> {}
            }

            impl<T> Drop for DrainRange<'_, T> {
                fn drop(&mut self) {}
            }

            impl<T> Debug for DrainRange<'_, T> {
                fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {}
            }
        }
    }
}

Overview

There are two key points in DrainRange:

  • it must act like an iterator, i.e., it must provide capability of iterations over elements bounded by range;
  • it is intended to do range removal after .drain(range) is called.

Let's highlight the two points one by one.

Range Iteration

DrainRange iterators the elements in range like what the Range and RangeMut does, except that DrainRange takes the KVs' ownership and returns to its caller. But for efficiency consideration, unlike DrainFIlter, it does not immediately remove the node from the tree when calling next and next_back. Instead, it leaves the iterated nodes uninitialized. The drained nodes are actually removed when dropping the DrainRange iterator.

Range Removal

Range removal is generally more efficient that single-node removal because of fewer balancing operations. That is the reason why we choose range removal but not multiple single-node removals.

Before removal

After draining, there is a big "hole" in the tree, and what we do is to remove the middle empty nodes, and collapsing the left and right half nodes, for each layers.

         n2 +----+----+----+
            | A  | E  | I  |
            +----+----+----+
                /      \
               :   ...  :
              :          :
n1           /            \ n4
+----+----+----+   n3   +----+----+----+
| B  | C  | D  |   ...  | F  |  G |  H |
+----+----+----+        +----+----+----+

Imagine on the B-Tree as above, after draining a range (e.g. D..=F), the following changes will happen:

  • For the lowest common ancestor (LCA) layer (e.g. n1), a subrange (e.g. E) of no more than one node is emptied, between one edge node (excluded, called left edge bound, or left edge) to another edge node (excluded right edge).
  • For any layer except the LCA layer, there is a left node (eg. n1) of which a suffix is emptied, and a right node (eg. n4) of which a prefix is emptied. All the nodes between the left node and the right node is completely emptied. Call the empty and non-empty bound of the left node and the right node as left edge and right edge, respectively as the LCA layer.
  • For any layer, any edge node between two empty KV node is emptied, as well as its descendants if it is non-leaf.

Bottom-up collapse algorithm

As we already know the left leaf edge and the right leaf edge after iteration, range removal can be achieived in a bottom-up way.

Here the terminology "collapse" means a combination of merging and stealing, similar with the implementation of single KV removal in BTree, but more complicated. The next section is about how to collapse the two half nodes into one.

To clarify, here are some other important terminologies:

  • left_node and right_node represent the two nodes which have an empty suffix and an empty prefix, respectively. left_node and right_node might be the same node in the LCA layer.
  • left_edge and right_edge represent the edge bound of empty and non-empty nodes of left_node and right_node, respectively.
  • left_parent and right_parent represent the parent edges of left_node and right_node, respectively. They are both optional.
  • collapsed represent the edge which left_edge and right_edge collapsed into. It might contains descendants for non-leaf layers.

The collapse algorithm is as the following pseudocode:

fn remove_range(left_edge: Edge, right_edge: Edge) -> Edge {
    let mut left_node = left_edge.to_node();
    let mut right_node = right_edge.to_node();

    let leaf = collapsed(left_node, right_node);

    left_node = left_edge.get_parent_node();
    right_node = right_edge.get_parent_node();
    while (left_node, right_node) not None {
        let collapsed: Edge = collapse(left_node, right_node);

        let left_parent = left_edge.get_parent_edge();
        let right_parent = right_edge.get_parent_edge();
        collapsed.link_as_left_child(right_parent.left_kv());

        left_node = left_parent.to_node();
        right_node = right_parent.to_node();
    }

    leaf
}

After and before each collapse, collapsed edge is linked as a left child of right_parent, and left_parent has no right child. This makes the left KV node of left_parent (i.e. left_edge of the next collapse) free to move, and it is very useful during collapse.

The collapse

Collapse is the most complicated part of this algorithm. Let's reveal its details.

The first thing is to recursively deallocate all the internal edge nodes (including the empty left_edge) located on left_node and right_node, since they are empty.

Successively, to reuse the code of BalancingContext to rebalence the tree, we shift all the KVs and edges at the right of right_edge (inclusive) to compress the right_node, and then using the merging or stealing strategy to rebalance the tree. We will handle this later.

Collapse of non-LCA

For non-LCA layers, take the following figure for example, where we the drained range is D..=G.

   +---------+====+     +====+---------+
   | A       | E  | ... | F  | K       |
   +---------+====+     +====+---------+
`left_parent`/                \  `right_parent`
            /                  \
`left_node`/                    \  `right_node`
+----+----+====+            +====+----+----+
| B  | C  | D  |    ...     | G  | I  |  J |
+----+----+====+            +====+----+----+
/    |    :                     /     |    |
B'   C'                        H      I'   J'
         (`left_edge` is empty) (H is `right_edge` as a child)

First check if sum of the length of drained left_node and right_node is no larger than CAPACITY. If yes, it means we can simply merge left_node and right_node into a single node, make it a child of right_parent, and leave another empty node as the child of left_parent, being sure that the next collapse iteration will deallocate it. (In the figure above, mrege |B|C| with |I|J| into a single node, and let it be a left child of K.)

Note that this merge operation is quite simple, and unlike that of single KV removal, it does not touch any nodes in the parent layer at all (keep in mind that left_edge is always empty before and after the collapse). There are left_node.len() + right_node.len() KV data, and left_node.len() + right_node.len() + 1 edge children before the merge, which is consistent with after the merge. (In the figure above, merged node |B|C|I|J| has exactly 5 children B', C', H, I', J'.)

Things are more complicated when left_node and right_node cannot be simply merged into one, which means there are more than CAPACITY KV data in left_node plus right_node.

We construct a BalanceContext by pushing up the right most non-empty KV datum of left_node to the left of right_parent, naming the pushed KV as right_parent_kv, then linking the left_node as a left child of right_parent_kv, and we get a valid BalanceContext consisting of right_parent_kv, left_node, and right_node. (In the above figure, push C up to replace F, and the BalanceContext is C as parent, |B| as left child, and I|J as right child.)

After that, rebalance left_node and right_node via stealing. Here left_node and right_node are not really collapsed into a single node, but we can simply regard the left_node as the collapsed new node, move the right_parent pointer to its left adjancent edge (actually the left_node before), then go on the next collapse.

Collapse of LCA

For the LCA node, simply shifting all edges and KVs of the right part can make it a single node. If it underflows after collapse, rebalance the node merging or (bulk) stealing like that in single KV removal. Go on to fix its ancestors and notify the caller to handle pop internal root if the root node becomes empty.

Cornor cases

When pushing up KV datum beside left_edge, if the right_parent node happens to be full, we first try to push up to the left_parent and do stealing in a semmetric way. If both left_parent and right_parent is full, push it up to the nearest ancestor which has an empty KV to place it, then skip all the passed-by layers where both left_node and right_node are full (there are actually no work to do with these layers), and continue the next collapse on the layer where the descentant KV has been pushed up to.

If pushing up to its ancestor leaves left_node and right_node underflowing, do stealing on each of them as where we do in single KV removol.

Drawbacks

  • The recursive empty edge deallocation operation might not be efficient enough and has risk of running out of stack space.
  • Compared to Vec::drain, VecDeque::drain and String::drain, which take a index range as input parameter, this function signature might be confusing where the input range parameter means a range of keys, instead of indices.

Alternatives

Interval Deletion in B-trees proposed a top-down range deletion algorithm. In this algorithm, before descendant to the children during the binary search, it first deletes nodes with key bounded within the range, which makes ordered iterating of the deleted range is not possible.

Recursively deallocating the empty edge nodes can be done by a traverse from the left_edge to the right_edge, which deallocates all the empty edge nodes before collapsing in advance. Or even this can be done when iterating the DrainRange.

Unresolved questions

  • Should the method be named drain_range?
2 Likes

Doesn't that imply that mem::forgetting a partially iterated DrainRange is unsound?

I think this would need to be considered in the context of drain_filter:

1 Like

I'm afraid yes. DrainFilter removes the node while iterating by tracking the deletion pointer. DrainRange can also do it similarly, but it is much more complicated to track two pointers other than one.

Oh, I didn't catch the full context of what people previously discussed about. There seems to be a lot of unresolved problems of draining.

Take a step back, what if only adding a remove_range method but allowing the deleted elements visited in order via a FnMut(K, V)? I think this RFC should make sense.

That's a bit unclear to me. Are you focusing on the "range removal" bit or on the "after"? In the latter case, obviously nothing can happen before the call. So do you mean it's not allowed to happen before the call returns?

The algorithm you describe is sounds pretty much like BTree: add drain and split_off_range methods by ssomers · Pull Request #81075 · rust-lang/rust · GitHub.

Take a step back, what if only adding a remove_range method but allowing the deleted elements visited in order via a FnMut(K, V) ?

It would definitely be easier to write, just an alternative/extra wrapper to the core algorithm, like drain or split_off. And I think it would be easier to understand and use too.

1 Like

I think building drain on top of split_off_range actually makes a lot of sense, especially for multithreaded code. As described in the PR @ssomers referenced, there is no need to visit all of the nodes in the subtree that is being split off. As a result, splitting off a range should be O(log n) time. This means that the original tree spends less time being locked, potentially allowing other threads to make progress while the drain is running.

My mistake. Here it should be "do range removal when DrainRange is dying".

Thank you for the information, I didn't see this PR, and it seems almost completed. I think it's time to close this topic now.

I'm not sure exactly what that operation is, or how my one-year-old PR works for that matter, but you're right that stack space is apparently a concern in other BTree algorithms. But it's quite doable to let the code descend iteratively rather than through recursion, even if it has to ascend again later. The tree traversal was made bidirectional (somewhat against Rust's spirit) many years ago for that reason. So this isn't an actual drawback.

Here it should be "do range removal when DrainRange is dying".

But why the focus on dying? I genuinely don't understand where this idea comes from. Just do range removal during the call to drain/drain_range/remove_range/…

I didn't see this PR, and it seems almost completed.

It's almost completely forgotten. Your topic is most welcome to revive interest.

1 Like

I'm not sure where it origins, but mentioned about drain, the first thing that come up in my mind is, to take something one by one out of a collection, and then fix the original collection to shrink.

Having read your PR, I see in your .drain(range) method, first split off the original tree into two new ones, and then consume the new tree.

The semantics looks much clearer than my version, but I mind a little where it must make a new tree before using the drained items. I mean it might need extra allocations when splitting the nodes.

Yes it does, though only from the LCA down to the leaves, along the lower and upper bound of the range. When draining a large range from a vast tree, most of the new tree consists of old tree nodes that haven't even been visited. There's the O(log(N)). Also, the new tree isn't made BTreeMap-conformant unless needed, i.e., when it's turned into a split off tree. I guess more optimization is possible (for instance, when all the elements in a node go right, they are all moved out - perhaps swapping the node itself would be faster).

However there's no doubt that if what the caller wants is simply remove_and_drop_range, it's wasteful. Though even remove_and_drop_range in place is hard: you have to withstand a lone panic in a drop handler and continue the drain or repair the tree.

I don't follow. Surely, the tree is locked (under &mut) as long as the Drain iterator is around? Could

pub fn drain<T: ?Sized, R>(&mut self, range: R) -> Drain<K, V>

be altered such that rustc knows that Drain is self-contained? I forget (or never fully understood) how rustc works this out.

The idea is that the drain iterator will take ownership of the tree returned by split_off_range(). The parent tree only needs to be locked while split_off_range() is being executed, once the subtree has been extracted from the parent, they are separate trees. That lets you unlock the parent, while still having a draining iterator.

It's all about lifetimes. If it's

pub fn drain<T: ?Sized, R>(&mut self, range: R) -> Drain<'_, K, V>

// or equivalently

pub fn drain<'a, T: ?Sized, R>(&'a mut self, range: R) -> Drain<'a, K, V>

// which can (unfortunately and confusingly) also be written

pub fn drain<T: ?Sized, R>(&mut self, range: R) -> Drain<K, V>
// unless you activate `#![warn(elided_lifetimes_in_paths)]`

(which is basically the kind of API that Vec::drain offers) then the lifetime parameter of Drain is coupled to the lifetime of &mut self, so the tree remains locked (i.e. mutably borrowed) for the duration of Drain's existence. If it's

pub fn drain<T: ?Sized, R>(&mut self, range: R) -> Drain<K, V>

(and Drain doesn't have any (elided) lifetime parameter) then the map is immediately unlocked (i.e. no longer mutably borrowed) after the drain call. So you could release a lock on Mutex<BTreeMap<...>> immediately after the drain call and before doing the iteration, if the API was like this. This is similar to BTreeMap::split_off which returns a BTreeMap<K, V>, i.e. a type without any lifetime parameter, hence &mut self does not need to stay alive for as long as the returend BTreeMap exists).

2 Likes

Could this be solved with this humorously named technique?

1 Like

Yes, we can, with the cost of leaking the whole container before draining, like what DrainFilter does for other containers.

Over the days, I have been tried figuring out a way to solve the problems with drop-on-drain, and found the following idea. Unfortunately, I'm afraid it is nearly impossible to find a reliable way to make sure the clean up work being done before the destruction of any RAII objects.