Applying function to all enum variants

1

I would like to perform an operation all variants of an enum.
Currently this requires quite long winded code like


enum Number{
    One(i64),
    Two(i64),
    Three(i64),
    Four(i64)
}

fn do_stuff() {
       match self {
            Number::One(n) =>
                blah(n),
            Number::Two(n) =>
                blah(n),
            Number::Three(n) =>
                blah(n),
            Number::Four(n) =>
                blah(n)
        }
}

In this case, blah would be a generic method which handles all of the
variants of Number. Now I know this can't be avoided in rust, but I
was wondering whether it could at least be written once. My thought
was something like this:

impl Number {

    fn apply<F, T, R>(&self, f: F) -> R
        where F: Fn(impl T) -> R,
    {
        match self {
            Number::One(n) =>
                f(n),
            Number::Two(n) =>
                f(n),
            Number::Three(n) =>
                f(n),
            Number::Four(n) =>
                f(n)
        }
    }
}

The point here is that I could write this once when I define my data
structure (Number). The in downstream code, I should be able to do
something like:

trait Foo {

    fn frob(&self) -> T;
}

impl Foo for i64 {

    fn from(&self) -> Result<(),String>{
        Ok(())
    }
}

fn do_foo(n:Number) -> Result<(), String> {
    n.apply(|n| n.from())

}

Of course, I would still need to write the underlying implementation
for each of the types wrapped in the enum (which wouldn't be i64 and
might not all be the same), but I wouldn't have to write the enum
dispatch functionality everytime.

Nice idea, but fails because I can't do impl T where T is some
trait because Rust is unhappy that T is generic.

Is there a solution here?

2

Is this acceptable?

enum Number {
    One(i64),
    Two(i64),
    Three(i64),
    Four(i64)
}

fn blah(_: i64) {}

impl Number {
    fn do_stuff(&self) {
        use Number::*;

        match self {
            One(n) | Two(n) | Three(n) | Four(n) => blah(*n)
        }
    }
}

fn main() {}

I would still need to write the underlying implementation
for each of the types wrapped in the enum (which wouldn’t be i64 and
might not all be the same),

I missed this at first, so nevermind...

3

The either crate has an internal either! macro which is used to reduce a lot of manual repetition, like in the implementation of Iterator.

4

I frequently have functions like

impl Number {
    fn fold<F, R>(&self, f: F) -> R
        where F: FnOnce(i64) -> R,
    {
        match *self {
            Number::One(n) |
            Number::Two(n) |
            Number::Three(n) |
            Number::Four(n) => f(n),
        }
    }

    fn map<F, R>(&self, f: F) -> Self
        where F: FnOnce(i64) -> i64,
    {
        match *self {
            Number::One(n) => Number::One(f(n)),
            Number::Two(n) => Number::Two(f(n)),
            Number::Three(n) => Number::Three(f(n)),
            Number::Four(n) => Number::Four(f(n)),
        }
    }
}

...unfortunately...

Now that's a rough situation to be in! Closures in rust unfortunately cannot be generic.

I will present a number of possible solutions, but no silver bullets.

Possible solution 1: Dynamic dispatch

This is the easiest solution, but it's not always feasible.

If the trait is object safe, you can just take F: Fn(&dyn Trait) -> R.

Object safety is a difficult condition to maintain. If your trait frequently changes to accommodate new requirements (and you value your sanity), this may not be suitable.

Possible solution 2: Manual generic closures

This is the most general solution, and probably the most correct solution for public APIs, but it requires a lot of effort. Before going this route you should ask yourself: Is it worth it?

Basically, you can define your own trait to use in place of the Fn traits, and then explicitly define a struct and impl the trait when you need to use it. For instance:

enum Heterogenous {
    Red(RedThing),
    Blue(BlueThing),
}

trait Thing {
    fn number(&self) -> u32;
}
impl Thing for RedThing { ... }
impl Thing for BlueThing { ... }

// Dedicated Fn replacement for Things
trait ThingFunc {
    type Output;

    fn call(self, thing: &impl Thing) -> Self::Output;
}

impl Heterogenous {
    fn fold<F: ThingFunc>(&self, f: F) -> F::Output {
        match self {
            Heterogenous::Red(x) => f(x),
            Heterogenous::Blue(x) => f(x),
        }
    }
}

Using it will be painful. You need to define a struct at every call site.

fn example(h: Heterogenous) {
    struct Folder {
        // struct fields are whatever your closure would have been
        // closed over from the environment
        data: Vec<i64>,
    }

    impl ThingFunc for Folder {
        type Output = ();

        fn call(self, thing: &impl Thing) {
            println!("closed over data: {:?}", data);
            println!("thing number: {}", thing.number());
        }
    }

    let data = vec![1, 5, 6];
    h.fold(Folder { data })
}

All I can say is you gotta weigh the costs against the benefits for this one.

There are also variations on this pattern, such as syn's AST folding and visitor traits which have a bajillion methods for specific node types because it turns out that's an incredibly effective way to organize code that recursively walks the AST.

Possible solution 3: Generate many closures using a macro

This is the hackiest of the solutions, best suited for private use and for quickly making parts of your implementation DRY without too much maintenance expense. It won't work if you need to mutably borrow or move something out of your closure's environment.

impl Heterogenous {
    fn fold(
        &self,
        red_f: impl Fn(&RedThing) -> R),
        blue_f: impl Fn(&BlueThing) -> R),
    ) -> R {
        match self {
            Heterogenous::Red(x) => red_f(x),
            Heterogenous::Blue(x) => blue_f(x),
        }
    }
}

Yep, you make your function take multiple closures. How can you use it while keeping DRY? Why, using macros, of course!

fn example(h: Heterogenous) {
    let data = vec![1, 5, 6];

    // using a macro to simulate a generic closure
    macro_rules! f {
        () => {
            |thing| {
                println!("closed over data: {:?}", data);
                println!("thing number: {}", thing.number());
            }
        }
    }

    h.fold(f!(), f!()) // call the macro twice to create two identical closures
}
5

I've had to deal with this problem a lot when writing operating on some sort of AST. The best solution I found was to write a defer!() macro which will pattern match on each variant and apply some operation to it.

I can't find the original macro definition, but the idea was you could write something like this:

enum Foo {
    Bar(i32),
    Baz(f64),
    Quux(String),  
    SomethingElseWhichIsIgnored,
}

fn main() {
    let foo = Foo::Quux(String::from("Hello World"));
    
    defer!(foo, Bar | Baz | Quux => |item| println!("{}", item));   
}

And it'd get expanded to

fn main() {
      let foo = Foo::Quux(String::from("Hello World"));

    match foo {
        Foo::Bar(ref item) => println!("{}", item),
        Foo::Baz(ref item) => println!("{}", item),
        Foo::Quux(ref item) => println!("{}", item),
        _ => {},
     }
}
1 Like
6

I think dynamic dispatch sounds straight forward, but it seems a shame to pick a less performant solution to save typing. The other two solutions are interesting but I worry about their complexity. My code is long, but at least it's simple!

7

Yeah, I am not quite doing an AST but it's very similar. I think you have the right solution, actually, or at least nearly. Although I have shown four variants, one of my enums has, in reality, about 50 variants. So, the solution is not a defer macro, but one bespoke for the enum

macro_rules! on_number {
    ($on:ident, $with:ident, $body:tt) => {
        match $on {
            Number::One(n) =>
                $body
            Number::Two(n) =>
               $body 
            etc...
        }
    }
}

fn main() {
    let x = Number::One(10);

    on_number!{x, n,
               {
                   fred("{}",n);
               }
    }
}

In my case, this is even easier because some of my enums are generated from using a macro already, so I should be able to have these macros generate another macro.

1 Like