How to group enum members but keep match exhaustivity check

65 Views Asked by At

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.

2

There are 2 best solutions below

1
drewtato On

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
0
cafce25 On

If all your groups are distinct, you can add another layer of enums:

enum Numbers {
    Prime(Prime),
    NonPrime(NonPrime),
}
enum Prime {
    Two,
    Three,
    Five,
}
enum NonPrime {
    One,
    Four,
    Six,
}

Then you can match on the outer layer:

use Number::*;
use Prime::*;
use NonPrime::*;
match number {
    Prime(_) => todo!("do something for primes"),
    NonPrime(One) => todo!("do something for one"),
    NonPrime(_) => todo!("do something for other non primes"),
}

is_prime becomes very trivial, too:

impl Number {
    fn is_prime(self) -> bool {
        matches!(self, Number::Prime(_))
    }
}