checkpoint with initial assembly getRoot func

dmfxyz · Mar 20, 2024 · 374a161 · 374a161
1 parent 6a3f9fc
commit 374a161
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 23 deletions.
diff --git a/src/CompleteMerkle.sol b/src/CompleteMerkle.sol
@@ -26,32 +26,57 @@ contract CompleteMerkle is MurkyBase {
                mstore(0x20, right)
            }
            _hash := keccak256(0x0, 0x40)
-            _hash := keccak256(0x0, 0x40)
         }
     }
 
-    function _getTree(bytes32[] memory data) public pure returns (bytes32[] memory) {
-        require(data.length > 1, "won't generate root for single leaf");
-        bytes32[] memory tree = new bytes32[](2 * data.length - 1);
-         assembly {
-            //let pos := mload(tree)
-            let iter := mload(data)
-            let y := add(sub(mload(tree), iter), 1)
-            for {let i := add(data, mul(iter, 0x20))} gt(i, iter) { i := sub(i, 0x20)} {
-                mstore(add(tree, mul(y, 0x20)), mload(i))
-                y := add(y, 1)
+    function getTree(bytes32[] memory data) public pure returns (bytes32[] memory) {
+        require(data.length > 1, 'wont generate root for single leaf');
 
+        bytes32[] memory tree = new bytes32[](2 * data.length - 1);
+        assembly {
+            let dlen := mload(data)
+            let titer := add(sub(mload(tree), dlen), 1)
+            for {let i := add(data, mul(dlen, 0x20))} gt(i, data) { i := sub(i, 0x20)} {
+                mstore(add(tree, mul(titer, 0x20)), mload(i))
+                titer := add(titer, 1)
             }
-        }
+         }
+        //  for (uint256 i = 0; i < data.length; ++i) {
+        //     tree[tree.length - i - 1] = data[i];
+        // }
         return tree;
     }
 
     function getRoot(bytes32[] memory data) public pure override returns (bytes32) {
-        bytes32[] memory tree = _getTree(data);
-        for (uint256 i = tree.length - 1; i > 0; i -= 2) {
-            uint256 posToWrite = (i - 1) / 2;
-            tree[posToWrite] = hashLeafPairs(tree[i - 1], tree[i]);
+        bytes32[] memory tree = getTree(data);
+        assembly {
+            function hash_leafs(left, right) -> _hash {
+                switch lt(left, right)
+                case 0 {
+                    mstore(0x0, right)
+                    mstore(0x20, left)
+                }
+                default {
+                    mstore(0x0, left)
+                    mstore(0x20, right)
+                }
+                _hash := keccak256(0x0, 0x40)
+            }
+            for {let i := mload(tree)} gt(i, 1) {i := sub(i, 2)} {
+                // TODO: clean all this up, mainly broken out for early understanding and debugging
+                let left := mload(add(tree, mul(sub(i, 1), 0x20)))
+                let right := mload(add(tree, mul(i, 0x20)))
+                let indexToWrite := div(sub(i, 1), 2)
+                let posToWrite := add(tree, mul(indexToWrite, 0x20))
+                mstore(posToWrite, hash_leafs(left, right))
+            }
         }
+        // for (uint256 i = tree.length - 1; i > 0; i -= 2) {
+        //     uint256 posToWrite = (i - 1) / 2;
+        //     tree[posToWrite] = hashLeafPairs(tree[i - 1], tree[i]);
+        // }
         return tree[0];
     }
+
+
 }
diff --git a/src/test/OZMerkle.t.sol b/src/test/OZMerkle.t.sol
@@ -109,7 +109,7 @@ contract ContractTest is Test {
         data[0] = 0xcf0e8c2fa63ea2b3726dbea696df21baec00c4cdb37deeab03a15f190659544c; // sha3(alice, address(0x0))
         data[1] = 0xafeb0ccce3b008968e7ffbfc7482d85551f5f90e713b8441449d808d25e9cc64; // sha3(bob, address(0x0))
         data[2] = 0x7f37e8358c0d3959e3e800344c357551753c24a55d5d987cef461e933b137a02; // sha3(charlie, address(0x0))
-        bytes32[] memory tree = cm._getTree(data);
+        //bytes32[] memory tree = cm._getTree(data);
         bytes32 root = cm.getRoot(data);
         assertEq(root, bytes32(0x10d772be44576d1c88a8039062ea3cd6163862f2deca298213e35b18cbd85490));
     }
@@ -123,11 +123,11 @@ contract ContractTest is Test {
         data[3] = 0xde1820ee7887b5ae922f14f423bb2e7a6595e423f1a0c0a82a2ddeed09a92a25;
         data[4] = 0xcac6fc160d04af9e1fd8f0c71cf8d333453b39589d3846524462ee7737bd728d;
         bytes32 elem;
-        assembly {
-            elem := mload(data)
-        }
-        emit log_bytes32((elem));
-        bytes32[] memory tree = cm._getTree(data);
+        // assembly {
+        //     elem := mload(data)
+        // }
+        // emit log_bytes32((elem));
+        // bytes32[] memory tree = cm._getTree(data);
         // for (uint256 i = 0; i < tree.length; ++i) {
         //     emit log_bytes32(tree[i]);
         // }
@@ -144,9 +144,24 @@ contract ContractTest is Test {
         data[3] = 0xde1820ee7887b5ae922f14f423bb2e7a6595e423f1a0c0a82a2ddeed09a92a25;
         data[4] = 0xcac6fc160d04af9e1fd8f0c71cf8d333453b39589d3846524462ee7737bd728d;
         data[5] = 0x1688f29243f54ddded6dedcbbc8dae64ef939f0b967d0fa56a6e5938febb5f79;
-        bytes32[] memory tree = cm._getTree(data);
+        //bytes32[] memory tree = cm._getTree(data);
         bytes32 root = cm.getRoot(data);
         assertEq(root, bytes32(0x82e830e6e22a0b487eac4fc6868317898b8c2131d21821dd2c8d723306afeb0a));
 
     }
+
+    function testRootGenerationRegularMerkle6ForGas() public {
+        bytes32[] memory data = new bytes32[](6);
+
+        data[0] = 0xcf0e8c2fa63ea2b3726dbea696df21baec00c4cdb37deeab03a15f190659544c; // sha3(alice, address(0x0))
+        data[1] = 0xafeb0ccce3b008968e7ffbfc7482d85551f5f90e713b8441449d808d25e9cc64; // sha3(bob, address(0x0))
+        data[2] = 0x7f37e8358c0d3959e3e800344c357551753c24a55d5d987cef461e933b137a02; // sha3(charlie, address(0x0))
+        data[3] = 0xde1820ee7887b5ae922f14f423bb2e7a6595e423f1a0c0a82a2ddeed09a92a25;
+        data[4] = 0xcac6fc160d04af9e1fd8f0c71cf8d333453b39589d3846524462ee7737bd728d;
+        data[5] = 0x1688f29243f54ddded6dedcbbc8dae64ef939f0b967d0fa56a6e5938febb5f79;
+        //bytes32[] memory tree = cm._getTree(data);
+        bytes32 root = m.getRoot(data);
+        assertEq(true, true);
+
+    }
 }