Conditional enum variants in Rust

Posted on:

Recently, I had to implement a set of types (AST nodes) in Rust that would be generic over boolean flags. More specifically, they had to be statically dependant on compile-time boolean parameters, where instances of the same type with false and true values of such parameter would be incompatible, as these flags would not only determine behaviour of impl but also affect e.g. which branches of enum are allowed. Pseudo code sample for demonstration purposes:

enum MyEnum<AllowB: bool> {
  A,
  B if AllowB, // this won't be allowed on MyEnum<false>
}

impl<AllowB: bool> MyEnum<AllowB> {
  fn is_a(&self) -> bool {
    match *self {
      MyEnum::A => true,
      _ => false,
    }
  }
}

impl MyEnum<true> {
  // this won't be allowed on MyEnum<false>, giving a helpful compile-time error
  fn is_b(&self) -> bool {
    match *self {
      MyEnum::B => true,
      _ => false,
    }
  }
}

The first thing that comes to mind is the RFC #2000 for const generics which would allow parametrising types over any compile-time constant values. Unfortunately, while the RFC has been accepted, it's not implemented even in nightly yet, and some implementation details are still in flux. Apart from that, it's unlikely to help with the enum variants requirement, even though it does ensure the other two.

What can we do instead? Well, why don't we try and implement booleans using the Rust type system? Then we will be able to pass them around, check and manipulate purely at compile-time.

First of all, let's define trait for our Bool which would determine what we can do with such boolean flags.

For the beginning, let's say we want to have just conversion to actual runtime boolean value. This will be done through an associated constant as our boolean variants are statically determined at compile-time.

pub trait Bool {
  const VALUE: bool;
}

Now, let's implement actual variants of our Bool as zero-size types implementing this trait:

pub struct True;

impl Bool for True {
  const VALUE: bool = true;
}

pub struct False;

impl Bool for False {
  const VALUE: bool = false;
}

Great! Now we can use these flags to implement our enum. Note that if someFlag is still not a valid Rust syntax, so instead we just store the flag inside of the variant payload:

pub enum MyEnum<AllowB: Bool> {
  A,
  B(AllowB), // TODO: disallow on MyEnum<False>
}

impl<AllowB: Bool> MyEnum<AllowB> {
  pub fn is_a(&self) -> bool {
    match *self {
      MyEnum::A => true,
      _ => false,
    }
  }
}

impl MyEnum<True> {
  // this won't be allowed on MyEnum<False>, giving a helpful compile-time error
  pub fn is_b(&self) -> bool {
    match *self {
      MyEnum::B(_) => true,
      _ => false,
    }
  }
}

If we try it, we can see that impl blocks work as expected:

[playground]

let e: MyEnum<False> = MyEnum::A;
println!("{}", e.is_a()); // prints "true"
println!("{}", e.is_b()); // error[E0599]: no method named `is_b` found for type `MyEnum<False>` in the current scope

Great! One issue remains though: while these flags are visible in the generated docs and should help the developer using our API understand that they should never use the B variant when AllowB is false-ish, this is not a strong enough guarantee which might still lead to accidental mistakes:

[playground]

let e: MyEnum<False> = MyEnum::B(False); // compiles while it shouldn't

Apart from that, the compiler doesn't know either that B variant simply can't exist on MyEnum<False> and so can't provide better optimisations. If we look at the generated code for, say, MyEnum<False>::is_a, we can see that it still tries to check the tag at the beginning of our enum at runtime even though it should be able to statically determine the result as true since only variant A can exist on such type!

[disassembly]

example::test_is_a:
  xor dil, 1
  mov eax, edi
  ret

To solve these two issues, we can use a trick with so-called uninhabited type:

pub enum Void {}

The trick is, while this type definition is perfectly valid in Rust, it has no variants, which means that you can't actually construct values of this type (unlike in C, Rust enums have much stronger guarantees and are not just glorified integer constants).

Moreover, such type "infects" all its usages - since you can't construct values of this type, any type containing Void becomes uninhabited too, and, in turn, any place where you're dealing with values of any such type, becomes unreachable and can be optimised out as dead code.

This is often used as a placeholder for generic params which we know can't be ever instantiated. For example, it's used by standard library as a ParseError for FromStr implementation on String - indeed, since string can be always converted to string, there is no way an error could be produced by this operation, which means we can provide a placeholder for Err part of Result which can't be actually constructed, and compiler will know that this operation is infallible.

Let's see how we can use this property for our original task. Assuming that our desired flag is always true-ish (and if not, we can implement unary not operator), we can make False an uninhabited type, which in turn will make any variants containing it uninhabited too. This might sound complicated, but the change is as simple as:

-pub struct False;
+pub enum False {}

Now, if we try this with our last example, we will immediately get a compile error as False is an enum now:

[playground]

error[E0423]: expected value, found enum `False`
  --> src/main.rs:42:36
   |
42 |   let _e: MyEnum<True> = MyEnum::B(False);
   |                                    ^^^^^ not a value

And, since it has no variants, there is nothing you can put on the right-hand side of this assignment to make it work. Neither you can assign MyEnum::B(True) to MyEnum<False> variable:

[playground]

42 |   let _e: MyEnum<False> = MyEnum::B(True);
   |                                     ^^^^ expected enum `False`, found struct `True`
   |
   = note: expected type `False`
              found type `True`

However, if you set the AllowB to True, you can start accepting MyEnum::B as expected:

[playground]

let _e: MyEnum<True> = MyEnum::B(True); // compiles

Apart from type safety, the advanced optimisations are now possible too. In particular, MyEnum<False>::is_a now just returns constant true as we wanted:

[disassembly]

example::test_is_a:
  mov al, 1
  ret

So we can see everything works as expected, even if with the little annoyance of having to store the True flag on the variant.

If we wish, we could now implement all the logical operators on top of these values (after all, Rust's type system is Turing-complete), but I wouldn't go that far in this post.

Hope this has been useful. Let me know on Twitter if you have any thoughts about the approach.


More posts: