From fcf983d5ab78128636718d0827cc4a218017e5d0 Mon Sep 17 00:00:00 2001
From: Ratan Kaliani <ratankaliani@berkeley.edu>
Date: Tue, 7 Nov 2023 23:29:35 -0800
Subject: [PATCH] refactor(plonky2x): `get_root_from_hashed_leaves` (#282)

---
 plonky2x/core/src/frontend/merkle/mod.rs       |  1 +
 .../core/src/frontend/merkle/tendermint.rs     | 16 ++++++++++++++--
 plonky2x/core/src/frontend/merkle/utils.rs     | 18 ++++++++++++++++++
 plonky2x/core/src/utils/lido/mod.rs            |  1 +
 4 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 plonky2x/core/src/frontend/merkle/utils.rs

diff --git a/plonky2x/core/src/frontend/merkle/mod.rs b/plonky2x/core/src/frontend/merkle/mod.rs
index 1bbe2f313..584c7f193 100644
--- a/plonky2x/core/src/frontend/merkle/mod.rs
+++ b/plonky2x/core/src/frontend/merkle/mod.rs
@@ -1,2 +1,3 @@
 pub mod tendermint;
 pub mod tree;
+pub mod utils;
diff --git a/plonky2x/core/src/frontend/merkle/tendermint.rs b/plonky2x/core/src/frontend/merkle/tendermint.rs
index a78ea730d..c3badf736 100644
--- a/plonky2x/core/src/frontend/merkle/tendermint.rs
+++ b/plonky2x/core/src/frontend/merkle/tendermint.rs
@@ -1,7 +1,10 @@
+use ethers::types::H256;
 use itertools::Itertools;
+use num::pow;
 
 use super::tree::MerkleInclusionProofVariable;
 use crate::backend::circuit::PlonkParameters;
+use crate::frontend::merkle::utils::log2_ceil_usize;
 use crate::frontend::vars::Bytes32Variable;
 use crate::prelude::{
     ArrayVariable, BoolVariable, ByteVariable, BytesVariable, CircuitBuilder, CircuitVariable,
@@ -127,17 +130,26 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
         leaf_hashes: Vec<Bytes32Variable>,
         leaves_enabled: Vec<BoolVariable>,
     ) -> Bytes32Variable {
-        assert!(NB_LEAVES.is_power_of_two());
         assert!(leaf_hashes.len() == NB_LEAVES);
         assert!(leaves_enabled.len() == NB_LEAVES);
 
+        let empty_bytes = Bytes32Variable::constant(self, H256::from_slice(&[0u8; 32]));
+
+        // Extend leaf_hashes and leaves_enabled to be a power of 2.
+        let padded_nb_leaves = pow(2, log2_ceil_usize(NB_LEAVES));
+        assert!(padded_nb_leaves >= NB_LEAVES && padded_nb_leaves.is_power_of_two());
+
         // Hash each of the validators to get their corresponding leaf hash.
+        // Pad the leaves to be a power of 2.
         let mut current_nodes = leaf_hashes.clone();
+        current_nodes.resize(padded_nb_leaves, empty_bytes);
 
         // Whether to treat the validator as empty.
+        // Pad the enabled array to be a power of 2.
         let mut current_node_enabled = leaves_enabled.clone();
+        current_node_enabled.resize(padded_nb_leaves, self._false());
 
-        let mut merkle_layer_size = NB_LEAVES;
+        let mut merkle_layer_size = padded_nb_leaves;
 
         // Hash each layer of nodes to get the root according to the Tendermint spec, starting from the leaves.
         while merkle_layer_size > 1 {
diff --git a/plonky2x/core/src/frontend/merkle/utils.rs b/plonky2x/core/src/frontend/merkle/utils.rs
new file mode 100644
index 000000000..03b60c688
--- /dev/null
+++ b/plonky2x/core/src/frontend/merkle/utils.rs
@@ -0,0 +1,18 @@
+/// Computes the ceiling of the base 2 log of a `usize`.
+pub fn log2_ceil_usize(x: usize) -> usize {
+    if x <= 1 {
+        // log2(0) and log2(1) are both 0.
+        return 0;
+    }
+
+    let mut result = 0;
+    // Subtract 1 to ensure rounding up for powers of 2.
+    let mut value = x - 1;
+
+    while value > 0 {
+        value >>= 1;
+        result += 1;
+    }
+
+    result as usize
+}
diff --git a/plonky2x/core/src/utils/lido/mod.rs b/plonky2x/core/src/utils/lido/mod.rs
index ac791c020..585cec2bd 100644
--- a/plonky2x/core/src/utils/lido/mod.rs
+++ b/plonky2x/core/src/utils/lido/mod.rs
@@ -167,6 +167,7 @@ mod tests {
     use super::*;
     use crate::utils;
     #[tokio::test]
+    #[cfg_attr(feature = "ci", ignore)]
     async fn test_lido_metadata() {
         utils::setup_logger();
         let rpc_url = "https://eth.llamarpc.com";