refactor: VariableMerkleTree readability improvements

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<let N: u32>(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());
     }
 }
+