Compile-time `for loop` -like macro

help
1

Hi.
I'm writing a macro so that I can be sure that my code will be unrolled and hence, I can have it defined inside a const fn.

I'd like to build a macro that generates the following (notice we're already inside another macro):

let (r0, carry) = mac(0, val[0], $r2.0[0], 0);
let (r1, carry) = mac(0, val[0], $r2.0[1], carry);
let (r2, carry) = mac(0, val[0], $r2.0[2], carry);
let (r3, r4) = mac(0, val[0], $r2.0[3], carry);

let (r1, carry) = mac(r1, val[1], $r2.0[0], 0);
let (r2, carry) = mac(r2, val[1], $r2.0[1], carry);
let (r3, carry) = mac(r3, val[1], $r2.0[2], carry);
let (r4, r5) = mac(r4, val[1], $r2.0[3], carry);

let (r2, carry) = mac(r2, val[2], $r2.0[0], 0);
let (r3, carry) = mac(r3, val[2], $r2.0[1], carry);
let (r4, carry) = mac(r4, val[2], $r2.0[2], carry);
let (r5, r6) = mac(r5, val[2], $r2.0[3], carry);

let (r3, carry) = mac(r3, val[3], $r2.0[0], 0);
let (r4, carry) = mac(r4, val[3], $r2.0[1], carry);
let (r5, carry) = mac(r5, val[3], $r2.0[2], carry);
let (r6, r7) = mac(r6, val[3], $r2.0[3], carry);

Where the number of rs and blocks varies depending on a top-level macro attribute called $num_limbs.

What I did?

#[macro_use]
use paste::paste;

#[macro_export]
macro_rules! mac_round {
    ($num_limbs:expr, $idx:expr, $val:expr, $r2:expr, $carry:expr) => {
        for (i, j) in ($idx..($num_limbs + $idx - 2)).into_iter().enumerate() {
            paste::paste! { let (crate::r!(j), $carry) = mac(crate::r!(j), $val[$idx], $r2.0[i], $carry); }
        }
        paste::paste! { let (crate::r!($num_limbs + $idx -1), crate::r!($num_limbs + $idx)) = mac(crate::r!($num_limbs + $idx -1), $val[$idx], $r2.0[$limb_len + $idx - 1], $carry); }
    };
}

#[macro_export]
macro_rules! mac0 {
    ( $num_limbs:expr, $val:expr, $r2:expr, $carry:expr) => {
        paste::paste! { let (crate::r!(i), carry) = mac(0, $val[0], $r2.0[0], 0); }
        for i in (1..$num_limbs).into_iter() {
            paste::paste! { let (crate::r!(i), carry) = mac(0, $val[0], $r2.0[i], $carry); }
        }
    };
}

#[macro_export]
macro_rules! mac {
    ($num_limbs:expr, $val:expr, $r2:expr) => {
        // for idx in 0..2 * num_limbs {
        //     paste::paste! { let r!(idx) = 0u64;}
        // }
        paste::paste! { let carry = 0u64; };
        crate::mac0!($num_limbs, $val, $r2, carry);

        for i in (1..$num_limbs).into_iter() {
            crate::mac_round!($num_limbs, i, $val, $r2, carry)
        }
    };
}

#[macro_export]
macro_rules! r {
    ($idx: expr) => {
        paste::paste! {[<r $idx>]}
    };
}

So far, I was expecting my for loops to be executed at compile time resulting on the pasting of the code X times.
But cargo-expand shows:

let carry = 0u64;
let (ri, carry) = mac(0, val[0], R2.0[0], 0);
for i in (1..4).into_iter() {
    let (ri, carry) = mac(0, val[0], R2.0[i], carry);
}
for i in (1..4).into_iter() {
    for (i, j) in (i..(4 + i - 2)).into_iter().enumerate() {
        let (rj, carry) = mac(rj, val[i], R2.0[i], carry);
    }
}

The question then is immediate. Is there any way (I can port the loops to be iterators) on which I can run this for loop on any other way (recursive macro that decreases/increases a number or whatever). So that I can print rust code at compile time based on a constant that is provided also at compile time to a macro?

2

That will be easiest using a procedural macro, probably. It's hard to perform value-level counting in declarative macros.

1 Like
3

You might be interested in the seq-macro crate.

4 Likes
4

That's really really close to what I need. Thank you so much for the suggestion!

The only issue I see is that iteration range is defined by variables of the top level macro.

Hence, If I try to operate with them as:

#[macro_export]
macro_rules! mac_round {
    ($num_limbs:expr, $idx:expr, $val:expr, $r2:expr, $carry:expr) => {
        seq_macro::seq!(j in $idx..$num_limbs + $idx - 2 {
            paste::paste! { let (crate::r!(j), $carry) = mac(crate::r!(j), $val[$idx], $r2.0[i], $carry); }
        });
        paste::paste! { let (crate::r!($num_limbs + $idx -1), crate::r!($num_limbs + $idx)) = mac(crate::r!($num_limbs + $idx -1), $val[$idx], $r2.0[$num_limbs + $idx - 1], $carry); }
    };
}

I get:

error: expected curly braces
   --> src/derive/mod.rs:193:47
    |
191 |    / macro_rules! mac_round {
192 |    |     ($num_limbs:expr, $idx:expr, $val:expr, $r2:expr, $carry:expr) => {
193 |    |         seq_macro::seq!(j in $idx..$num_limbs + $idx - 2 {
    |    |                                               ^
194 |    |             paste::paste! { let (crate::r!(j), $carry) = mac(crate::r!(j...
...      |
197 |    |     };
198 |    | }
    |    |_- in this expansion of `crate::mac_round!` (#3)
...
211 |    / macro_rules! mac {
212 |    |     ($num_limbs:expr, $val:expr, $r2:expr) => {{
213 |    |         // for idx in 0..2 * num_limbs {
214 |    |         //     paste::paste! { let r!(idx) = 0u64;}
...      |
220 |    |             crate::mac_round!($num_limbs, i, $val, $r2, carry);
    |    |             -------------------------------------------------- in this macro invocation (#3)
221 |    |         });
222 |    |     }};
223 |    | }
    |    |_- in this expansion of `crate::mac!` (#2)
    |
   ::: src/derive/field.rs:15:1
    |
15  | /    macro_rules! field_common {
16  | |        (
17  | |            $field:ident,
18  | |            $num_limbs:expr,
...   |
75  | |                        crate::mac!($num_limbs, val, $r2);
    | |                        --------------------------------- in this macro invocation (#2)
...   |
350 | |        };
351 | |    }
    | |____- in this expansion of `field_common!` (#1)
    |
   ::: src/secp256r1/fq.rs:128:1
    |
128 |    / field_common!(
129 |    |     Fq,
130 |    |     4,
131 |    |     MODULUS,
...      |
140 |    |     R3
141 |    | );
    |    |_- in this macro invocation (#1)

Do you know if there's anyway around this?
I tried to do:

#[macro_export]
macro_rules! mac_round {
    ($num_limbs:expr, $idx:expr, $val:expr, $r2:expr, $carry:expr) => {
        let range_max = $idx..$num_limbs + $idx - 2;
        seq_macro::seq!(j in $idx..range_max {
            paste::paste! { let (crate::r!(j), $carry) = mac(crate::r!(j), $val[$idx], $r2.0[i], $carry); }
        });
        paste::paste! { let (crate::r!($num_limbs + $idx -1), crate::r!($num_limbs + $idx)) = mac(crate::r!($num_limbs + $idx -1), $val[$idx], $r2.0[$num_limbs + $idx - 1], $carry); }
    };
}

But then the issue is that the range is expected to be a literal.

So I presume either I need to request the bounds/range as a macro input or there's no way around it?

5

Then problem is that $idx and $num_limbs are tokens, you can think of them as text, strings. They are not values you can sum. You can produce the text where you join them with the string +, but that won't result in summing them, it will result in a program containing a sum.

You have to either try to rewrite your logic in such a way that no math is needed when computing what the source code should look like, or use a proc macro (which, mind you, can still only work in limited cases, namely when your mac_round is called with literals as arguments)

2 Likes
6

Thank you so much for the help and the explanations. I think I might go for a build.rs approach to solve this.

7

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.