Add Field::shifted_powers and some iterator niceties (#1599)

* Add Field::shifted_powers and iterator niceties

* Remove comment
This commit is contained in:
Gio 2024-06-13 09:53:00 -06:00 committed by GitHub
parent ed6f452294
commit 7ba3c1f298
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -427,9 +427,13 @@ pub trait Field:
}
fn powers(&self) -> Powers<Self> {
self.shifted_powers(Self::ONE)
}
fn shifted_powers(&self, start: Self) -> Powers<Self> {
Powers {
base: *self,
current: Self::ONE,
current: start,
}
}
@ -563,6 +567,7 @@ pub trait PrimeField64: PrimeField + Field64 {
}
/// An iterator over the powers of a certain base element `b`: `b^0, b^1, b^2, ...`.
#[must_use = "iterators are lazy and do nothing unless consumed"]
#[derive(Clone, Debug)]
pub struct Powers<F: Field> {
base: F,
@ -577,6 +582,24 @@ impl<F: Field> Iterator for Powers<F> {
self.current *= self.base;
Some(result)
}
fn size_hint(&self) -> (usize, Option<usize>) {
(usize::MAX, None)
}
fn nth(&mut self, n: usize) -> Option<F> {
let result = self.current * self.base.exp_u64(n.try_into().unwrap());
self.current = result * self.base;
Some(result)
}
fn last(self) -> Option<F> {
panic!("called `Iterator::last()` on an infinite sequence")
}
fn count(self) -> usize {
panic!("called `Iterator::count()` on an infinite sequence")
}
}
impl<F: Field> Powers<F> {
@ -592,3 +615,26 @@ impl<F: Field> Powers<F> {
}
}
}
#[cfg(test)]
mod tests {
use super::Field;
use crate::goldilocks_field::GoldilocksField;
#[test]
fn test_powers_nth() {
type F = GoldilocksField;
const N: usize = 10;
let powers_of_two: Vec<F> = F::TWO.powers().take(N).collect();
for (n, &expect) in powers_of_two.iter().enumerate() {
let mut iter = F::TWO.powers();
assert_eq!(iter.nth(n), Some(expect));
for &expect_next in &powers_of_two[n + 1..] {
assert_eq!(iter.next(), Some(expect_next));
}
}
}
}