chore: introduce for_loop

noir_stdlib/src/collections/bounded_vec.nr

@@ -320,19 +320,12 @@ impl<T, let MaxLen: u32> BoundedVec<T, MaxLen> {
         let new_len = self.len + append_len;
         assert(new_len <= MaxLen, "extend_from_bounded_vec out of bounds");
 
-        if is_unconstrained() {
-            for i in 0..append_len {
-                self.storage[self.len + i] = vec.get_unchecked(i);
-            }
-        } else {
-            let mut exceeded_len = false;
-            for i in 0..Len {
-                exceeded_len |= i == append_len;
-                if !exceeded_len {
-                    self.storage[self.len + i] = vec.get_unchecked(i);
-                }
-            }
-        }
+        let storage_ref = &mut self.storage;
+        for_loop::<Len, _>(
+            0,
+            append_len,
+            |i| { storage_ref[self.len + i] = vec.get_unchecked(i); },
+        );
         self.len = new_len;
     }
 
@@ -395,20 +388,7 @@ impl<T, let MaxLen: u32> BoundedVec<T, MaxLen> {
     /// ```
     pub fn any<Env>(self, predicate: fn[Env](T) -> bool) -> bool {
         let mut ret = false;
-        if is_unconstrained() {
-            for i in 0..self.len {
-                ret |= predicate(self.storage[i]);
-            }
-        } else {
-            let mut ret = false;
-            let mut exceeded_len = false;
-            for i in 0..MaxLen {
-                exceeded_len |= i == self.len;
-                if !exceeded_len {
-                    ret |= predicate(self.storage[i]);
-                }
-            }
-        }
+        for_loop::<MaxLen, _>(0, self.len, |i| { ret |= predicate(self.storage[i]); });
         ret
     }
 
@@ -426,19 +406,12 @@ impl<T, let MaxLen: u32> BoundedVec<T, MaxLen> {
     pub fn map<U, Env>(self, f: fn[Env](T) -> U) -> BoundedVec<U, MaxLen> {
         let mut ret = BoundedVec::new();
         ret.len = self.len();
-
-        if is_unconstrained() {
-            for i in 0..self.len() {
-                ret.storage[i] = f(self.get_unchecked(i));
-            }
-        } else {
-            for i in 0..MaxLen {
-                if i < self.len() {
-                    ret.storage[i] = f(self.get_unchecked(i));
-                }
-            }
-        }
-
+        let storage_ref = &mut ret.storage;
+        for_loop::<MaxLen, _>(
+            0,
+            self.len(),
+            |i| { storage_ref[i] = f(self.get_unchecked(i)); },
+        );
         ret
     }
 

noir_stdlib/src/collections/map.nr

@@ -1,6 +1,7 @@
 use crate::cmp::Eq;
 use crate::collections::bounded_vec::BoundedVec;
 use crate::default::Default;
+use crate::for_loop;
 use crate::hash::{BuildHasher, Hash, Hasher};
 use crate::option::Option;
 
@@ -311,14 +312,17 @@ impl<K, V, let N: u32, B> HashMap<K, V, N, B> {
         // docs:end:iter_mut
         let mut entries = self.entries();
         let mut new_map = HashMap::with_hasher(self._build_hasher);
+        let new_map_ref = &mut new_map;
 
-        for i in 0..N {
-            if i < self._len {
+        for_loop::<N, _>(
+            0,
+            self._len,
+            |i| {
                 let entry = entries.get_unchecked(i);
                 let (key, value) = f(entry.0, entry.1);
-                new_map.insert(key, value);
-            }
-        }
+                new_map_ref.insert(key, value);
+            },
+        );
 
         self._table = new_map._table;
     }
@@ -349,14 +353,17 @@ impl<K, V, let N: u32, B> HashMap<K, V, N, B> {
         // docs:end:iter_keys_mut
         let mut entries = self.entries();
         let mut new_map = HashMap::with_hasher(self._build_hasher);
+        let new_map_ref = &mut new_map;
 
-        for i in 0..N {
-            if i < self._len {
+        for_loop::<N, _>(
+            0,
+            self._len,
+            |i| {
                 let entry = entries.get_unchecked(i);
                 let (key, value) = (f(entry.0), entry.1);
-                new_map.insert(key, value);
-            }
-        }
+                new_map_ref.insert(key, value);
+            },
+        );
 
         self._table = new_map._table;
     }

noir_stdlib/src/hash/keccak.nr

@@ -1,4 +1,4 @@
-use crate::runtime::is_unconstrained;
+use crate::{for_loop, runtime::is_unconstrained};
 
 global BLOCK_SIZE_IN_BYTES: u32 = 136; //(1600 - BITS * 2) / WORD_SIZE;
 global WORD_SIZE: u32 = 8; // Limbs are made up of u64s so 8 bytes each.
@@ -15,17 +15,8 @@ pub(crate) fn keccak256<let N: u32>(input: [u8; N], message_size: u32) -> [u8; 3
     // Copy input to block bytes. For that we'll need at least input bytes (N)
     // but we want it to be padded to a multiple of BLOCK_SIZE_IN_BYTES.
     let mut block_bytes = [0; ((N / BLOCK_SIZE_IN_BYTES) + 1) * BLOCK_SIZE_IN_BYTES];
-    if is_unconstrained() {
-        for i in 0..message_size {
-            block_bytes[i] = input[i];
-        }
-    } else {
-        for i in 0..N {
-            if i < message_size {
-                block_bytes[i] = input[i];
-            }
-        }
-    }
+    let block_bytes_ref = &mut block_bytes;
+    for_loop::<N, _>(0, message_size, |i| { block_bytes_ref[i] = input[i]; });
 
     //1. format_input_lanes
     let max_blocks = (N + BLOCK_SIZE_IN_BYTES) / BLOCK_SIZE_IN_BYTES;

noir_stdlib/src/hash/poseidon2.nr

@@ -1,4 +1,5 @@
 use crate::default::Default;
+use crate::for_loop;
 use crate::hash::Hasher;
 
 comptime global RATE: u32 = 3;
@@ -25,13 +26,10 @@ impl Poseidon2 {
 
     fn perform_duplex(&mut self) {
         // add the cache into sponge state
-        for i in 0..RATE {
-            // We effectively zero-pad the cache by only adding to the state
-            // cache that is less than the specified `cache_size`
-            if i < self.cache_size {
-                self.state[i] += self.cache[i];
-            }
-        }
+        // We effectively zero-pad the cache by only adding to the state
+        // cache that is less than the specified `cache_size`
+        let state_ref = &mut self.state;
+        for_loop::<RATE, _>(0, self.cache_size, |i| { state_ref[i] += self.cache[i]; });
         self.state = crate::hash::poseidon2_permutation(self.state, 4);
     }
 
@@ -67,11 +65,8 @@ impl Poseidon2 {
         let two_pow_64 = 18446744073709551616;
         let iv: Field = (in_len as Field) * two_pow_64;
         let mut sponge = Poseidon2::new(iv);
-        for i in 0..input.len() {
-            if i < in_len {
-                sponge.absorb(input[i]);
-            }
-        }
+        let sponge_ref = &mut sponge;
+        for_loop::<N, _>(0, in_len, |i| { sponge_ref.absorb(input[i]); });
 
         // In the case where the hash preimage is variable-length, we append `1` to the end of the input, to distinguish
         // from fixed-length hashes. (the combination of this additional field element + the hash IV ensures

noir_stdlib/src/lib.nr

@@ -121,6 +121,32 @@ where
 #[builtin(as_witness)]
 pub fn as_witness(x: Field) {}
 
+/// A for loop like this:
+///
+/// ```noir
+/// for i in start..min(Max, end) {
+///     f(i);
+/// }
+/// ```
+///
+/// that is implemented using constrained and (an optimized) unconstrained versions.
+pub(crate) fn for_loop<let Max: u32, Env>(start: u32, end: u32, f: fn[Env](u32) -> ()) {
+    if crate::runtime::is_unconstrained() {
+        assert(end <= Max);
+        for i in start..end {
+            f(i);
+        }
+    } else {
+        let mut exceeded_len = false;
+        for i in start..Max {
+            exceeded_len |= i == end;
+            if !exceeded_len {
+                f(i);
+            }
+        }
+    }
+}
+
 mod tests {
     #[test(should_fail_with = "custom message")]
     fn test_static_assert_custom_message() {