feat: merkle tree verifier implementation to support all numbers of leaves

src/lib/tree/Utils.sol

@@ -42,3 +42,34 @@ function pathLengthFromKey(uint256 key, uint256 numLeaves) pure returns (uint256
         return 1 + pathLengthFromKey(key - numLeavesLeftSubTree, numLeaves - numLeavesLeftSubTree);
     }
 }
+
+/// @notice Returns the minimum number of bits required to represent `x`; the
+/// result is 0 for `x` == 0.
+/// @param x Number.
+function _bitsLen(uint256 x) pure returns (uint256) {
+    uint256 count = 0;
+
+    while (x != 0) {
+        count++;
+        x >>= 1;
+    }
+
+    return count;
+}
+
+/// @notice Returns the largest power of 2 less than `x`.
+/// @param x Number.
+function _getSplitPoint(uint256 x) pure returns (uint256) {
+    // Note: since `x` is always an unsigned int * 2, the only way for this
+    // to be violated is if the input == 0. Since the input is the end
+    // index exclusive, an input of 0 is guaranteed to be invalid (it would
+    // be a proof of inclusion of nothing, which is vacuous).
+    require(x >= 1);
+
+    uint256 bitLen = _bitsLen(x);
+    uint256 k = 1 << (bitLen - 1);
+    if (k == x) {
+        k >>= 1;
+    }
+    return k;
+}

src/lib/tree/binary/BinaryMerkleTree.sol

@@ -42,65 +42,56 @@ library BinaryMerkleTree {
             }
         }
 
-        uint256 height = 1;
-        uint256 stableEnd = proof.key;
+        bytes32 computedHash = computeRootHash(proof.key, proof.numLeaves, digest, proof.sideNodes);
 
-        // While the current subtree (of height 'height') is complete, determine
-        // the position of the next sibling using the complete subtree algorithm.
-        // 'stableEnd' tells us the ending index of the last full subtree. It gets
-        // initialized to 'key' because the first full subtree was the
-        // subtree of height 1, created above (and had an ending index of
-        // 'key').
-
-        while (true) {
-            // Determine if the subtree is complete. This is accomplished by
-            // rounding down the key to the nearest 1 << 'height', adding 1
-            // << 'height', and comparing the result to the number of leaves in the
-            // Merkle tree.
-
-            uint256 subTreeStartIndex = (proof.key / (1 << height)) * (1 << height);
-            uint256 subTreeEndIndex = subTreeStartIndex + (1 << height) - 1;
-
-            // If the Merkle tree does not have a leaf at index
-            // 'subTreeEndIndex', then the subtree of the current height is not
-            // a complete subtree.
-            if (subTreeEndIndex >= proof.numLeaves) {
-                break;
-            }
-            stableEnd = subTreeEndIndex;
-
-            // Determine if the key is in the first or the second half of
-            // the subtree.
-            if (proof.sideNodes.length <= height - 1) {
-                return false;
-            }
-            if (proof.key - subTreeStartIndex < (1 << (height - 1))) {
-                digest = nodeDigest(digest, proof.sideNodes[height - 1]);
-            } else {
-                digest = nodeDigest(proof.sideNodes[height - 1], digest);
-            }
+        return (computedHash == root);
+    }
 
-            height += 1;
+    /// @notice Use the leafHash and innerHashes to get the root merkle hash.
+    /// If the length of the innerHashes slice isn't exactly correct, the result is nil.
+    /// Recursive impl.
+    function computeRootHash(uint256 key, uint256 numLeaves, bytes32 leafHash, bytes32[] memory sideNodes)
+        private
+        pure
+        returns (bytes32)
+    {
+        if (numLeaves == 0) {
+            revert("cannot call computeRootHash with 0 number of leaves");
         }
-
-        // Determine if the next hash belongs to an orphan that was elevated. This
-        // is the case IFF 'stableEnd' (the last index of the largest full subtree)
-        // is equal to the number of leaves in the Merkle tree.
-        if (stableEnd != proof.numLeaves - 1) {
-            if (proof.sideNodes.length <= height - 1) {
-                return false;
+        if (numLeaves == 1) {
+            if (sideNodes.length != 0) {
+                revert("unexpected inner hashes");
             }
-            digest = nodeDigest(digest, proof.sideNodes[height - 1]);
-            height += 1;
+            return leafHash;
         }
-
-        // All remaining elements in the proof set will belong to a left sibling\
-        // i.e proof sideNodes are hashed in "from the left"
-        while (height - 1 < proof.sideNodes.length) {
-            digest = nodeDigest(proof.sideNodes[height - 1], digest);
-            height += 1;
+        if (sideNodes.length == 0) {
+            revert("expected at least one inner hash");
         }
+        uint256 numLeft = _getSplitPoint(numLeaves);
+        bytes32[] memory sideNodesLeft = slice(sideNodes, 0, sideNodes.length - 1);
+        if (key < numLeft) {
+            bytes32 leftHash = computeRootHash(key, numLeft, leafHash, sideNodesLeft);
+            return nodeDigest(leftHash, sideNodes[sideNodes.length - 1]);
+        }
+        bytes32 rightHash = computeRootHash(key - numLeft, numLeaves - numLeft, leafHash, sideNodesLeft);
+        return nodeDigest(sideNodes[sideNodes.length - 1], rightHash);
+    }
 
-        return (digest == root);
+    /// @notice creates a slice of bytes32 from the data slice of bytes32 containing the elements
+    /// that correspond to the provided range.
+    /// It selects a half-open range which includes the begin element, but excludes the end one.
+    /// @param _data The slice that we want to select data from.
+    /// @param _begin The beginning of the range (inclusive).
+    /// @param _end The ending of the range (exclusive).
+    /// @return _ the sliced data.
+    function slice(bytes32[] memory _data, uint256 _begin, uint256 _end) internal pure returns (bytes32[] memory) {
+        if (_begin > _end) {
+            revert("Invalid range: _begin is greater than _end");
+        }
+        bytes32[] memory out = new bytes32[](_end-_begin);
+        for (uint256 i = _begin; i < _end; i++) {
+            out[i - _begin] = _data[i];
+        }
+        return out;
     }
 }

src/lib/tree/binary/test/BinaryMerkleTree.t.sol

@@ -101,6 +101,36 @@ contract BinaryMerkleProofTest is DSTest {
         assertTrue(isValid);
     }
 
+    function testVerifyLeafTwoOfEight() external {
+        bytes32 root = 0xc1ad6548cb4c7663110df219ec8b36ca63b01158956f4be31a38a88d0c7f7071;
+        bytes32[] memory sideNodes = new bytes32[](3);
+        sideNodes[0] = 0xb413f47d13ee2fe6c845b2ee141af81de858df4ec549a58b7970bb96645bc8d2;
+        sideNodes[1] = 0x78850a5ab36238b076dd99fd258c70d523168704247988a94caa8c9ccd056b8d;
+        sideNodes[2] = 0x4301a067262bbb18b4919742326f6f6d706099f9c0e8b0f2db7b88f204b2cf09;
+
+        uint256 key = 1;
+        uint256 numLeaves = 8;
+        BinaryMerkleProof memory proof = BinaryMerkleProof(sideNodes, key, numLeaves);
+        bytes memory data = hex"02";
+        bool isValid = BinaryMerkleTree.verify(root, proof, data);
+        assertTrue(isValid);
+    }
+
+    function testVerifyLeafThreeOfEight() external {
+        bytes32 root = 0xc1ad6548cb4c7663110df219ec8b36ca63b01158956f4be31a38a88d0c7f7071;
+        bytes32[] memory sideNodes = new bytes32[](3);
+        sideNodes[0] = 0x4f35212d12f9ad2036492c95f1fe79baf4ec7bd9bef3dffa7579f2293ff546a4;
+        sideNodes[1] = 0x6bcf0e2e93e0a18e22789aee965e6553f4fbe93f0acfc4a705d691c8311c4965;
+        sideNodes[2] = 0x4301a067262bbb18b4919742326f6f6d706099f9c0e8b0f2db7b88f204b2cf09;
+
+        uint256 key = 2;
+        uint256 numLeaves = 8;
+        BinaryMerkleProof memory proof = BinaryMerkleProof(sideNodes, key, numLeaves);
+        bytes memory data = hex"03";
+        bool isValid = BinaryMerkleTree.verify(root, proof, data);
+        assertTrue(isValid);
+    }
+
     function testVerifyLeafSevenOfEight() external {
         bytes32 root = 0xc1ad6548cb4c7663110df219ec8b36ca63b01158956f4be31a38a88d0c7f7071;
         bytes32[] memory sideNodes = new bytes32[](3);
@@ -130,4 +160,25 @@ contract BinaryMerkleProofTest is DSTest {
         bool isValid = BinaryMerkleTree.verify(root, proof, data);
         assertTrue(isValid);
     }
+
+    // Test vectors:
+    // 0x00
+    // 0x01
+    // 0x02
+    // 0x03
+    // 0x04
+    function testVerifyProofOfFiveLeaves() external {
+        bytes32 root = 0xb855b42d6c30f5b087e05266783fbd6e394f7b926013ccaa67700a8b0c5a596f;
+        bytes32[] memory sideNodes = new bytes32[](3);
+        sideNodes[0] = 0x96a296d224f285c67bee93c30f8a309157f0daa35dc5b87e410b78630a09cfc7;
+        sideNodes[1] = 0x52c56b473e5246933e7852989cd9feba3b38f078742b93afff1e65ed46797825;
+        sideNodes[2] = 0x4f35212d12f9ad2036492c95f1fe79baf4ec7bd9bef3dffa7579f2293ff546a4;
+
+        uint256 key = 1;
+        uint256 numLeaves = 5;
+        BinaryMerkleProof memory proof = BinaryMerkleProof(sideNodes, key, numLeaves);
+        bytes memory data = bytes(hex"01");
+        bool isValid = BinaryMerkleTree.verify(root, proof, data);
+        assertTrue(isValid);
+    }
 }

src/lib/tree/namespace/NamespaceMerkleTree.sol

@@ -218,37 +218,6 @@ library NamespaceMerkleTree {
         return count;
     }
 
-    /// @notice Returns the minimum number of bits required to represent `x`; the
-    /// result is 0 for `x` == 0.
-    /// @param x Number.
-    function _bitsLen(uint256 x) private pure returns (uint256) {
-        uint256 count = 0;
-
-        while (x != 0) {
-            count++;
-            x >>= 1;
-        }
-
-        return count;
-    }
-
-    /// @notice Returns the largest power of 2 less than `x`.
-    /// @param x Number.
-    function _getSplitPoint(uint256 x) private pure returns (uint256) {
-        // Note: since `x` is always an unsigned int * 2, the only way for this
-        // to be violated is if the input == 0. Since the input is the end
-        // index exclusive, an input of 0 is guaranteed to be invalid (it would
-        // be a proof of inclusion of nothing, which is vacuous).
-        require(x >= 1);
-
-        uint256 bitLen = _bitsLen(x);
-        uint256 k = 1 << (bitLen - 1);
-        if (k == x) {
-            k >>= 1;
-        }
-        return k;
-    }
-
     /// @notice Computes the NMT root recursively.
     /// @param proof Namespace Merkle multiproof for the leaves.
     /// @param leafNodes Leaf nodes for which inclusion is proven.