Rust: Add drain_filter for BinaryHeap

[RalfJung]

std::collections::BinaryHeap is missing a drain_filter method (rust-lang/rfcs#2140).

[dtolnay]

Sounds good to me. I agree that this part of the API would be good to keep consistent between HashMap/HashSet and BTreeMap/BTreeSet. Please submit a PR.

[nbro]

An issue opened about 1 year ago. 17 thumbs up. Is this issue being addressed?

[Mark-Simulacrum]

I don't believe there's any current in-progress work, but we're happy to take PRs. If you're interested, we can probably connect you with a mentor or provide some instructions.

[nbro]

@Mark-Simulacrum Honestly, right now, I don't think I am qualified enough to address this issue. I would leave it someone else more qualified. Maybe, in the future, I may feel comfortable enough with Rust to address it.

[Ryan1729]

If someone ever gets around to this, it would be a shame to repeat the mistake of passing &T like in the Vec retain implementation. These methods should pass &mut T to the closure, like the HashMap implementation does.

[jq-rs]

@Mark-Simulacrum I could take a look to BtreeMap retain(), if there is some guidance/mentoring available. As a first Rust lang PR it may be otherwise a bit too steep effort.

[Mark-Simulacrum]

cc @rust-lang/libs would someone be able to provide some mentoring for @jq-rs?

[jq-rs]

Getting nowhere so far, so I'll pass this to future volunteers.

[jonhoo]

BinaryHeap now has drain, although for some reason it removes elements in _arbitrary_ order (https://github.com/rust-lang/rust/issues/59278) :/

[hellow554]

[...] it removes elements in _arbitrary_ order (#59278) :/

Is that really a problem? I can't think of a situation where you want to remove elements in order.

Ah sure, it's about iterating in a specific order, not removing (was confused with retain)

[jonhoo]

I mean, it's a heap... If I want to re-use the heap, it seems pretty reasonable to use .drain to walk the heap rather than a while let Some(_) = heap.pop loop. I think this is _especially_ surprising for IntoIterator -- it'd be like a BTreeMap yielding elements in a random order when you do for (k, v) in btreemap.

[czipperz]

Is it acceptable to implement this using brute force and then redo it later for performance gains? Just do something like:

let this = mem::replace(self, Default::default());
*self = this.into_iter().filter(...).collect();

[gendx]

Regarding BTreeMap/Set::drain, it looks like most of the implementation can be derived from IntoIter (apart from emptying the map instead of mem::forget-ing it). I can give it a try if that's the case.

[ssomers]

Subject says it all...

Actually, not to me... which drain? Vec & VecDeque have a range argument; HashMap and BinaryHeap don't (because there is no order or it's not exploitable). BTreeMap has order and exposes range concepts already.

I assumed a range argument was implied, and I believe @arthurprs's comment reveals the same assumption. But then, even if you know how to sew up the remaining tree, you can't simply adjust its length because there is no length concept on ranges.

If drain without argument is what we're after, then I still don't quite understand @gendx's wording in the last comment. Isn't a plain drain _always_ the same as the IntoIterator implementation, apart from the cleanup of self? And in this case, since that cleanup happens in into_iter and outside of IntoIter, can we not just use IntoIter as is, as in my attempt?

[gendx]

I was thinking about the no-range-argument case. There might exist a clever optimization for drain, but indeed given how easy it is to derive drain from into_iter it's unlikely that such optimization wouldn't apply to into_iter as well. So your attempt looks good! (At least from a high-level perspective)

As for the sub-range case, that would definitely be a different and more complex implementation.

[ssomers]

It turns out that a drain without range is trivial. An excerpt from simple code on playground comparing to an existing drain:

        for i in v[0].drain(..) {}
        for i in std::mem::replace(&mut m[0], BTreeSet::new()).into_iter() {}

So I think it's only worth the trouble if it operates on a range.

Secondly, I wondered why retain doesn't take a range either, but that's what drain_filter is for (the #59618 mentioned above). It might die like #1353 did, but more likely it means retain will be useless and we want drain_filter instead.

[ssomers]

Actually, thanks to #61129, in the 1.40 beta version of Rust, the above alternative to drain can be shortened to:

std::mem::take(&mut m[0]).into_iter()

[mbrubeck]

All but two of the methods originally requested here have since been implemented. I'm splitting this issue so that we have one ticket for per method. This ticket is for BinaryHeap::drain_filter.

BTreeMap/BTreeSet::drain_filter are tracked at #70530. BinaryHeap::retain is tracked at #71503. BTreeMap/BTreeSet::retain are part of rust-lang/rfcs#1338.