Get `&mut Map` back from entry API

Consider the following API:

use std::collections::HashMap;
use std::hash::Hash;

pub fn get_priority_mut<'a, K, V>(
    map: &'a mut HashMap<K, V>,
    key1: &K,
    key2: &K,
) -> Option<&'a mut V>
where
    K: Eq + Hash,
{
    // Check for the first key without mutably borrowing the map
    if map.contains_key(key1) {
        map.get_mut(key1)
    } else {
        // If the first isn't there, we are free to mutably borrow for the second
        map.get_mut(key2)
    }
}

I would like to avoid the double-lookup of key1, but we can't because of the usual borrow checker limitation. The usual fix for that is to use entry api, but even that doesn't work in this case:

use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::hash::Hash;

// ❌ THIS WILL NOT COMPILE
pub fn get_priority_mut_entry<'a, K, V>(
    map: &'a mut HashMap<K, V>,
    key1: K,
    key2: K,
) -> Option<&'a mut V>
where
    K: Eq + Hash,
{
    match map.entry(key1) { 
        Entry::Occupied(occupied) => Some(occupied.into_mut()),
        Entry::Vacant(_vacant) => {
            match map.entry(key2) {
                Entry::Occupied(occupied2) => Some(occupied2.into_mut()),
                Entry::Vacant(_) => None,
            }
        }
    }
}
error[E0499]: cannot borrow `*map` as mutable more than once at a time
  --> src/lib.rs:17:19
   |
 6 | pub fn get_priority_mut_entry<'a, K, V>(
   |                               -- lifetime `'a` defined here
...
14 |     match map.entry(key1) { 
   |           --- first mutable borrow occurs here
15 |         Entry::Occupied(occupied) => Some(occupied.into_mut()),
   |                                      ------------------------- returning this value requires that `*map` is borrowed for `'a`
16 |         Entry::Vacant(_vacant) => {
17 |             match map.entry(key2) {
   |                   ^^^ second mutable borrow occurs here

I'd like to propose a potential solution: Add a method to [Vacant|Occupied]Entry that consumes the entry and returns &mut HashMap. If you had this API, you could write this:

match map.entry(key1) { 
    Entry::Occupied(occupied) => Some(occupied.into_mut()),
    Entry::Vacant(vacant) => {
        match vacant.into_map().entry(key2) {
            Entry::Occupied(occupied2) => Some(occupied2.into_mut()),
            Entry::Vacant(_) => None,
        }
    }
}

which does not run into the borrow checker limitation.

4 Likes

Doesn't a simple if let else cover this?

    let a_or_b = if let a @ Some(_) = map.get_mut("a") {
        a
    } else {
        map.get_mut("b")
    };

A method like Entry::or_get(alternative_key) for expressing this succinctly in cases where you are already using the Entry API would be nice to have, that's for sure.

1 Like

Your playground is missing the essential detail for running into the borrow-checker limitation: Returning the Option<&mut String> value from a function :wink:

It doesn't have a reference to give back to you -- note that Entry<'_, K, V> isn't even parameterized for the original S hasher in HashMap<K, V, S>.

6 Likes

Damn!

2 Likes

But, this way of writing it at least works with polonius.

The suggestion might still be valid for hashbrown though, since its VacantEntry does have a S parameter and contain a &mut HashMap.

5 Likes

True, and FWIW its hash_table entries already have an into_table method like that.

1 Like