Add Field::shifted_powers and some iterator niceties

[gio256]

Adds the Field::shifted_powers method from Plonky3 along with a few additions to the Powers iterator implementation.

Please pick and choose what's useful.

<Powers<F> as Iterator>::nth in particular might require some tweaking:

• It currently panics if n overflows u64. This seems at least as good as the current behavior of hanging trying to compute u64::MAX field multiplications. But, maybe it's worth the performance hit to use Field::exp_biguint.
• The current implementation incurs one extra multiplication in the case of n = 1 (e.g. generator.powers().skip(1)). It's possible this is the most commonly expected case and nth should just be removed altogether.

[wborgeaud]

It's possible this is the most commonly expected case and nth should just be removed altogether

Why do you think this is the most commonly expected case? In general, I'm not sure where nth could be useful. Do you have use cases in mind?

[gio256]

The reason for the PR as a whole was to hopefully make it easier to replace as many instances of F::from_canonical_usize(1 << n) as possible (for safety and readability reasons). Originally, I thought nth would be useful for bit shifts for example.

// n is the (hypothetical) number of bits to shift by.
// f_n is a flag indicating whether this is really the n to shift by.
// All bits in the output with index < n are 0.
// m in {n..32} is the LE index of a possibly nonzero bit in the output.
// We want to line up the nth bit of the output with the 0th bit of the input,
// so we take `output_bits[m] = input_bits[m - n]`.
let shift_left: P = lv
    .shift_by_indices
    .into_iter()
    .enumerate()
    .flat_map(|(n, f_n)| {
        (n..32)
            .zip(P::Scalar::TWO.powers().skip(n))
            .map(move |(m, base)| f_n * lv.input_bits[m - n] * base)
    })
    .sum();

But, this is more efficiently done with Field::shifted_powers anyway. nth is used under the hood by iterator adapters like skip and step_by, so if those seem useful in the context of powers() than maybe it's worth keeping. Otherwise, I'm happy to remove it.

[gio256]

Thinking about this more, I gave a misleading example as using Field::shifted_powers in the code above doesn't actually accomplish anything.

let shift_left: P = lv
    .shift_by_indices
    .into_iter()
    .enumerate()
    .flat_map(|(n, f_n)| {
        (n..32)
            .zip(P::Scalar::TWO.powers_shifted(P::Scalar::from_canonical_u32(1 << n)))
            .map(move |(m, base)| f_n * lv.input_bits[m - n] * base)
    })
    .sum();

It's a small thing, but this still requires verifying that at no point in the loop does 1 << n exceed u32::MAX or the order of the field. For that reason, I think the version using Iterator::skip is slightly better if the overhead is negligible. So, I'll update my (soft) pitch for adding nth to the Iterator implementation:

• It makes the overhead of using Field::powers more intuitive and predictable, in that advancing the iterator by some large n costs one exponentiation rather than n muls. This enables the use of Field::powers to (efficiently) replace a mildly dangerous pattern.

• On the other hand, this means using exponentiation even if n = 1 (e.g. generator.powers().skip(1)). It's also far from obvious that P::Scalar::TWO.powers().skip(n) translates to (2^n, 2^n+1, ...), so you could argue that even this use case is contrived.

[matthiasgoergens]

Nice, I miss shifted_powers!

[wborgeaud]

Looks good to me!

But, this is more efficiently done with Field::shifted_powers anyway. nth is used under the hood by iterator adapters like skip and step_by, so if those seem useful in the context of powers() than maybe it's worth keeping. Otherwise, I'm happy to remove it.

For step_by we could actually do better by returning Powers {base: self.base.exp_u64(step as u64), current: self.current}, but I don't think it's possible to override this method and panicking seems a bit extreme.

I think we can keep your implementation of nth though. It's much faster for large n and shouldn't be called that often for n=1.

[gio256]

I merged upstream (#1601) so the tests would run, but let me know if I should squash.

Otherwise, this should be ready to go. Thanks for the feedback.

[Nashtare]

LGTM too!

I merged upstream (#1601) so the tests would run, but let me know if I should squash.

No problem about this, PRs are squashed into main anyway!