✨ Deterministic SSTORE2 (#303)

Co-authored-by: James Wenzel <[email protected]>

Author: Vectorized

.gas-snapshot

@@ -422,11 +422,11 @@ MulticallableTest:testMulticallablePreservesMsgSender() (gas: 11118)
 MulticallableTest:testMulticallablePreservesMsgValue() (gas: 37569)
 MulticallableTest:testMulticallablePreservesMsgValueUsedTwice() (gas: 39417)
 MulticallableTest:testMulticallableReturnDataIsProperlyEncoded() (gas: 11634)
-MulticallableTest:testMulticallableReturnDataIsProperlyEncoded(string,string,uint256) (runs: 256, μ: 9708, ~: 7470)
+MulticallableTest:testMulticallableReturnDataIsProperlyEncoded(string,string,uint256) (runs: 256, μ: 9637, ~: 7464)
 MulticallableTest:testMulticallableReturnDataIsProperlyEncoded(uint256,uint256,uint256,uint256) (runs: 256, μ: 11802, ~: 11802)
 MulticallableTest:testMulticallableRevertWithCustomError() (gas: 10234)
 MulticallableTest:testMulticallableRevertWithMessage() (gas: 11946)
-MulticallableTest:testMulticallableRevertWithMessage(string) (runs: 256, μ: 12609, ~: 12655)
+MulticallableTest:testMulticallableRevertWithMessage(string) (runs: 256, μ: 12611, ~: 12655)
 MulticallableTest:testMulticallableRevertWithNothing() (gas: 10100)
 MulticallableTest:testMulticallableWithNoData() (gas: 6290)
 OwnableRolesTest:testGrantAndRemoveRolesDirect(address,uint256,uint256) (runs: 256, μ: 37138, ~: 40212)
@@ -458,26 +458,27 @@ OwnableRolesTest:testTransferOwnership() (gas: 19569)
 OwnableRolesTest:testTransferOwnership(address,bool,bool) (runs: 256, μ: 13821, ~: 13035)
 OwnableRolesTest:test__codesize() (gas: 20728)
 SSTORE2Test:testReadInvalidPointerCustomBoundsReverts() (gas: 3197)
-SSTORE2Test:testReadInvalidPointerCustomBoundsReverts(address,uint256,uint256) (runs: 256, μ: 724242, ~: 638383)
+SSTORE2Test:testReadInvalidPointerCustomBoundsReverts(address,uint256,uint256) (runs: 256, μ: 718904, ~: 624429)
 SSTORE2Test:testReadInvalidPointerCustomStartBoundReverts() (gas: 3241)
-SSTORE2Test:testReadInvalidPointerCustomStartBoundReverts(address,uint256) (runs: 256, μ: 726004, ~: 638486)
-SSTORE2Test:testReadInvalidPointerRevert(address) (runs: 256, μ: 750001, ~: 638362)
+SSTORE2Test:testReadInvalidPointerCustomStartBoundReverts(address,uint256) (runs: 256, μ: 745800, ~: 624451)
+SSTORE2Test:testReadInvalidPointerRevert(address) (runs: 256, μ: 670217, ~: 624561)
 SSTORE2Test:testReadInvalidPointerReverts() (gas: 3215)
 SSTORE2Test:testWriteRead() (gas: 76106)
-SSTORE2Test:testWriteRead(bytes) (runs: 256, μ: 801187, ~: 682523)
+SSTORE2Test:testWriteRead(bytes) (runs: 256, μ: 770750, ~: 668136)
 SSTORE2Test:testWriteReadCustomBounds() (gas: 34609)
-SSTORE2Test:testWriteReadCustomBounds(bytes,uint256,uint256) (runs: 256, μ: 743319, ~: 671096)
-SSTORE2Test:testWriteReadCustomBoundsOutOfRangeReverts(bytes,uint256,uint256) (runs: 256, μ: 802916, ~: 676174)
+SSTORE2Test:testWriteReadCustomBounds(bytes,uint256,uint256) (runs: 256, μ: 700018, ~: 655118)
+SSTORE2Test:testWriteReadCustomBoundsOutOfRangeReverts(bytes,uint256,uint256) (runs: 256, μ: 788965, ~: 662281)
 SSTORE2Test:testWriteReadCustomStartBound() (gas: 34843)
-SSTORE2Test:testWriteReadCustomStartBound(bytes,uint256) (runs: 256, μ: 814983, ~: 679249)
-SSTORE2Test:testWriteReadCustomStartBoundOutOfRangeReverts(bytes,uint256) (runs: 256, μ: 757497, ~: 675943)
+SSTORE2Test:testWriteReadCustomStartBound(bytes,uint256) (runs: 256, μ: 792462, ~: 665571)
+SSTORE2Test:testWriteReadCustomStartBoundOutOfRangeReverts(bytes,uint256) (runs: 256, μ: 721337, ~: 662060)
+SSTORE2Test:testWriteReadDeterministic(bytes) (runs: 256, μ: 872386, ~: 741858)
 SSTORE2Test:testWriteReadEmptyBound() (gas: 33835)
 SSTORE2Test:testWriteReadEmptyOutOfBoundsReverts() (gas: 36473)
 SSTORE2Test:testWriteReadFullBoundedRead() (gas: 76146)
 SSTORE2Test:testWriteReadFullStartBound() (gas: 35075)
 SSTORE2Test:testWriteReadOutOfBoundsReverts() (gas: 36473)
 SSTORE2Test:testWriteReadOutOfStartBoundReverts() (gas: 36477)
-SSTORE2Test:test__codesize() (gas: 8362)
+SSTORE2Test:test__codesize() (gas: 10082)
 SafeCastLibTest:testSafeCastToInt(int256) (runs: 256, μ: 4443, ~: 3415)
 SafeCastLibTest:testSafeCastToIntBench() (gas: 383456)
 SafeCastLibTest:testSafeCastToUint() (gas: 10560)

src/utils/SSTORE2.sol

@@ -34,7 +34,6 @@ library SSTORE2 {
 
     /// @dev Writes `data` into the bytecode of a storage contract and returns its address.
     function write(bytes memory data) internal returns (address pointer) {
-        // Note: The assembly block below does not expand the memory.
         /// @solidity memory-safe-assembly
         assembly {
             let originalDataLength := mload(data)
@@ -83,6 +82,73 @@ library SSTORE2 {
         }
     }
 
+    /// @dev Writes `data` into the bytecode of a storage contract with `salt`
+    /// and returns its deterministic address.
+    function writeDeterministic(bytes memory data, bytes32 salt)
+        internal
+        returns (address pointer)
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let originalDataLength := mload(data)
+            let dataSize := add(originalDataLength, DATA_OFFSET)
+
+            mstore(data, or(0x61000080600a3d393df300, shl(0x40, dataSize)))
+
+            // Deploy a new contract with the generated creation code.
+            pointer := create2(0, add(data, 0x15), add(dataSize, 0xa), salt)
+
+            // If `pointer` is zero, revert.
+            if iszero(pointer) {
+                // Store the function selector of `DeploymentFailed()`.
+                mstore(0x00, 0x30116425)
+                // Revert with (offset, size).
+                revert(0x1c, 0x04)
+            }
+
+            // Restore original length of the variable size `data`.
+            mstore(data, originalDataLength)
+        }
+    }
+
+    /// @dev Returns the initialization code hash of the storage contract for `data`.
+    /// Used for mining vanity addresses with create2crunch.
+    function initCodeHash(bytes memory data) internal pure returns (bytes32 hash) {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let originalDataLength := mload(data)
+            let dataSize := add(originalDataLength, DATA_OFFSET)
+
+            mstore(data, or(0x61000080600a3d393df300, shl(0x40, dataSize)))
+
+            hash := keccak256(add(data, 0x15), add(dataSize, 0xa))
+
+            // Restore original length of the variable size `data`.
+            mstore(data, originalDataLength)
+        }
+    }
+
+    /// @dev Returns the address of the storage contract for `data`
+    /// deployed with `salt` by `deployer`.
+    function predictDeterministicAddress(bytes memory data, bytes32 salt, address deployer)
+        internal
+        pure
+        returns (address predicted)
+    {
+        bytes32 hash = initCodeHash(data);
+        /// @solidity memory-safe-assembly
+        assembly {
+            // Compute and store the bytecode hash.
+            mstore8(0x00, 0xff) // Write the prefix.
+            mstore(0x35, hash)
+            mstore(0x01, shl(96, deployer))
+            mstore(0x15, salt)
+            predicted := keccak256(0x00, 0x55)
+            // Restore the part of the free memory pointer that has been overwritten.
+            mstore(0x35, 0)
+        }
+    }
+
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
     /*                         READ LOGIC                         */
     /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/

test/SSTORE2.t.sol

@@ -152,4 +152,21 @@ contract SSTORE2Test is TestPlus {
         vm.expectRevert(SSTORE2.ReadOutOfBounds.selector);
         SSTORE2.read(pointer, startIndex, endIndex);
     }
+
+    function testWriteReadDeterministic(bytes calldata testBytes) public brutalizeMemory {
+        bytes32 salt = bytes32(_random());
+        address deployer = address(this);
+        if (_random() % 8 == 0) {
+            (, deployer) = _randomSigner();
+        }
+        vm.prank(deployer);
+        address deterministicPointer = SSTORE2.writeDeterministic(testBytes, salt);
+        assertEq(SSTORE2.read(deterministicPointer), testBytes);
+        assertEq(
+            SSTORE2.predictDeterministicAddress(testBytes, salt, deployer), deterministicPointer
+        );
+
+        address pointer = SSTORE2.write(testBytes);
+        assertEq(pointer.code, deterministicPointer.code);
+    }
 }