How to find or insert into a Vec

Question

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.

Answer 1

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:

• Double mutable borrow error in a loop happens even with NLL on
• "Variable does not live long enough" when returning a Result containing a reference but it does live long enough

Answer 2

Vec is an unordered, not very structured type. It has no way to look up the exact position of an item within it; the closest the default functions get to is contains(), which only tells you if the item is contained.

Furthermore, due to the fact that a Vec is not a Set, the behavior of "find the item or append and return " is undefined - "find the item", if there is a duplicate, needs to be defined further.

To solve this problem without changing to the correct type (HashSet is the type you really want for this. Note the existence of get_or_insert(), which is literally what you are after. It pays to use the proper structure for the job, rather than to try to make everything fit a Vec), we're going to have to build it ourselves. Keeping to your signature, it looks like this (Playground):

trait VecSeekOrAppend<T:PartialEq>:Sized {
    fn get_or_insert(&mut self, item: T) -> &mut T;
}

impl<T> VecSeekOrAppend<T> for Vec<T> 
    where T: PartialEq + Clone {

    fn get_or_insert(&mut self, item: T) -> &mut T {
        if !self.contains(&item) {
            self.push(item.clone());
        }
        for i in self.iter_mut() {
            if i == &mut item {
                return i;
            }
        }
        unreachable!();
    }
}

The reason your initial version does not work is due to the returned lifetime requirement; all methods returning a reference from a Vec require a lifetime validity for the duration of use. By returning such a &mut reference, if you attempt to do it in one go, the mutation of the Vec<_> will happen while there is already a mutable borrow of it.

Splitting the cycle in two, and performing the insertion (without keeping a reference) to then find the eference, allows us to sidestep this problem. Another way to perform this is to store items by a serializable or hashable identifier (the exact way HashMap and HashSet work) in order to innately provide this layer of indirection.

There is a rust feature in the works to ease some of this pain (non-lexical lifetimes), but, as you can see from the github issue, it's not in the near future.