Rust 的 HashMap 中“插入或更新”操作的最佳实践是什么？

Question

I'm looking for a best practice for the insert-or-update operation, because it is very commonly used, I think we need to optimize both its writing style and efficiency.

Assume the following scenario: I got a hashmap

let classes: HashMap(String, HashSet<String>) = HashMap::new()

which is used to store students and whos classes.

The data is expected to be in the following form:

{ key: "ClassA", value: {"Bob", "Marry", "Jack"}},
{ key: "ClassB", value: {"Lee", "Tom"}},

Now I got a new set of data of a student and his/her class, let's take:

{ name: "Alice", class: "ClassC"}

Since I'm not sure if the class already appears in the classes HashMap, I need to figure out whether it exists first, if it is, I'm going to update the value, and if it's not, I'm going to add a new key->value pair.

What's the correct way to do it with out any unnecessary move or copy? Based on other answers, I tried to use std::collections::hash_map::Entry but I failed.

Thanks!

Answer 1

The idiomatic way to use that map would be something like this:

use std::collections::HashSet;
use std::collections::HashMap;

fn main() {
    let mut classes: HashMap<String, HashSet<String>> = HashMap::new();
    
    let e = classes.entry(String::from("ClassA"));
    e.or_default().insert(String::from("Alice"));

    let e = classes.entry(String::from("ClassA"));
    e.or_default().insert(String::from("Bob"));
    
    dbg!(&classes);
}

The HashMap::entry() function returns an Entry value, that represents either the value contained in the map or the place it would be if it were in the map. This Entry type has a lot of functions to access the contained value, and create it if needed. In your case the easiest function to use is or_default that creates a default value (an empty set) if the value is not yet in the map.

Then, since you have a mutable reference to the set inside the map, just insert the desired value.

Answer 2

To get an insert or update behavior I would do it in this fluent way

classes.entry(String::from("ClassC"))
    .or_default()
    .insert(String::from("Alice"));

A test example which covers your scenario

#[test]
fn test(){
    // Arrange
    let mut classes: HashMap<String, HashSet<String>> = HashMap::from([
        (String::from("ClassA"), HashSet::from([String::from("Bob"), String::from("Marry"), String::from("Jack")])),
        (String::from("ClassB"), HashSet::from([String::from("Lee"), String::from("Tom")])),
    ]);
    let class_c = String::from("ClassC");
    let alice = String::from("Alice");

    // Act
    classes.entry(class_c.clone())
        .or_default()
        .insert(alice.clone());

    // Assert
    let contains_class = classes.contains_key(&class_c);
    let contains_alice = classes.get(&class_c)
        .unwrap().contains(&alice);
    assert_eq!(true, contains_class);
    assert_eq!(true, contains_alice);
}