How to find or insert into a Vec

I'm trying to write a function that finds returns a mutable reference to an existing element in a Vec, or inserts it if it doesn't exist and returns a mutable reference to the new element.

I've tried a couple of times, but the borrow checker isn't convinced. I've simplified the code I was trying to write to the example below, which gives the same errors.

fn mut_find_or_insert<T: PartialEq>(vec: &mut Vec<T>, val: T) -> &mut T {
    if let Some(u) = vec.iter_mut().find(|u| **u == val) {
        u
    } else {
        vec.push(val);
        vec.last_mut().unwrap()
    }
}

Playground link: https://play.rust-lang.org/?version=stable&mode=debug&edition=2018&gist=cb12c38bcf3682b15a247d14aab48b6b

Rust gives me the following compiler error (full message through the playground link):

error[E0499]: cannot borrow `*vec` as mutable more than once at a time

This seems to be something that should be possible to implement in rust, however it's not clear to me how I reimplement this to avoid borrow checker errors.

Solution

The reason this doesn't work as written is because of a limitation in the current borrow checker. This is very similar to NLL case #3, in which the compiler borrows somewhat overzealously for an entire match statement when the borrow is only used in one of the branches. With the experimental "Polonius" borrow checker (available on the nightly compiler with the -Z polonius flag), your code is accepted as-is.

Working in the stable compiler, it's probably a good idea to redesign your data structures as Sébastien Renauld's answer also suggests, but if you need to make this work with a Vec, you can work around it by briefly using an index to end the borrow:

fn mut_find_or_insert<T: PartialEq>(vec: &mut Vec<T>, val: T) -> &mut T {
    if let Some(i) = vec.iter().position(|each| *each == val) {
        &mut vec[i]
    } else {
        vec.push(val);
        vec.last_mut().unwrap()
    }
}

This works because the result of calling position is not a reference, so the borrow of vec is not held during the if let.

This is similar to the following questions, which manage to find the same limitation using early return from a loop: