Add index as a parameter to each()

vlayer-xyz · Jun 7, 2024 · 8754d5c · 8754d5c
1 parent d5b0cc4
commit 8754d5c
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 31 deletions.
diff --git a/ethereum/circuits/lib/src/misc/bounded_vecs.nr b/ethereum/circuits/lib/src/misc/bounded_vecs.nr
@@ -13,16 +13,16 @@ pub fn bounded_vec_map<U, T, N, Env>(
     f: fn[Env](T) -> U
 ) -> BoundedVec<U, N> {
     let result: &mut BoundedVec<U, N> = &mut BoundedVec::new();
-    bounded_vec.for_each(|x| result.push(f(x)));
+    bounded_vec.each(|x, _| result.push(f(x)));
 
     *result
 }
 
 impl<T, N> Iterator<T> for BoundedVec<T, N> {
-    fn for_each<Env>(self, f: fn[Env](T) -> ()) {
+    fn each<Env>(self, f: fn[Env](T, u64) -> ()) {
         for i in 0..N {
             if i < self.len {
-                f(self.get(i));
+                f(self.get(i), i);
             }
         }
     }

diff --git a/ethereum/circuits/lib/src/misc/bytes.nr b/ethereum/circuits/lib/src/misc/bytes.nr
@@ -12,17 +12,17 @@ pub fn bytes_to_nibbles<MAX_NIBBLE_LEN>(bytes: Fragment<MAX_NIBBLE_LEN, u8>) ->
         MAX_NIBBLE_LEN >= 2 * bytes.length, "Bytes implicit data length must be at least 2 times larger than bytes explicit length"
     );
 
-    let mut nibbles = Fragment::new_with_length(2 * bytes.length, [0; MAX_NIBBLE_LEN]);
-    for i in 0..MAX_NIBBLE_LEN {
-        if i < bytes.length {
-            let (hi, lo) = byte_to_nibbles(bytes.at(i));
-
-            nibbles.set(2 * i, hi);
-            nibbles.set(2 * i + 1, lo);
-        }
+    let nibbles = &mut Fragment::new_with_length(2 * bytes.length, [0; MAX_NIBBLE_LEN]);
+    bytes.each(
+        |byte, i| {
+        let (hi, lo) = byte_to_nibbles(byte);
+
+        nibbles.set(2 * i, hi);
+        nibbles.set(2 * i + 1, lo);
     }
+    );
 
-    nibbles
+    *nibbles
 }
 
 pub fn nibbles_to_byte(upper: u8, lower: u8) -> u8 {

diff --git a/ethereum/circuits/lib/src/misc/fragment.nr b/ethereum/circuits/lib/src/misc/fragment.nr
@@ -39,7 +39,7 @@ impl<MAX_DATA_LEN, T> Fragment<MAX_DATA_LEN, T> {
     pub fn to_bounded_vec<N>(self) -> BoundedVec<T, N> {
         assert(self.length <= N, "Fragment length exceeds BoundedVec max length");
         let bounded_vec: &mut BoundedVec<T, N> = &mut BoundedVec::new();
-        self.for_each(|x| bounded_vec.push(x));
+        self.each(|x, _| bounded_vec.push(x));
 
         *bounded_vec
     }
@@ -118,41 +118,42 @@ impl<MAX_DATA_LEN, T> Fragment<MAX_DATA_LEN, T> {
     }
 
     pub fn eq<OTHER_MAX_DATA_LEN>(self, other: Fragment<OTHER_MAX_DATA_LEN, T>) -> bool where T: Eq {
-        let mut res = true;
+        let res = &mut true;
         if (self.length != other.length) {
-            res = false;
+            *res = false;
         }
-        for i in 0..MAX_DATA_LEN {
-            if i < self.length & i < other.length {
-                res &= self.at(i) == other.at(i);
+        self.each(
+            |el, i| {
+            if i < other.length {
+                *res &= el == other.at(i);
             }
         }
-        res
+        );
+
+        *res
     }
 }
 
 impl<MAX_DATA_LEN, T> Eq for Fragment<MAX_DATA_LEN, T> where T: Eq {
     fn eq(self, other: Fragment<MAX_DATA_LEN, T>) -> bool {
-        let mut res = true;
+        let res = &mut true;
         if self.length != other.length {
-            res = false;
+            *res = false;
         } else {
-            for i in 0..MAX_DATA_LEN {
-                if i < self.length {
-                    res &= self.at(i) == other.at(i);
-                }
-            }
-        }
+            self.each(|el, i| {
+                *res &= el == other.at(i);
+            });
+        };
 
-        res
+        *res
     }
 }
 
 impl<MAX_DATA_LEN, T> Iterator<T> for Fragment<MAX_DATA_LEN, T> {
-    fn for_each<Env>(self, f: fn[Env](T) -> ()) {
+    fn each<Env>(self, f: fn[Env](T, u64) -> ()) {
         for i in 0..MAX_DATA_LEN {
             if i < self.length {
-                f(self.at(i));
+                f(self.at(i), i);
             }
         }
     }

diff --git a/ethereum/circuits/lib/src/misc/iterator.nr b/ethereum/circuits/lib/src/misc/iterator.nr
@@ -1,3 +1,3 @@
 trait Iterator<T> {
-    fn for_each<Env>(self, f: fn[Env](T) -> ());
+    fn each<Env>(self, f: fn[Env](T, u64) -> ());
 }