Added directionBits to CMT proof (#129)

* Added directionBits to CMT proof * review fixes * review fixes in test * U512 lib: switched function from internal to private * updated version and lock file
dl-solarity · Jan 29, 2025 · 97aba26 · 97aba26
1 parent 6c2955d
commit 97aba26
Show file tree

Hide file tree

Showing 5 changed files with 752 additions and 432 deletions.
diff --git a/contracts/libs/bn/U512.sol b/contracts/libs/bn/U512.sol
@@ -1708,7 +1708,7 @@ library U512 {
      * @dev The modulus `m_` must be a prime number.
      * @dev Computes `(a_ * b_^(-1)) % m_` and stores the result in `r_`.
      */
-    function _moddiv(call512 call_, uint512 a_, uint512 b_, uint512 m_, uint512 r_) internal view {
+    function _moddiv(call512 call_, uint512 a_, uint512 b_, uint512 m_, uint512 r_) private view {
         unchecked {
             uint512 buffer_ = _buffer(call_);
 
@@ -1732,7 +1732,7 @@ library U512 {
      * @notice Performs bitwise AND of two 512-bit unsigned integers.
      * @dev Computes `a_ & b_` and stores the result in `r_`.
      */
-    function _and(uint512 a_, uint512 b_, uint512 r_) internal pure {
+    function _and(uint512 a_, uint512 b_, uint512 r_) private pure {
         unchecked {
             assembly {
                 mstore(r_, and(mload(a_), mload(b_)))
@@ -1745,7 +1745,7 @@ library U512 {
      * @notice Performs bitwise OR of two 512-bit unsigned integers.
      * @dev Computes `a_ | b_` and stores the result in `r_`.
      */
-    function _or(uint512 a_, uint512 b_, uint512 r_) internal pure {
+    function _or(uint512 a_, uint512 b_, uint512 r_) private pure {
         unchecked {
             assembly {
                 mstore(r_, or(mload(a_), mload(b_)))
@@ -1758,7 +1758,7 @@ library U512 {
      * @notice Performs bitwise XOR of two 512-bit unsigned integers.
      * @dev Computes `a_ ^ b_` and stores the result in `r_`.
      */
-    function _xor(uint512 a_, uint512 b_, uint512 r_) internal pure {
+    function _xor(uint512 a_, uint512 b_, uint512 r_) private pure {
         unchecked {
             assembly {
                 mstore(r_, xor(mload(a_), mload(b_)))
@@ -1771,7 +1771,7 @@ library U512 {
      * @notice Performs bitwise NOT of a 512-bit unsigned integer.
      * @dev Computes `~a_` and stores the result in `r_`.
      */
-    function _not(uint512 a_, uint512 r_) internal pure {
+    function _not(uint512 a_, uint512 r_) private pure {
         unchecked {
             assembly {
                 mstore(r_, not(mload(a_)))
@@ -1784,7 +1784,7 @@ library U512 {
      * @notice Performs left shift of a 512-bit unsigned integer.
      * @dev Computes `a_ << b_` and stores the result in `r_`.
      */
-    function _shl(uint512 a_, uint8 b_, uint512 r_) internal pure {
+    function _shl(uint512 a_, uint8 b_, uint512 r_) private pure {
         unchecked {
             assembly {
                 mstore(r_, or(shl(b_, mload(a_)), shr(sub(256, b_), mload(add(a_, 0x20)))))
@@ -1797,7 +1797,7 @@ library U512 {
      * @notice Performs right shift of a 512-bit unsigned integer.
      * @dev Computes `a_ >> b_` and stores the result in `r_`.
      */
-    function _shr(uint512 a_, uint8 b_, uint512 r_) internal pure {
+    function _shr(uint512 a_, uint8 b_, uint512 r_) private pure {
         unchecked {
             assembly {
                 mstore(

diff --git a/contracts/libs/data-structures/CartesianMerkleTree.sol b/contracts/libs/data-structures/CartesianMerkleTree.sol
@@ -30,7 +30,7 @@ pragma solidity ^0.8.4;
  *
  * uintTreaple.getRoot();
  *
- * CartesianMerkleTree.Proof memory proof = uintTreaple.getProof(100);
+ * CartesianMerkleTree.Proof memory proof = uintTreaple.getProof(100, 0);
  *
  * uintTreaple.getNodeByKey(100);
  *
@@ -590,6 +590,7 @@ library CartesianMerkleTree {
      * @param root The root hash of the Cartesian Merkle tree.
      * @param siblings An array of sibling hashes can be used to get the Cartesian Merkle Root.
      * @param siblingsLength The number of siblings to be used for evidence.
+     * @param directionBits A path from the root to the node.
      * @param existence Indicates the presence (true) or absence (false) of the node.
      * @param key The key associated with the node.
      * @param nonExistenceKey The non-existence key of the auxiliary node in case when existence is false.
@@ -598,6 +599,7 @@ library CartesianMerkleTree {
         bytes32 root;
         bytes32[] siblings;
         uint256 siblingsLength;
+        uint256 directionBits;
         bool existence;
         bytes32 key;
         bytes32 nonExistenceKey;
@@ -775,6 +777,7 @@ library CartesianMerkleTree {
             root: _rootMerkleHash(treaple),
             siblings: new bytes32[](desiredProofSize_),
             siblingsLength: 0,
+            directionBits: 0,
             existence: false,
             key: key_,
             nonExistenceKey: ZERO_HASH
@@ -787,20 +790,23 @@ library CartesianMerkleTree {
         Node storage node;
         uint256 currentSiblingsIndex_;
         uint256 nextNodeId_ = treaple.merkleRootId;
+        uint256 directionBits_;
 
         while (true) {
             node = treaple.nodes[uint64(nextNodeId_)];
 
             if (node.key == key_) {
-                _addProofSibling(
-                    proof_,
-                    currentSiblingsIndex_++,
-                    treaple.nodes[node.childLeft].merkleHash
-                );
-                _addProofSibling(
-                    proof_,
-                    currentSiblingsIndex_++,
-                    treaple.nodes[node.childRight].merkleHash
+                bytes32 leftHash_ = treaple.nodes[node.childLeft].merkleHash;
+                bytes32 rightHash_ = treaple.nodes[node.childRight].merkleHash;
+
+                _addProofSibling(proof_, currentSiblingsIndex_++, leftHash_);
+                _addProofSibling(proof_, currentSiblingsIndex_++, rightHash_);
+
+                proof_.directionBits = _calculateDirectionBit(
+                    directionBits_,
+                    currentSiblingsIndex_,
+                    leftHash_,
+                    rightHash_
                 );
 
                 proof_.existence = true;
@@ -820,15 +826,17 @@ library CartesianMerkleTree {
             }
 
             if (nextNodeId_ == 0) {
-                _addProofSibling(
-                    proof_,
-                    currentSiblingsIndex_++,
-                    treaple.nodes[node.childLeft].merkleHash
-                );
-                _addProofSibling(
-                    proof_,
-                    currentSiblingsIndex_++,
-                    treaple.nodes[node.childRight].merkleHash
+                bytes32 leftHash_ = treaple.nodes[node.childLeft].merkleHash;
+                bytes32 rightHash_ = treaple.nodes[node.childRight].merkleHash;
+
+                _addProofSibling(proof_, currentSiblingsIndex_++, leftHash_);
+                _addProofSibling(proof_, currentSiblingsIndex_++, rightHash_);
+
+                proof_.directionBits = _calculateDirectionBit(
+                    directionBits_,
+                    currentSiblingsIndex_,
+                    leftHash_,
+                    rightHash_
                 );
 
                 proof_.nonExistenceKey = node.key;
@@ -843,6 +851,13 @@ library CartesianMerkleTree {
                 currentSiblingsIndex_++,
                 treaple.nodes[otherNodeId_].merkleHash
             );
+
+            directionBits_ = _calculateDirectionBit(
+                directionBits_,
+                currentSiblingsIndex_,
+                treaple.nodes[uint64(nextNodeId_)].merkleHash,
+                treaple.nodes[otherNodeId_].merkleHash
+            );
         }
 
         return proof_;
@@ -875,6 +890,23 @@ library CartesianMerkleTree {
         proof_.siblings[currentSiblingsIndex_] = siblingToAdd_;
     }
 
+    function _calculateDirectionBit(
+        uint256 directionBits_,
+        uint256 currentSiblingsIndex_,
+        bytes32 leftHash_,
+        bytes32 rightHash_
+    ) private pure returns (uint256) {
+        if (currentSiblingsIndex_ != 2) {
+            directionBits_ <<= 1;
+        }
+
+        if (leftHash_ > rightHash_) {
+            directionBits_ |= 1;
+        }
+
+        return directionBits_;
+    }
+
     function _hashNodes(CMT storage treaple, uint256 nodeId_) private view returns (bytes32) {
         Node storage node = treaple.nodes[uint64(nodeId_)];