Added proof verification functions to CMT and IMT (#158)

RuslanProgrammer · web-flow · commit 7b34e8f4741f · 2025-08-26T12:40:52.000+03:00
* Added proof verification functions to CMT and IMT

* fixed review

* Updated changelog
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,7 +2,8 @@
 
 ## [none]
 
-- Fixed `pre-release` GitHub CI workflow
+- Fixed `pre-release` GitHub CI workflow.
+- Added `verifyProof` function to the `CartesianMerkleTree` and the `IncrementalMerkleTree` libraries.
 
 ## [3.2.0]
 
diff --git a/contracts/libs/data-structures/CartesianMerkleTree.sol b/contracts/libs/data-structures/CartesianMerkleTree.sol
@@ -5,7 +5,7 @@ pragma solidity ^0.8.21;
  * @notice Cartesian Merkle Tree Module
  *
  * A magnificent ZK-friendly data structure based on a Binary Search Tree + Heap + Merkle Tree. Short names: CMT, Treaple.
- * Possesses deterministic and idemponent properties. Can be used as a substitute for a Sparse Merkle Tree (SMT).
+ * Possesses deterministic and idempotent properties. Can be used as a substitute for a Sparse Merkle Tree (SMT).
  *
  * Gas usage for adding and removing 1,000 elements to a CMT with the keccak256 and poseidon hash functions is detailed below:
  *
@@ -32,6 +32,8 @@ pragma solidity ^0.8.21;
  *
  * CartesianMerkleTree.Proof memory proof = uintTreaple.getProof(100, 0);
  *
+ * uintTreaple.verifyProof(proof);
+ *
  * uintTreaple.getNodeByKey(100);
  *
  * uintTreaple.remove(100);
@@ -142,6 +144,21 @@ library CartesianMerkleTree {
         return _proof(treaple._treaple, bytes32(key_), desiredProofSize_);
     }
 
+    /**
+     * @notice The function to verify the proof for inclusion or exclusion of a node in the CMT.
+     * Complexity is O(log(n)), where n is the max depth of the treaple.
+     *
+     * @param treaple self.
+     * @param proof_ The CMT proof struct.
+     * @return True if the proof is valid, false otherwise.
+     */
+    function verifyProof(
+        UintCMT storage treaple,
+        Proof memory proof_
+    ) internal view returns (bool) {
+        return _verifyProof(treaple._treaple, proof_);
+    }
+
     /**
      * @notice The function to get the root of the Cartesian Merkle Tree.
      * Complexity is O(1).
@@ -306,6 +323,21 @@ library CartesianMerkleTree {
         return _proof(treaple._treaple, key_, desiredProofSize_);
     }
 
+    /**
+     * @notice The function to verify the proof for inclusion or exclusion of a node in the CMT.
+     * Complexity is O(log(n)), where n is the max depth of the treaple.
+     *
+     * @param treaple self.
+     * @param proof_ The CMT proof struct.
+     * @return True if the proof is valid, false otherwise.
+     */
+    function verifyProof(
+        Bytes32CMT storage treaple,
+        Proof memory proof_
+    ) internal view returns (bool) {
+        return _verifyProof(treaple._treaple, proof_);
+    }
+
     /**
      * @notice The function to get the root of the Cartesian Merkle Tree.
      * Complexity is O(1).
@@ -470,6 +502,21 @@ library CartesianMerkleTree {
         return _proof(treaple._treaple, _fromAddressToBytes32(key_), desiredProofSize_);
     }
 
+    /**
+     * @notice The function to verify the proof for inclusion or exclusion of a node in the CMT.
+     * Complexity is O(log(n)), where n is the max depth of the treaple.
+     *
+     * @param treaple self.
+     * @param proof_ The CMT proof struct.
+     * @return True if the proof is valid, false otherwise.
+     */
+    function verifyProof(
+        AddressCMT storage treaple,
+        Proof memory proof_
+    ) internal view returns (bool) {
+        return _verifyProof(treaple._treaple, proof_);
+    }
+
     /**
      * @notice The function to get the root of the Cartesian Merkle Tree.
      * Complexity is O(1).
@@ -871,6 +918,49 @@ library CartesianMerkleTree {
         return proof_;
     }
 
+    function _verifyProof(CMT storage treaple, Proof memory proof_) private view returns (bool) {
+        // invalid exclusion proof
+        if (!proof_.existence && proof_.nonExistenceKey == ZERO_HASH) {
+            return false;
+        }
+
+        bool directionBit_ = _extractDirectionBit(proof_.directionBits, 0);
+
+        bytes32 leftHash_ = proof_.siblings[proof_.siblingsLength - 2];
+        bytes32 rightHash_ = proof_.siblings[proof_.siblingsLength - 1];
+
+        if (directionBit_) {
+            (leftHash_, rightHash_) = (rightHash_, leftHash_);
+        }
+
+        bytes32 computedHash_ = _getNodesHash(
+            treaple,
+            proof_.existence ? proof_.key : proof_.nonExistenceKey,
+            leftHash_,
+            rightHash_
+        );
+
+        for (uint256 i = 2; i < proof_.siblingsLength; i += 2) {
+            directionBit_ = _extractDirectionBit(proof_.directionBits, i);
+
+            leftHash_ = computedHash_;
+            rightHash_ = proof_.siblings[proof_.siblingsLength - i - 1];
+
+            if (directionBit_) {
+                (leftHash_, rightHash_) = (rightHash_, leftHash_);
+            }
+
+            computedHash_ = _getNodesHash(
+                treaple,
+                proof_.siblings[proof_.siblingsLength - i - 2],
+                leftHash_,
+                rightHash_
+            );
+        }
+
+        return computedHash_ == proof_.root;
+    }
+
     function _newNode(CMT storage treaple, bytes32 key_) private returns (uint256) {
         uint64 nodeId_ = ++treaple.nodesCount;
 
@@ -914,6 +1004,17 @@ library CartesianMerkleTree {
         return directionBits_;
     }
 
+    /**
+     * @dev Extracts the direction bit from the direction mask.
+     * @return True if the node is on the right side of the parent, false if it's on the left.
+     */
+    function _extractDirectionBit(
+        uint256 directionBits_,
+        uint256 currentSiblingsIndex_
+    ) private pure returns (bool) {
+        return (directionBits_ & (1 << (currentSiblingsIndex_ / 2))) != 0;
+    }
+
     function _hashNodes(CMT storage treaple, uint256 nodeId_) private view returns (bytes32) {
         Node storage node = treaple.nodes[uint64(nodeId_)];
 
diff --git a/contracts/libs/data-structures/IncrementalMerkleTree.sol b/contracts/libs/data-structures/IncrementalMerkleTree.sol
@@ -110,6 +110,25 @@ library IncrementalMerkleTree {
         return _root(tree._tree);
     }
 
+    /**
+     * @notice The function to verify a proof of a leaf's existence in the uint256 tree.
+     * Complexity is O(log(n)), where n is the number of elements in the tree.
+     *
+     * @param tree self.
+     * @param siblings_ The siblings of the leaf.
+     * @param directionBits_ The direction bits of the leaf.
+     * @param leaf_ The leaf.
+     * @return True if the proof is valid, false otherwise.
+     */
+    function verifyProof(
+        UintIMT storage tree,
+        bytes32[] memory siblings_,
+        uint256 directionBits_,
+        bytes32 leaf_
+    ) internal view returns (bool) {
+        return _verifyProof(tree._tree, siblings_, directionBits_, leaf_);
+    }
+
     /**
      * @notice The function to return the height of the uint256 tree. Complexity is O(1).
      * @param tree self.
@@ -190,6 +209,25 @@ library IncrementalMerkleTree {
         return _root(tree._tree);
     }
 
+    /**
+     * @notice The function to verify a proof of a leaf's existence in the bytes32 tree.
+     * Complexity is O(log(n)), where n is the number of elements in the tree.
+     *
+     * @param tree self.
+     * @param siblings_ The siblings of the leaf.
+     * @param directionBits_ The direction bits of the leaf.
+     * @param leaf_ The leaf.
+     * @return True if the proof is valid, false otherwise.
+     */
+    function verifyProof(
+        Bytes32IMT storage tree,
+        bytes32[] memory siblings_,
+        uint256 directionBits_,
+        bytes32 leaf_
+    ) internal view returns (bool) {
+        return _verifyProof(tree._tree, siblings_, directionBits_, leaf_);
+    }
+
     /**
      * @notice The function to return the height of the bytes32 tree. Complexity is O(1).
      */
@@ -266,6 +304,25 @@ library IncrementalMerkleTree {
         return _root(tree._tree);
     }
 
+    /**
+     * @notice The function to verify a proof of a leaf's existence in the address tree.
+     * Complexity is O(log(n)), where n is the number of elements in the tree.
+     *
+     * @param tree self.
+     * @param siblings_ The siblings of the leaf.
+     * @param directionBits_ The direction bits of the leaf.
+     * @param leaf_ The leaf.
+     * @return True if the proof is valid, false otherwise.
+     */
+    function verifyProof(
+        AddressIMT storage tree,
+        bytes32[] memory siblings_,
+        uint256 directionBits_,
+        bytes32 leaf_
+    ) internal view returns (bool) {
+        return _verifyProof(tree._tree, siblings_, directionBits_, leaf_);
+    }
+
     /**
      * @notice The function to return the height of the address tree. Complexity is O(1).
      */
@@ -401,6 +458,29 @@ library IncrementalMerkleTree {
         return root_;
     }
 
+    function _verifyProof(
+        IMT storage tree,
+        bytes32[] memory siblings_,
+        uint256 directionBits,
+        bytes32 leaf_
+    ) private view returns (bool) {
+        function(bytes32, bytes32) view returns (bytes32) hash2_ = tree.isCustomHasherSet
+            ? tree.hash2
+            : _hash2;
+
+        bytes32 computedHash_ = leaf_;
+
+        for (uint256 i = 0; i < siblings_.length; ++i) {
+            if ((directionBits >> i) & 1 != 0) {
+                computedHash_ = hash2_(siblings_[i], computedHash_);
+            } else {
+                computedHash_ = hash2_(computedHash_, siblings_[i]);
+            }
+        }
+
+        return computedHash_ == _root(tree);
+    }
+
     function _height(IMT storage tree) private view returns (uint256) {
         return tree.branches.length;
     }
diff --git a/contracts/mock/libs/data-structures/CartesianMerkleTreeMock.sol b/contracts/mock/libs/data-structures/CartesianMerkleTreeMock.sol
@@ -96,6 +96,24 @@ contract CartesianMerkleTreeMock {
         return _addressCMT.getProof(key_, desiredProofSize_);
     }
 
+    function verifyUintProof(
+        CartesianMerkleTree.Proof memory proof_
+    ) external view returns (bool) {
+        return _uintCMT.verifyProof(proof_);
+    }
+
+    function verifyBytes32Proof(
+        CartesianMerkleTree.Proof memory proof_
+    ) external view returns (bool) {
+        return _bytes32CMT.verifyProof(proof_);
+    }
+
+    function verifyAddressProof(
+        CartesianMerkleTree.Proof memory proof_
+    ) external view returns (bool) {
+        return _addressCMT.verifyProof(proof_);
+    }
+
     function getUintRoot() external view returns (bytes32) {
         return _uintCMT.getRoot();
     }
diff --git a/contracts/mock/libs/data-structures/IncrementalMerkleTreeMock.sol b/contracts/mock/libs/data-structures/IncrementalMerkleTreeMock.sol
@@ -67,6 +67,30 @@ contract IncrementalMerkleTreeMock {
         return _addressTree.root();
     }
 
+    function verifyUintProof(
+        bytes32[] memory siblings_,
+        uint256 directionBits_,
+        bytes32 element_
+    ) external view returns (bool) {
+        return _uintTree.verifyProof(siblings_, directionBits_, element_);
+    }
+
+    function verifyBytes32Proof(
+        bytes32[] memory siblings_,
+        uint256 directionBits_,
+        bytes32 element_
+    ) external view returns (bool) {
+        return _bytes32Tree.verifyProof(siblings_, directionBits_, element_);
+    }
+
+    function verifyAddressProof(
+        bytes32[] memory siblings_,
+        uint256 directionBits_,
+        bytes32 element_
+    ) external view returns (bool) {
+        return _addressTree.verifyProof(siblings_, directionBits_, element_);
+    }
+
     function getUintTreeHeight() external view returns (uint256) {
         return _uintTree.height();
     }
diff --git a/test/libs/data-structures/CartesianMerkleTree.test.ts b/test/libs/data-structures/CartesianMerkleTree.test.ts
diff --git a/test/libs/data-structures/IncrementalMerkleTree.test.ts b/test/libs/data-structures/IncrementalMerkleTree.test.ts
