diff --git a/noir_stdlib/src/collections/bounded_vec.nr b/noir_stdlib/src/collections/bounded_vec.nr index c030544e791..c83ddd098e2 100644 --- a/noir_stdlib/src/collections/bounded_vec.nr +++ b/noir_stdlib/src/collections/bounded_vec.nr @@ -320,19 +320,12 @@ impl BoundedVec { 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::( + 0, + append_len, + |i| { storage_ref[self.len + i] = vec.get_unchecked(i); }, + ); self.len = new_len; } @@ -395,20 +388,7 @@ impl BoundedVec { /// ``` pub fn any(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::(0, self.len, |i| { ret |= predicate(self.storage[i]); }); ret } @@ -426,19 +406,12 @@ impl BoundedVec { pub fn map(self, f: fn[Env](T) -> U) -> BoundedVec { 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::( + 0, + self.len(), + |i| { storage_ref[i] = f(self.get_unchecked(i)); }, + ); ret } diff --git a/noir_stdlib/src/collections/map.nr b/noir_stdlib/src/collections/map.nr index bc0b80124db..aa4cabeef2e 100644 --- a/noir_stdlib/src/collections/map.nr +++ b/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 HashMap { // 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::( + 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 HashMap { // 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::( + 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; } diff --git a/noir_stdlib/src/hash/keccak.nr b/noir_stdlib/src/hash/keccak.nr index 93b366d1ec1..227ec9562b3 100644 --- a/noir_stdlib/src/hash/keccak.nr +++ b/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(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::(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; diff --git a/noir_stdlib/src/hash/poseidon2.nr b/noir_stdlib/src/hash/poseidon2.nr index ad4476f0d7c..47777669933 100644 --- a/noir_stdlib/src/hash/poseidon2.nr +++ b/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::(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::(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 diff --git a/noir_stdlib/src/lib.nr b/noir_stdlib/src/lib.nr index fb073516d29..9133368fa6f 100644 --- a/noir_stdlib/src/lib.nr +++ b/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(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() {