Wrap type in enum and return reference

Question

type Id = u8;

struct A {
    id: Id,
}
struct B {
    id: Id,
}
struct C {
    id: Id,
}

struct State {
    a_vec: Vec<A>,
    b_vec: Vec<B>,
    c_vec: Vec<C>,
}

impl State {
    fn new() -> Self {
        Self {
            a_vec: Vec::new(),
            b_vec: Vec::new(),
            c_vec: Vec::new(),
        }
    }

    fn get_e0(&self, id: Id) -> &E0 {

        if let Some(a) = self.a_vec.iter().find(|x| x.id==id) {
            &E0::A(a)
        } else if let Some(b) = self.b_vec.iter().find(|x| x.id==id) {
            &E0::B(b)
        } else {
            panic!("ahh that id doesn't exist everbody panic!!!")
        }
    }
    
    fn get_e0_mut(&mut self, id: Id) -> &mut E0 {
        if let Some(a) = self.a_vec.iter_mut().find(|x| x.id==id) {
            &mut E0::A(a)
        } else if let Some(b) = self.b_vec.iter_mut().find(|x| x.id==id) {
            &mut E0::B(b)
        } else {
            panic!("ahh that id doesn't exist everbody panic!!!")
        }
    }
}

enum E0 {
    A(A),
    B(B),
}

enum E1 {
    A(A),
    C(C),
}

fn main() {
    let state = State::new();

    let a0 = A { id: 0 };
    let a1 = A { id: 1 };
    let b0 = B { id: 2 };
    let c0 = C { id: 3 };
    state.a_vec.push(a0);
    state.a_vec.push(a1);
    state.b_vec.push(b0);
    state.c_vec.push(c0);

    let e5 = state.get_e0(1);
}

I'm looking for a way to implement the function get_e0 and get_e0_mut that wrap several types into an enum so the caller doesn't have to care which of A or B their id relates to, only that they will get an E0 . Yet an Vec of E0 's seems unfeasible as there might be separate grouping such as E1 .

If these functions are not possible then is there another method that could be used to reduce the overhead of searching all the respective Vec's individually each time.

It is guaranteed that the all id's are unique.

Answer 1

You cannot return a reference to a temporary. Instead, you can make your enums generic over their contents. You can therefore use a single enum:

enum E0<T, U> {
    A(T),
    B(U),
}

You can then use it like this:

fn get_e0(&self, id: Id) -> E0<&A, &B> {
    if let Some(a) = self.a_vec.iter().find(|x| x.id == id) {
        E0::A(a)
    } else if let Some(b) = self.b_vec.iter().find(|x| x.id == id) {
        E0::B(b)
    } else {
        panic!("ahh that id doesn't exist everbody panic!!!")
    }
}

fn get_e0_mut(&mut self, id: Id) -> E0<&mut A, &mut B> {
    if let Some(a) = self.a_vec.iter_mut().find(|x| x.id == id) {
        E0::A(a)
    } else if let Some(b) = self.b_vec.iter_mut().find(|x| x.id == id) {
        E0::B(b)
    } else {
        panic!("ahh that id doesn't exist everbody panic!!!")
    }
}

Thanks to lifetime elision rules, you don't have to specify lifetimes.

Playground link

Note that if you want to avoid the panic, your return type should express the notion that there can be no value found.

You can for example return an Option:

fn get_e0(&self, id: Id) -> Option<E0<&A, &B>> { ... }

Or alter the enum to have a None variant, similar to Option:

enum E0<T, U> {
    A(T),
    B(U),
    None,
}

And use it like this:

fn get_e0(&self, id: Id) -> E0<&A, &B> {
    if let Some(a) = self.a_vec.iter().find(|x| x.id==id) {
        E0::A(a)
    } else if let Some(b) = self.b_vec.iter().find(|x| x.id==id) {
        E0::B(b)
    } else {
        E0::None
    }
}

It is most of the time more idiomatic to express such situations using the type system instead of panicking.