How to drain some elements from `HashMap` with predicate

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!

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.

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))]
}

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.

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

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.

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.