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