diff --git a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/variable_merkle_tree.nr b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/variable_merkle_tree.nr index 701646552f2..22f13811692 100644 --- a/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/variable_merkle_tree.nr +++ b/noir-projects/noir-protocol-circuits/crates/types/src/merkle_tree/variable_merkle_tree.nr @@ -6,14 +6,54 @@ pub struct VariableMerkleTree { pub root: Field, } -unconstrained fn get_height(input: u32, start: u8) -> u8 { - let mut height = 0; +/// Recursively calculates the smallest power exponent `n` where `2^n` is greater than or equal +/// to the input value, starting from a given exponent. +/// +/// # Arguments +/// +/// * `input` - Target value to find the next power of 2 that exceeds or equals it +/// * `start` - Initial exponent to start checking from +/// +/// # Returns +/// +/// The smallest exponent `n` where `2^n >= input` +unconstrained fn get_next_power_exponent(input: u32, start: u8) -> u8 { + let mut next_power_exponent = 0; if input <= 2 << start { - height = start; + next_power_exponent = start; } else { - height = get_height(input, start + 1); + next_power_exponent = get_next_power_exponent(input, start + 1); } - height + next_power_exponent +} + +/// Calculates the previous power of 2. +/// +/// The previous power of 2 is the largest power of 2 that is smaller than `value`. +/// For example: +/// - For 7, previous power of 2 is 4 (2^2) +/// - For 10, previous power of 2 is 8 (2^3) +/// - For 16, previous power of 2 is 8 (2^3) +/// +/// # Arguments +/// * value - Integer value for which we need to find the previous power of 2 +/// +/// # Returns +/// The previous power of 2 +fn get_prev_power_2(value: u32) -> u32 { + let next_power_exponent = unsafe { + //@safety This is a hint that we'll use to compute the next and previous powers of two, which we check to be + // larger and smaller than respectively. The get_next_power_exponent function happens to return exactly what + // we need. + get_next_power_exponent(value, 0) + }; + + let next_power_2 = 2 << next_power_exponent; + let prev_power_2 = next_power_2 / 2; + assert((value == 0) | (value == 1) | (value > prev_power_2)); + assert(value <= next_power_2); + + prev_power_2 } // This calculates the root of the minimal size merkle tree required @@ -32,23 +72,12 @@ impl VariableMerkleTree { // | tx_0 | | tx_1 | // pub fn new_sha(leaves: [Field; N], num_non_empty_leaves: u32) -> Self { - // Find size of tree required - let height = unsafe { get_height(num_non_empty_leaves, 0) }; - let next_power_2 = 2 << height; - let prev_power_2 = next_power_2 / 2; - assert( - (num_non_empty_leaves == 0) - | (num_non_empty_leaves == 1) - | (num_non_empty_leaves > prev_power_2), - ); - assert(num_non_empty_leaves <= next_power_2); + let prev_power_2 = get_prev_power_2(num_non_empty_leaves); + // hash base layer // If we have no num_non_empty_leaves, we return 0 - let mut stop = if num_non_empty_leaves == 0 { - true - } else { - false - }; + let mut stop = num_non_empty_leaves == 0; + let mut nodes = [0; N]; for i in 0..N / 2 { // stop after non zero leaves @@ -61,7 +90,8 @@ impl VariableMerkleTree { } // hash the other layers - stop = if prev_power_2 == 1 { true } else { false }; + stop = prev_power_2 == 1; + let mut next_layer_end = prev_power_2 / 2; let mut next_layer_size = next_layer_end; let mut root = nodes[0]; @@ -190,3 +220,4 @@ pub mod tests { assert_eq(tree.root, full_tree.get_root()); } } +