CODEWITHSUNDEEP
genericsrustenumsrust-macros
MSR
MSR

Reputation: 2901

Writing const generic enum combinations using macro

Consider a toy struct with two const generic parameters:

pub struct Foo<const N: usize, const M: usize>([usize; N], [usize; M]);

impl<const N: usize, const M: usize> Foo<N, M> {
    pub fn bar(&self) -> usize {
        N * M
    }
}

Let's say all combinations of N and M between 1 and 5 are permitted, so that we could write the following enum:

pub enum FooEnum {
    Foo_1_1(Foo<1, 1>),
    Foo_1_2(Foo<1, 2>),
    Foo_2_1(Foo<2, 1>),
    Foo_2_2(Foo<2, 2>),
    // ... and so on.
}

impl FooEnum {
    pub fn bar(&self) -> usize {
        match self {
            Self::Foo_1_1(x) => x.bar(),
            Self::Foo_1_2(x) => x.bar(),
            Self::Foo_2_1(x) => x.bar(),
            Self::Foo_2_2(x) => x.bar(),
            // ... and so on.

        }
    }
}

My question is: Can we write a declarative macro to generate this, without manually writing out all the combinations? That is, something like impl_foo_enum!(1, 2, 3, 4, 5), rather than impl_foo_enum!(1;1, 1;2, 1;3, [...and so on]).

The latter macro I am able to write, using the paste crate:

macro_rules! impl_foo_enum {
    ($($n:literal;$m:literal),+) => {
        paste::paste! {
            pub enum FooEnum2 {
                $(
                    [<Foo _ $n _ $m>](Foo<$n, $m>)
                ),+
            }

            impl FooEnum2 {
                pub fn bar(&self) -> usize {
                    match self {
                        $(Self::[<Foo _ $n _ $m>](x) => x.bar()),+
                    }
                }
            }
        }
    }
}

impl_foo_enum!(1;1, 1;2, 2;1, 2;2);

(Playground)

To get the less tedious macro, there's a couple of related questions with helpful answers (1, 2) that I thought I could adapt, but in both cases a function call can be repeated inside a macro, which appears to simplify things. For example, using the approach in the first linked example, I started:

macro_rules! for_all_pairs {
    ($mac:ident: $($x:literal)*) => {
        for_all_pairs!(@inner $mac: $($x)*; $($x)*);
    };
    (@inner $mac:ident: ; $($x:literal)*) => {};
    (@inner $mac:ident: $head:literal $($tail:literal)*; $($x:literal)*) => {
        $(
            $mac!($head $x);
        )*
        for_all_pairs!(@inner $mac: $($tail)*; $($x)*);
    };
}

macro_rules! impl_foo_enum {
    ($n:literal $m:literal) => {
        paste::paste! { [<Foo _ $n _ $m>](Foo<$n, $m>) }
    }
}

pub enum FooEnum3 {
    for_all_pairs!(impl_foo_enum: 1 2)
}

(Playground)

Which does not compile, since the compiler does not expect a macro in the enum variant position (I believe).

(To be clear, I don't necessarily want to use the above for anything serious, I just ran into it while experimenting and got curious.)

Upvotes: 2

Views: 791

Answers (1)

phimuemue
phimuemue

Reputation: 36081

Here you go:

#![allow(non_camel_case_types)]
pub struct Foo<const N: usize, const M: usize>([usize; N], [usize; M]);
impl<const N: usize, const M: usize> Foo<N, M> {
    pub fn bar(&self) -> usize {
        N * M
    }
}

macro_rules! impl_foo_2{
    ($($n:literal)*) => {
        impl_foo_2!([] @orig($($n)*) ($($n)*) ($($n)*));
    };
    (
        [$(($n:literal $m:literal))*]
        @orig($($n_orig:literal)*)
        ($($n_unused:literal)*) ()
    ) => {
        paste::paste! {
            pub enum FooEnum2 {
                $([<Foo _ $n _ $m>](Foo<$n, $m>)),+
            }
            impl FooEnum2 {
                pub fn bar(&self) -> usize {
                    match self {
                        $(Self::[<Foo _ $n _ $m>](x) => x.bar()),+
                    }
                }
            }
        }
    };
    (
        [$($t:tt)*]
        @orig($($n_orig:literal)*)
        () ($m0:literal $($m:literal)*)
    ) => {
        impl_foo_2!(
            [$($t)*]
            @orig($($n_orig)*)
            ($($n_orig)*) ($($m)*)
        );
    };
    (
        [$($t:tt)*]
        @orig($($n_orig:literal)*)
        ($n0:literal $($n:literal)*) ($m0:literal $($m:literal)*)
    ) => {
        impl_foo_2!(
            [$($t)* ($n0 $m0)]
            @orig($($n_orig)*)
            ($($n)*) ($m0 $($m)*)
        );
    }
}

impl_foo_2!(1 2 3 4 5);

impl_foo_2 internally generates two identical copies of your number list. It then goes on and processes one m at a time, combining it with every n (it does so by repeatedly chopping off the first n). If the n-list is exhausted, it resets the n-list, and chops off the first m. All this is done until all n and m are exhausted.

The intermediate results are collected into the macro's first parameter which - at the end - is passed to your impl_foo_enum.

Upvotes: 3

Related Questions