From 8754d5c15f710905bf82310d04eb0d9096b2b002 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Izabela=20O=C5=BCd=C5=BCe=C5=84ska?= <izabelaozdz@gmail.com>
Date: Fri, 7 Jun 2024 10:49:21 +0200
Subject: [PATCH] Add index as a parameter to each()

---
 .../circuits/lib/src/misc/bounded_vecs.nr     |  6 +--
 ethereum/circuits/lib/src/misc/bytes.nr       | 18 ++++-----
 ethereum/circuits/lib/src/misc/fragment.nr    | 37 ++++++++++---------
 ethereum/circuits/lib/src/misc/iterator.nr    |  2 +-
 4 files changed, 32 insertions(+), 31 deletions(-)

diff --git a/ethereum/circuits/lib/src/misc/bounded_vecs.nr b/ethereum/circuits/lib/src/misc/bounded_vecs.nr
index ffd7a81b..a7d74414 100644
--- 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
index 5cfefa03..7bce767d 100644
--- 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
index 5806bfea..a1a054e7 100644
--- 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
index 244cdeb7..40287832 100644
--- 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) -> ());
 }