Get mut ref else insert

Question

Below is a simplified version of my problem. I want to get a mutable reference to an element of a vector, similar to .get_mut() except using comparison of the actual value, rather than just the index, and to add that value to the index if it is not found and then return a mut ref to it.

I am stuck on the fact that data seems to be getting moved into the iter for the for loop which means I can't then continue to use it. I have tried different things like putting the for loop in a {} block, using &data, using .for_each(), but can't find a solution.

fn get_or_push(data: &mut Vec<i32>, element: i32) -> &mut i32 {
    for x in data {
        if *x == element {
            return x;
        }
    }
    data.push(element);
    data.last_mut().unwrap()
}

fn main() {
    let mut data = vec![1, 2];

    // get mut ref to element if it exists or push it and return mut ref
    {
        // add block so two goes out of scope and is dropped before getting another mut ref for three
        let two = get_or_push(&mut data, 2);
        *two = *two * 2;
    }
    {
        let three = get_or_push(&mut data, 3);
        *three = *three * 2;
    }
    println!("{:?}", data); // should print [1, 4, 6]
}