Search code examples
rustenumsborrow-checker

Unwrap enum variant and return either owned value or reference


Suppose I have an enum with two variants, each of which wraps around a single value of a different type - and the enum is generic on these two types:

enum AorB<A, B> {
    A(A),
    B(B),
}

It frequently happens in my codebase that I know which variant I am receiving and I need to unwrap the inner value of the known type. I know this could raise some questions about the design, but let us put those aside. In any case, I would like to write convenience unwrapping methods in order to avoid having to match every time that situation arises.

That situation essentially takes two forms: I have a pointer to an AorB<A, B> and I want to get a pointer to the inner value; or when I am happy to consume the full AorB<A, B> object and want to get ownership of the inner value. This gives rise to four methods:

impl<A, B> AorB<A, B> {
    fn unwrap_a(self) -> A {
        match self {
            AorB::A(a) => a,
            _ => panic!(),
        }
    }

    fn unwrap_b(self) -> B {
        match self {
            AorB::B(b) => b,
            _ => panic!(),
        }
    }

    fn unwrap_a_ref(&self) -> &A {
        match self {
            AorB::A(a) => a,
            _ => panic!(),
        }
    }

    fn unwrap_b_ref(&self) -> &B {
        match self {
            AorB::B(b) => b,
            _ => panic!(),
        }
    }
}

My question is: is there a way to simplify these into four methods into two (one for each variant) and automatically handle the owned vs. reference cases elegantly?

Here are my current thoughts:

  1. I know there are traits such as AsRef and Borrow which are meant to flexibly handle various ownership situations, but I cannot see how to directly apply them to the aforementioned methods.

  2. A more or less obvious solution is to only leave the _ref methods in, and simply have the caller clone() the output whenever ownership is needed. I think if one then replaces e.unwrap_a() by e.unwrap_a_ref().clone(), e will be automatically deallocated by the borrow checker after the call as it is not used afterwards (otherwise e.unwrap_a() would not compile), so that memory usage is basically not duplicated by the switch to the _ref method. Is this correct? In any case, the call to clone() is slightly inconvenient on the caller side, and it does involve copying data in a way that unwrap_a did not need.

  3. A third option would be to only leave _ref methods in and add A: Deref and B: Deref in the AorB definition. This is very similar to 2, only the caller is spared the call to clone() (as far as I can tell). The slight issue with this approach is that A and B will be my own types in practice, and I would rather avoid implementing Deref for them if possible, as this can be a delicate trait.

Is there a better way than 2 or 3 above to achieve the method simplification I am seeking?


Solution

  • Option has a similar problem, it solves it by introducing a method as_ref which allows going from &Option<T> to Option<&T> I think it works here quite well, too:

    impl<A, B> AorB<A, B> {
        fn as_ref(&self) -> AorB<&A, &B> {
            match self {
                AorB::A(a) => AorB::A(a),
                AorB::B(b) => AorB::B(b),
            }
        }
    }
    

    then whenever you need to avoid consuming you just insert a call to as_ref:

    a_or_b.as_ref().unwrap_a();
    

    instead of calling .unwrap_a() directly.