Implement IntoIterator and use it for an example

ceno_zkvm/src/uint.rs

@@ -18,7 +18,7 @@ use constants::BYTE_BIT_WIDTH;
 use ff::Field;
 use ff_ext::ExtensionField;
 use goldilocks::SmallField;
-use itertools::Itertools;
+use itertools::{Itertools, enumerate};
 use std::{
     borrow::Cow,
     mem::{self, MaybeUninit},
@@ -35,15 +35,36 @@ pub enum UintLimb<E: ExtensionField> {
     Expression(Vec<Expression<E>>),
 }
 
-impl<E: ExtensionField> UintLimb<E> {
-    pub fn iter(&self) -> impl Iterator<Item = &WitIn> {
+impl<E: ExtensionField> IntoIterator for UintLimb<E> {
+    type Item = WitIn;
+    type IntoIter = std::vec::IntoIter<WitIn>;
+
+    fn into_iter(self) -> Self::IntoIter {
         match self {
-            UintLimb::WitIn(vec) => vec.iter(),
+            UintLimb::WitIn(wits) => wits.into_iter(),
             _ => unimplemented!(),
         }
     }
 }
 
+impl<'a, E: ExtensionField> IntoIterator for &'a UintLimb<E> {
+    type Item = &'a WitIn;
+    type IntoIter = std::slice::Iter<'a, WitIn>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        match self {
+            UintLimb::WitIn(wits) => wits.iter(),
+            _ => unimplemented!(),
+        }
+    }
+}
+
+impl<E: ExtensionField> UintLimb<E> {
+    pub fn iter(&self) -> impl Iterator<Item = &WitIn> {
+        self.into_iter()
+    }
+}
+
 impl<E: ExtensionField> Index<usize> for UintLimb<E> {
     type Output = WitIn;
 
@@ -801,8 +822,8 @@ impl<'a, T: Into<u64> + From<u32> + Copy + Default> Value<'a, T> {
         let mut c_limbs = vec![0u16; num_limbs];
         let mut carries = vec![0u64; num_limbs];
         let mut tmp = vec![0u64; num_limbs];
-        a_limbs.iter().enumerate().for_each(|(i, &a_limb)| {
-            b_limbs.iter().enumerate().for_each(|(j, &b_limb)| {
+        enumerate(a_limbs).for_each(|(i, &a_limb)| {
+            enumerate(b_limbs).for_each(|(j, &b_limb)| {
                 let idx = i + j;
                 if idx < num_limbs {
                     tmp[idx] += a_limb as u64 * b_limb as u64;

ceno_zkvm/src/uint/arithmetic.rs

@@ -267,14 +267,12 @@ impl<const M: usize, const C: usize, E: ExtensionField> UIntLimbs<M, C, E> {
         rhs: &UIntLimbs<M, C, E>,
     ) -> Result<IsEqualConfig, ZKVMError> {
         let n_limbs = Self::NUM_LIMBS;
-        let (is_equal_per_limb, diff_inv_per_limb): (Vec<WitIn>, Vec<WitIn>) = self
-            .limbs
-            .iter()
-            .zip_eq(rhs.limbs.iter())
-            .map(|(a, b)| circuit_builder.is_equal(a.expr(), b.expr()))
-            .collect::<Result<Vec<(WitIn, WitIn)>, ZKVMError>>()?
-            .into_iter()
-            .unzip();
+        let (is_equal_per_limb, diff_inv_per_limb): (Vec<WitIn>, Vec<WitIn>) =
+            izip!(&self.limbs, &rhs.limbs)
+                .map(|(a, b)| circuit_builder.is_equal(a.expr(), b.expr()))
+                .collect::<Result<Vec<(WitIn, WitIn)>, ZKVMError>>()?
+                .into_iter()
+                .unzip();
 
         let sum_expr = is_equal_per_limb
             .iter()

ceno_zkvm/src/uint/logic.rs

@@ -18,7 +18,7 @@ impl<const M: usize, E: ExtensionField> UIntLimbs<M, 8, E> {
         b: &Self,
         c: &Self,
     ) -> Result<(), ZKVMError> {
-        for (a_byte, b_byte, c_byte) in izip!(a.limbs.iter(), b.limbs.iter(), c.limbs.iter()) {
+        for (a_byte, b_byte, c_byte) in izip!(&a.limbs, &b.limbs, &c.limbs) {
             cb.logic_u8(rom_type, a_byte.expr(), b_byte.expr(), c_byte.expr())?;
         }
         Ok(())