test: Merkle libs

cartesi · Jan 28, 2025 · 4f42cfa · 4f42cfa
1 parent 4dc5078
commit 4f42cfa
Showing 1 changed file with 65 additions and 0 deletions.
diff --git a/cartesi-rollups/contracts/test/Merkle.t.sol b/cartesi-rollups/contracts/test/Merkle.t.sol
@@ -9,10 +9,75 @@ import {MerkleConstants} from "src/MerkleConstants.sol";
 import {PristineMerkleTree} from "src/PristineMerkleTree.sol";
 import {Merkle} from "src/Merkle.sol";
 
+contract PristineMerkleTreeContract {
+    function getNodeAtHeight(uint256 height) external pure returns (bytes32) {
+        return PristineMerkleTree.getNodeAtHeight(height);
+    }
+}
+
+contract MerkleContract {
+    function getHashOfLeafAtIndex(bytes calldata data, uint256 leafIndex) external pure returns (bytes32) {
+        return Merkle.getHashOfLeafAtIndex(data, leafIndex);
+    }
+}
+
 contract MerkleTest is Test {
     using Merkle for bytes;
 
+    PristineMerkleTreeContract pristineMerkle;
+    MerkleContract merkle;
+
+    function setUp() external {
+        pristineMerkle = new PristineMerkleTreeContract();
+        merkle = new MerkleContract();
+    }
+
     function testJoinEquivalence(bytes32 a, bytes32 b) external pure {
         assertEq(Merkle.join(a, b), keccak256(abi.encodePacked(a, b)));
     }
+
+    function testPristineMerkleTree() external pure {
+        bytes32 node = keccak256(abi.encode(0));
+        for (uint256 height; height <= MerkleConstants.TREE_HEIGHT; ++height) {
+            assertEq(PristineMerkleTree.getNodeAtHeight(height), node);
+            node = Merkle.join(node, node);
+        }
+    }
+
+    function testPristineMerkleTreeRevert() external {
+        vm.expectRevert("Height out of bounds");
+        pristineMerkle.getNodeAtHeight(MerkleConstants.TREE_HEIGHT + 1);
+    }
+
+    function testPristineMerkleTreeRevert(uint256 height) external {
+        height = bound(height, MerkleConstants.TREE_HEIGHT + 1, type(uint256).max);
+        vm.expectRevert("Height out of bounds");
+        pristineMerkle.getNodeAtHeight(height);
+    }
+
+    function testGetHashOfLeafAtIndex(bytes calldata data) external view {
+        // whole leaves
+        for (uint256 leafIndex; leafIndex < (data.length >> MerkleConstants.LOG2_LEAF_SIZE); ++leafIndex) {
+            uint256 start = leafIndex << MerkleConstants.LOG2_LEAF_SIZE;
+            uint256 end = (leafIndex + 1) << MerkleConstants.LOG2_LEAF_SIZE;
+            assertEq(merkle.getHashOfLeafAtIndex(data, leafIndex), keccak256(data[start:end]));
+        }
+        // first empty or pristine leaf
+        {
+            uint256 leafIndex = data.length >> MerkleConstants.LOG2_LEAF_SIZE;
+            uint256 start = leafIndex << MerkleConstants.LOG2_LEAF_SIZE;
+            assertEq(merkle.getHashOfLeafAtIndex(data, leafIndex), keccak256(abi.encodePacked(bytes32(data[start:]))));
+        }
+        // first or second pristine leaf
+        {
+            uint256 leafIndex = 1 + (data.length >> MerkleConstants.LOG2_LEAF_SIZE);
+            assertEq(merkle.getHashOfLeafAtIndex(data, leafIndex), pristineMerkle.getNodeAtHeight(0));
+        }
+    }
+
+    function testGetHashOfLeafAtIndex(bytes calldata data, uint256 leafIndex) external view {
+        uint256 leafIndexUpperBound = data.length >> MerkleConstants.LOG2_LEAF_SIZE;
+        leafIndex = bound(leafIndex, 1 + leafIndexUpperBound, type(uint256).max);
+        assertEq(merkle.getHashOfLeafAtIndex(data, leafIndex), pristineMerkle.getNodeAtHeight(0));
+    }
 }