Recursive impl Extend for HashMap<_, HashSet<_>> where sets are combined, instead of overwritten?

Question

I'm trying to adapt my code to a solution to my previous question. Basically, I have a HashMap<String, HashSet<String>> that should be generated as a result of rayon's par_extend. The problem is that keys repeat and in such a case, I want HashSets to be combined, as opposed to overwritten. In other words, is there a way to add a custom impl Extend here so that the following code executes properly?

use std::collections::{HashMap, HashSet};

fn main() {
    let mut d: HashMap<String, HashSet<String>> = HashMap::new();

    d.extend(vec![1, 2].iter().map(|x| {
        let mut z = HashSet::new();
        z.insert(x.to_string());
        return ("a".into(), z);
    }));

    assert_eq!(d.get("a").unwrap().len(), 2);
}

Does this answer your question? [How do I implement a trait I don't own for a type I don't own?](https://stackoverflow.com/questions/25413201/how-do-i-implement-a-trait-i-dont-own-for-a-type-i-dont-own) — Elias Holzmann, Jul 08 '21 at 16:49
(Title does not fit exactly, but the answers there should fit this question as well) — Elias Holzmann, Jul 08 '21 at 16:51

score 1 · Answer 1 · answered Jul 09 '21 at 15:11

Because of the orphan rules, you can't. You can, however, define a cheap wrapper, implement Extend for that wrapper, do the extending and then unwrap the original map again. Maybe something like this:

use std::collections::{HashMap, HashSet};

type MapOfSets = HashMap<String, HashSet<String>>;

fn main() {
    let mut d: ExtendWrapper = ExtendWrapper::new(HashMap::new());

    d.extend(vec![1, 2].iter().map(|x| {
        let mut z = HashSet::new();
        z.insert(x.to_string());
        return ("a".into(), z);
    }));
    
    let d = d.into_inner();

    assert_eq!(d.get("a").unwrap().len(), 2);
}

struct ExtendWrapper(MapOfSets);

impl ExtendWrapper {
    fn new(map: MapOfSets) -> Self {
        Self(map)
    }

    fn into_inner(self) -> MapOfSets {
        self.0
    }
}

impl Extend<(String, HashSet<String>)> for ExtendWrapper {
    fn extend<T>(&mut self, iter: T)
    where T: IntoIterator<Item = (String, HashSet<String>)>
    {
        for (key, set) in iter {
            // lifetimes make it infeasible to use the entry api here :(
            if let Some(s) = self.0.get_mut(&key) {
                s.extend(set);
                continue;
            }
            self.0.insert(key, set);
        }
    }
}

Playground

Recursive impl Extend for HashMap<_, HashSet<_>> where sets are combined, instead of overwritten?

1 Answers1