Simplifying Rust macro rules for nested looping

Question

I have a simple macro with three very similar rules:

macro_rules! c {

($exp:expr, for $i:ident in $iter:expr) => (
    {
        let mut r = vec![];
        for $i in $iter {
            r.push($exp);
        }
        r
    }
);

($exp:expr, for $i:ident in $iter:expr, for $i2:ident in $iter2:expr) => (
    {
        let mut r = vec![];
        for $i2 in $iter2 {
            for $i in $iter {
                r.push($exp);
            }
        }
        r
    }
);

($exp:expr, for $i:ident in $iter:expr, for $i2:ident in $iter2:expr, for $i3:ident in $iter3:expr) => (
    {
        let mut r = vec![];
        for $i in $iter {
            for $i2 in $iter2 {
                for $i3 in $iter3 {
                    r.push($exp);
                }
            }
        }
        r
    }
);

}

Each rule differs from the others by the number of for $i:ident in $iter:exp patterns being matched. The logic is similarly the same.

Is there a way to simplify these rules into one using repetition patterns such as $(...)* or $(...)+ and still be able to express the nested looping in the macro logic?

Playground link

Answer 1

You can use a recursive TT (token tree) munching macro:

macro_rules! c {
    (@loop $v:ident, $exp:expr, for $i:ident in $iter:expr) => (
        for $i in $iter {
            $v.push($exp);
        }
    );

    (@loop $v:ident, $exp:expr, for $i:ident in $iter:expr, $($tail:tt)*) => (
        for $i in $iter {
            c!(@loop $v, $exp, $($tail)*);
        }
    );

    ($exp:expr, $(for $i:ident in $iter:expr),*) => (
        {
            let mut r = vec![];
            c!(@loop r, $exp, $(for $i in $iter),*);
            r
        }
    );
}

The rules labelled with @loop do all the work.

A TT munching recursive macro is very similar to a recursive function. At every invocation, it processes (munches) only a portion of the input, generates intermediate output, and sends the remaining "unmunched" input tail to another macro invocation. Eventually, the input is small enough to not require any more macro invocations and reaches the base case at which the recursion is terminated.

Here, the recursive @loop rule captures a single token tree matching to for $i:ident in $iter:expr, and stores the remaining input (other such for $i in $iter expressions) in a $($tail:tt)*. The macro rule then generates the loop for the captured for $i in $iter expression and generates the loop body by invoking the same rule with the unmunched input ($($tail)*).

Eventually, $($tail)* contains only one token tree that can be matched to for $i:ident in $iter:expr. In that case, the base case @loop rule is called, generating the innermost loop which pushes the expression onto the Vec.

This macro should work for an arbitrary number of for $i in $iter expressions as long as it stays within the macro recursion limit. If you do find yourself up against the recursion limit, the number of recursive invocations can be reduced by processing two for $i:ident in $iter:expr expressions at once in the recursive @loop rule.

Rust Playground