More concise way to convert an enum variant to another in this code

I'm curious if there is a more concise way to write Buzz::wuzz in the below code. Ignore the simplicity of the types as this code is a simple analog of actual code.

I'm open to certain changes; however Foo and Bar must be treated as opaque types. Buzz must contain an Option. The signatures of the functions must remain as such. The general idea is that Foo is the final target that Buzz will return. Bar must be used to add Foos; however it's "expensive" and we want to optimize for the situation that we can keep a Foo without ever converting it to a Bar.

use core::mem;
struct Foo(usize);
impl Foo {
    const fn new() -> Self {
        Self(0)
    }
}
struct Bar(usize);
impl Bar {
    const fn new() -> Self {
        Self(0)
    }
    const fn add(&mut self, foo: &Foo) {
        self.0 = self.0.checked_add(foo.0).expect("overflow");
    }
    const fn into_foo(self) -> Foo {
        Foo(self.0)
    }
}
enum Fizz {
    Foo(Foo),
    Bar(Bar),
}
struct Buzz(Option<Fizz>);
impl Buzz {
    const fn new() -> Self {
        Self(Some(Fizz::Foo(Foo::new())))
    }
    /// Must never be called after [`Self::to_foo`] is called.
    fn wuzz(&mut self, foo: Foo) {
        enum Action {
            FooToFoo,
            FooToBar,
            None,
        }
        if foo.0 != 0 {
            let s = self
                .0
                .as_mut()
                .unwrap_or_else(|| panic!("Buzz::wuzz must not be called after Buzz::to_foo"));
            match match *s {
                Fizz::Foo(ref f) => {
                    if f.0 == 0 {
                        Action::FooToFoo
                    } else {
                        Action::FooToBar
                    }
                }
                Fizz::Bar(ref mut b) => {
                    b.add(&foo);
                    Action::None
                }
            } {
                Action::FooToFoo => {
                    if let Fizz::Foo(ref mut f) = *s {
                        *f = foo;
                    } else {
                        unreachable!("bug in code");
                    }
                }
                Action::FooToBar => {
                    if let Fizz::Foo(f) = mem::replace(s, Fizz::Bar(Bar::new())) {
                        if let Fizz::Bar(ref mut b) = *s {
                            b.add(&f);
                            b.add(&foo);
                        } else {
                            unreachable!("bug in code");
                        }
                    } else {
                        unreachable!("bug in core::mem::replace");
                    }
                }
                Action::None => (),
            }
        }
    }
    /// Must be called at most once.
    fn to_foo(&mut self) -> Foo {
        match self
            .0
            .take()
            .unwrap_or_else(|| panic!("Buzz::to_foo must not be called more than once"))
        {
            Fizz::Foo(f) => f,
            Fizz::Bar(b) => b.into_foo(),
        }
    }
}

If I understand, you want this operation:

impl Fizz {
    const fn combine_foo(&mut self, foo: Foo) {
        match self {
            Self::Foo(cur @ Foo(0)) => *cur = foo,
            Self::Foo(cur) => {
                let mut bar = Bar::new();
                bar.add(cur);
                bar.add(&foo);
                *self = Self::Bar(bar);
            }
            Self::Bar(bar) => {
                bar.add(&foo);
            }
        }
    }
}

And then you can

impl Buzz {
    /// Must never be called after [`Self::to_foo`] is called.
    fn wuzz(&mut self, foo: Foo) {
        if foo.0 != 0 {
            let Some(ref mut fizz) = self.0 else {
                panic!("Buzz::wuzz must not be called after Buzz::to_foo");
            };
            
            fizz.combine_foo(foo);
        }
    }
}
3 Likes