Get mutable reference to element of Vec or create new element and get that reference

Question

I have a Vec<State> list and want to search for an element and get a mutable reference to it. If it doesn't exist, a new default element should be created and added to the list:

struct State {
    a: usize,
}

fn print_states(states: &Vec<State>) {
    for state in states {
        print!("State{{a:{}}} ", state.a);
    }
    println!();
}

fn main() {
    let mut states = vec![State { a: 1 }, State { a: 2 }, State { a: 3 }];

    print_states(&states);

    let mut state = match states.iter_mut().find(|state| state.a == 2) {
        Some(state) => state,
        None => {
            let new_state = State { a: 3 };
            states.push(new_state);
            states.last().unwrap()
        }
    };
    state.a = 4;
    drop(state);
    print_states(&states);
}

This leads to:

error[E0594]: cannot assign to `state.a` which is behind a `&` reference
  --> src/main.rs:25:5
   |
17 |     let mut state = match states.iter_mut().find(|state| state.a == 2) {
   |         --------- help: consider changing this to be a mutable reference: `&mut State`
...
25 |     state.a = 4;
   |     ^^^^^^^^^^^ `state` is a `&` reference, so the data it refers to cannot be written

The problem is the None path. When using None => panic!() without the creation of this new default element I can modify the found element

What do I need to change to make this work?

Answer 1

Your problem is the state.last().unwrap()-line. The method .last() on Vec returns a &State, which causes the compiler to infer the type of state to be &State (which the &mut State from the Some()-case can be coerced to). This is why you can't change state on line 28.

Change the line to state.last_mut().unwrap() and state will be a &mut State instead of &State. Your example compiles after this.