Rust function that accepts either HashMap and BtreeMap

Question

fn edit_map_values(
            map1: &mut HashMap<String, i128> || &mut BTreeMap<String, i128>){
    for tuple in map1.iter_mut() {
        if !map1.contains_key(&"key1") {
             *tuple.1 += 1;
        }
    }
    map1.insert(&"key2", 10);
}

How do I write one function that accepts either HashMap and BtreeMap like in the example above?

Answer 1

It is possible to abstract over types by using traits and for your specific use-case, you can take a look at this more constrained example.

use core::{borrow::Borrow, hash::Hash};
use std::collections::{BTreeMap, HashMap};

trait GenericMap<K, V> {
    fn contains_key<Q>(&self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Hash + Eq + Ord;

    fn each_mut<F>(&mut self, cb: F)
    where
        F: FnMut((&K, &mut V));

    fn insert(&mut self, key: K, value: V) -> Option<V>;
}

impl<K, V> GenericMap<K, V> for HashMap<K, V>
where
    K: Eq + Hash,
{
    fn contains_key<Q>(&self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Hash + Eq + Ord,
    {
        self.contains_key(k)
    }

    fn each_mut<F>(&mut self, mut cb: F)
    where
        F: FnMut((&K, &mut V)),
    {
        self.iter_mut().for_each(|x| cb(x))
    }

    fn insert(&mut self, key: K, value: V) -> Option<V> {
        self.insert(key, value)
    }
}

impl<K, V> GenericMap<K, V> for BTreeMap<K, V>
where
    K: Ord,
{
    fn contains_key<Q>(&self, k: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Hash + Eq + Ord,
    {
        self.contains_key(k)
    }

    fn each_mut<F>(&mut self, mut cb: F)
    where
        F: FnMut((&K, &mut V)),
    {
        self.iter_mut().for_each(|x| cb(x))
    }

    fn insert(&mut self, key: K, value: V) -> Option<V> {
        self.insert(key, value)
    }
}

fn edit_map_values<T: GenericMap<String, i128>>(map: &mut T) {
    map.each_mut(|(k, v)| {
        if k != "key1" {
            *v += 1;
        }
    });
    map.insert("key2".into(), 10);
}

fn main() {
    let mut hm: HashMap<String, i128> = [("One".into(), 1), ("Two".into(), 2)]
        .iter()
        .cloned()
        .collect();
    let mut btm: BTreeMap<String, i128> = [("Five".into(), 5), ("Six".into(), 6)]
        .iter()
        .cloned()
        .collect();
    dbg!(&hm);
    dbg!(&btm);
    edit_map_values(&mut hm);
    edit_map_values(&mut btm);
    dbg!(&hm);
    dbg!(&btm);
}

Answer 2

Way back before the 1.0 release, there used to be Map and MutableMap traits, but they have been removed before stabilization. The Rust type system is currently unable to express these traits in a nice way due to the lack of higher kinded types.

The eclectic crate provides experimental collection traits, but they haven't been updated for a year, so I'm not sure they are still useful for recent versions of Rust.

Further information:

• Does Rust have Collection traits?
• No common trait for Map types? (Rust language forum)
• Associated type constructors, part 1: basic concepts and introduction (blog post by Niko Matsakis)
• Generic associated type RFC

Answer 3

While there is no common Map trait, you could use a combination of other traits to operate on an Iterator to achieve similar functionality. Although this might not be very memory efficient due to cloning, and also a bit involved depending on the kind of operation you are trying to perform. The operation you tried to do may be implemented like this:

fn edit_map_values<I>(map: &mut I)
where
    I: Clone + IntoIterator<Item = (String, i128)> + std::iter::FromIterator<(String, i128)>,
{
    // Since into_iter consumes self, we have to clone here.
    let (keys, _values): (Vec<String>, Vec<_>) = map.clone().into_iter().unzip();

    *map = map
        .clone()
        .into_iter()
        // iterating while mutating entries can be done with map
        .map(|mut tuple| {
            if !keys.contains(&"key1".to_string()) {
                tuple.1 += 1;
            }
            tuple
        })
        // inserting an element can be done with chain and once
        .chain(std::iter::once(("key2".into(), 10)))
        .collect();
        // removing an element could be done with filter
        // removing and altering elements could be done with filter_map
        // etc.
}

fn main() {
    use std::collections::{BTreeMap, HashMap};

    {
        let mut m = HashMap::new();
        m.insert("a".to_string(), 0);
        m.insert("key3".to_string(), 1);

        edit_map_values(&mut m);

        println!("{:#?}", m);
    }

    {
        let mut m = BTreeMap::new();
        m.insert("a".to_string(), 0);
        m.insert("key3".to_string(), 1);

        edit_map_values(&mut m);

        println!("{:#?}", m);
    }
}

Both times the output is the same, except for the order of the HashMap of course:

{
    "a": 1,
    "key2": 10,
    "key3": 2,
}