How to drain some elements from `HashMap` with predicate

[fyl2xp1]

I feel rather dumb not being able to solve this by myself.

I have a HashMap where I'd like to modify and maybe remove several entries in order to process them elsewhere.

This is one of several horrible solutions I came up with:

use std::collections::{HashMap, hash_map::Entry};

#[derive(Debug)]
struct SomeObject(u32);

fn main() {

    let mut my_map = HashMap::<String, SomeObject>::default();
    my_map.insert("xyz".to_string(), SomeObject(123));
    my_map.insert("abc".to_string(), SomeObject(321));
    
    // clone(!!) all keys
    let keys = my_map.keys().cloned().collect::<Vec<String>>();
    
    // removed objects go here
    let mut drained_entries = Vec::new();
    
    for key in keys {
        // unnecessary lookup
        match my_map.entry(key) {
            Entry::Occupied(mut entry) => {
                let some_object = entry.get_mut();
                some_object.0 -= 1;
                if some_object.0 < 200 {
                    drained_entries.push(entry.remove());
                }
            }
            Entry::Vacant(_entry) => {
                // 
                unreachable!()
            }
        }
    }
    
    println!("retained {:?}", my_map);
    println!("drained {:?}", drained_entries);
    
    // prints:
    // retained {"abc": SomeObject(320)}
    // drained [SomeObject(122)]
}

(playground: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=d56a8d242df3c3979f6d5bdb1250a04d )

The following ideas didn't work out:

• drain() doesn't allow filtering and will remove every entry
• iter_mut() doesn't allow to remove entries
• retain() doesn't allow to move out the value without cloning
• entry() doesn't provide an iterator

If there was something like filter_drain() or occupied_entries() it would be easy to implement without cloning, double lookups, or excessive memory consumption.

Thank you for helping me out!

[cuviper]

There is drain_filter, but this was added very recently, so it's not yet ready for stable use. In the meantime, you could use hashbrown::HashMap, which is the basis for std's.

[fyl2xp1]

Thanks a lot for your quick and helpful reply!

This works as expected:

#![feature(hash_drain_filter)]
use std::collections::{HashMap};

#[derive(Debug)]
struct SomeObject(u32);

fn main() {

    let mut my_map = HashMap::<String, SomeObject>::default();
    my_map.insert("xyz".to_string(), SomeObject(123));
    my_map.insert("abc".to_string(), SomeObject(321));
    
    let drained_entries = my_map.drain_filter(|_key, some_object| {
        some_object.0 -= 1;
        some_object.0 < 200
    }).collect::<Vec<_>>();

    println!("retained {:?}", my_map);
    println!("drained {:?}", drained_entries);
    
    // prints:
    // retained {"abc": SomeObject(320)}
    // drained [("xyz", SomeObject(122))]
}

[H2CO3]

As a side note, I was going to suggest a manual solution that iterates over Entryes of the map and removes them if necessary.

However, I couldn't find a method on HashMap for acquiring an Iterator<Item=Entry> – this could be an interesting addition to allow manually implementing this and similar operations in general.

[fyl2xp1]

I expected the same to exist. It should, however, be an Iterator<Item=OccupiedEntry> as all entries are going to be occupied.

std::collections::hash_map::OccupiedEntry - Rust

Edit: Ah, I believe it's possible to edit the key of the entry. That would cause havoc when iterating through the set. One would need another type which prohibits editing the key.

[cuviper]

An entry iterator would allow multiple entries to exist at once for the same map. But removal may need to shift items in the table (like the old robin-hood one did), which may invalidate other entries.