Search code examples
rustenumspredicate

How to group enum members but keep match exhaustivity check


In Rust is there a way to "group" enum members in a way that I can get a predicate function AND use that predicate in a match statement?

Let's say I have this enum:

enum Number {
  One,
  Two,
  Three,
  Four,
  Five,
  Six
}

And this predicate:

impl Number {
  fn is_prime(&self) -> bool {
    self == Number::Two || self == Number::Three || self == Number::Five
  }
}

Then my problem is that I have several match statements around my codebase that handles things differently if the number is prime:

match number {
  Number::Two | Number::Three | Number::Five => {/* do something for primes*/}
  Number::One => ...
  Number::Four => ...
  Number::Six => ... 
}

I would like to have only one source of truth of what is a prime. Of course I could add a predicate check before each match statement:

if number.is_prime() {
  /* do something for primes*/
} else {
  match number {
    Number::One => ...
    Number::Four => ...
    Number::Six => ...
    _ => {} // should never happen
  }
}

But then I lose the exhaustivity check. I need to add a catch-all arm, which doesn't provide me compile safety to ensure that anyone that later adds member to Number has to explicitly handle it in each match statement.

Ideally I would like to do something like this:

match number {
  is_prime() => {/* do something for primes*/}
  Number::One => ...
  Number::Four => ...
  Number::Six => ... 
}

Is there a way to achieve the same goal? I guess there should be a way to do it through a macro but I try to avoid macros as much as I can as they make the code less explicit.


Solution

  • You can make a macro:

    #[macro_export]
    macro_rules! primes {
        () => {
            $crate::Number::Two | $crate::Number::Three | $crate::Number::Five
        };
    }
    
    impl Number {
        pub fn factors(&self) -> u8 {
            use Number::*;
            match self {
                primes!() => 2,
                One => 1,
                Four => 3,
                Six => 4,
            }
        }
    }
    

    This macro assumes Number is publicly available in the crate root. Use the appropriate path if it is somewhere else. I don't think there's a way to avoid writing the whole path out for every variant since you can't have use statements in a pattern, although if it is extraordinarily long you could make another macro for just the path part. More about $crate here.

    On the upside, it does work as part of a larger pattern:

    primes!() | Number::One