From 7ba3c1f298faddb99543b59fe24d898f092f78e8 Mon Sep 17 00:00:00 2001
From: Gio <102917377+gio256@users.noreply.github.com>
Date: Thu, 13 Jun 2024 09:53:00 -0600
Subject: [PATCH] Add `Field::shifted_powers` and some iterator niceties
 (#1599)

* Add Field::shifted_powers and iterator niceties

* Remove comment
---
 field/src/types.rs | 48 +++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/field/src/types.rs b/field/src/types.rs
index b2e43667..18c4b82d 100644
--- a/field/src/types.rs
+++ b/field/src/types.rs
@@ -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));
+            }
+        }
+    }
+}