rust - HashMap borrow issue when trying to implement find or insert

Question

I tried to implement own analogue of find_or_insert method that looks like this:

use std::collections::HashMap;

pub struct SomeManager {
    next: i32,
    types: HashMap<i32, i32>,
}

impl SomeManager {
    pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
        match self.types.get(&k) {
            Some(ref x) => return *x,
            None => {
                self.types.insert(k, self.next);
                self.next += 1;
                return self.types.get(&k).unwrap();
            }
        }
    }
}

fn main() {}

Error:

error[E0502]: cannot borrow `self.types` as mutable because it is also borrowed as immutable
  --> src/main.rs:13:17
   |
10 |         match self.types.get(&k) {
   |               ---------- immutable borrow occurs here
...
13 |                 self.types.insert(k, self.next);
   |                 ^^^^^^^^^^ mutable borrow occurs here
...
18 |     }
   |     - immutable borrow ends here

I know that there are some standard methods that implement this functionality, but I want this method to be as light as possible - it will be called very-very often and almost all of the time the values will already exist.

As I understand it, when we call self.types.get we borrow it to scope of match statement, so we can't call self.types.insert here. I have tried to move methods from None branch out of match statement, but it also fails.

The only working solution that I found requires invoking get twice:

pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
    let is_none = match self.types.get(&k) {
        Some(ref x) => false,
        None => true,
    };
    if is_none {
        self.types.insert(k, self.next);
        self.next += 1;
    }
    self.types.get(&k).unwrap()
}

How can I work around such situations?

Answer 1

There are a handful of methods on HashMap to achieve these sorts of complex cases. Most notably, for your case, HashMap::entry and Entry::or_insert_with:

pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
    self.types.entry(k).or_insert_with(|| {
        let value = self.next;
        self.next += 1;
        value
    })
}

In your case, however, there’s the borrow of self inside, so that won’t do.

We thus shift the borrow of self.next outside of the closure so the compiler can reason about it as disjoint from self.types. Problem solved with only one lookup, as it should be.

pub fn get_type<'a>(&'a mut self, k: i32) -> &'a i32 {
    let next = &mut self.next;

    self.types.entry(k).or_insert_with(|| {
        let value = *next;
        *next += 1;
        value
    })
}