Entry or_insert with Result or Option

[MellowMelon]

My aim is to write a function

fn get_mut_or_insert_result<K, V, F>(
    map: &mut BTreeMap<K, V>,
    key: K,
    or_insert: F,
) -> Result<&mut V>
where
    K: Ord,
    F: FnOnce() -> Result<V>,

which will do exactly one of:

• when the key exists, return the mutable reference to the existing value without calling or_insert.
• when the key is missing, insert the or_insert return (if Ok) into the map and return a mutable reference to it.
• return Err if or_insert is called and returns Err.

In this topic, the same question was asked, guessing a manual match was the needed workaround. The replies reached the same conclusion. I did too in my independent attempt. But the topic did not show the workaround. When trying to implement it myself, I got bitten by some conservative borrow checking, and I didn't like my final solution.

Here was my first attempt:

// (function types omitted, see first snippet)
fn get_mut_or_insert_result(map, key, or_insert) {
    if let Some(value) = map.get_mut(&key) {
        Ok(value)
    } else {
        let value = or_insert()?;
        // ERROR: second mutable borrow of map
        Ok(map.entry(key).or_insert(value))
    }
}

The compiler says map is being borrowed more than once at a time. This error remains even if a return is used in the first branch instead of else. I see that it's being conservative about what get_mut does with its mutable borrow, but even then I don't get why a returned None would be live in the else block.

My second attempt did work:

// (function types omitted, see first snippet)
fn get_mut_or_insert_result(map, key, or_insert) {
    let mut e = map.entry(key);
    match e {
        std::collections::btree_map::Entry::Vacant(_) => {
            let value = or_insert()?;
            Ok(e.or_insert(value))
        },
        std::collections::btree_map::Entry::Occupied(_) => {
            Ok(e.or_insert_with(|| unreachable!()))
        },
    }
}

That unreachable is ugly though, and or_insert_with is no help for understanding intent.

Trying to use the value inside Occupied, like its get_mut method, always ran into errors about borrowing from temporaries.

In total, the questions raised:

1. What's the best way to implement this function?
2. Why did the borrow checker fail the first attempt?
3. Why doesn't OccupiedEntry::get_mut return &'a mut V?

[R081n]

The into_mut method on OccupiedEntry should work

[MellowMelon]

... That was right in the documentation too. My bad for missing that. Thank you for the reminder.

[jendrikw]

@R081n Thanks, I was about to write the same thing.

Here a full example:

fn get_mut_or_insert_result<K, V, E, F>(
    map: &mut BTreeMap<K, V>,
    key: K,
    or_insert: F,
) -> Result<&mut V, E>
where
    K: Ord,
    F: FnOnce() -> Result<V, E>,
{
    let mut e = map.entry(key);
    match e {
        std::collections::btree_map::Entry::Vacant(ve) => {
            let value = or_insert()?;
            Ok(ve.insert(value))
        },
        std::collections::btree_map::Entry::Occupied(mut oe) => {
            Ok(oe.into_mut())
        },
    }
}

[steffahn]

@jendrikw

let mut e = map.entry(key);
match e {
    std::collections::btree_map::Entry::Vacant(ve) => {
        let value = or_insert()?;
        Ok(ve.insert(value))
    },
    std::collections::btree_map::Entry::Occupied(mut oe) => {
        Ok(oe.into_mut())
    },
}

I’d probably further compactify this a bit by use std::collections::btree_map::Entry; and inlining the lets, like

match map.entry(key) {
    Entry::Vacant(e) => Ok(e.insert(or_insert()?)),
    Entry::Occupied(e) => Ok(e.into_mut()),
}

[simonbuchan]

To tack onto this answer, I generally either import std::collections::hash_map or the like to make the type a bit clearer and avoid conflicts , similar to importing std::io so you can use io::Error; or I only import in the local function, which solves the same issues in a different way.