Merge pull request #19 from Arachnid/feature/create2

Add support for CREATE2 clones and address prediction
wighawag · Jan 17, 2024 · 0326de7 · 0326de7
2 parents 4343a1c + 2376882
commit 0326de7
Show file tree

Hide file tree

Showing 3 changed files with 127 additions and 9 deletions.
diff --git a/src/ClonesWithImmutableArgs.sol b/src/ClonesWithImmutableArgs.sol
@@ -3,7 +3,7 @@
 pragma solidity ^0.8.4;
 
 /// @title ClonesWithImmutableArgs
-/// @author wighawag, zefram.eth
+/// @author wighawag, zefram.eth, nick.eth
 /// @notice Enables creating clone contracts with immutable args
 library ClonesWithImmutableArgs {
     /// @dev The CREATE3 proxy bytecode.
@@ -18,6 +18,12 @@ library ClonesWithImmutableArgs {
     error CreateFail();
     error InitializeFail();
 
+    enum CloneType {
+        CREATE,
+        CREATE2,
+        PREDICT_CREATE2
+    }
+
     /// @notice Creates a clone proxy of the implementation contract, with immutable args
     /// @dev data cannot exceed 65535 bytes, since 2 bytes are used to store the data length
     /// @param implementation The implementation contract to clone
@@ -27,16 +33,76 @@ library ClonesWithImmutableArgs {
         internal
         returns (address payable instance)
     {
+        bytes memory creationcode = getCreationBytecode(implementation, data);
+        // solhint-disable-next-line no-inline-assembly
+        assembly {
+            instance := create(0, add(creationcode, 0x20), mload(creationcode))
+        }
+        if (instance == address(0)) {
+            revert CreateFail();
+        }
+    }
+
+    /// @notice Creates a clone proxy of the implementation contract, with immutable args,
+    ///         using CREATE2
+    /// @dev data cannot exceed 65535 bytes, since 2 bytes are used to store the data length
+    /// @param implementation The implementation contract to clone
+    /// @param data Encoded immutable args
+    /// @return instance The address of the created clone
+    function clone2(address implementation, bytes memory data)
+        internal
+        returns (address payable instance)
+    {
+        bytes memory creationcode = getCreationBytecode(implementation, data);
+        // solhint-disable-next-line no-inline-assembly
+        assembly {
+            instance := create2(0, add(creationcode, 0x20), mload(creationcode), 0)
+        }
+        if (instance == address(0)) {
+            revert CreateFail();
+        }
+    }
+
+    /// @notice Computes the address of a clone created using CREATE2
+    /// @dev data cannot exceed 65535 bytes, since 2 bytes are used to store the data length
+    /// @param implementation The implementation contract to clone
+    /// @param data Encoded immutable args
+    /// @return instance The address of the clone
+    function addressOfClone2(address implementation, bytes memory data)
+        internal
+        view
+        returns (address payable instance)
+    {
+        bytes memory creationcode = getCreationBytecode(implementation, data);
+        bytes32 bytecodeHash = keccak256(creationcode);
+        instance = payable(address(uint160(uint(keccak256(abi.encodePacked(
+            bytes1(0xff),
+            address(this),
+            bytes32(0),
+            bytecodeHash
+        ))))));
+    }
+
+    /// @notice Computes bytecode for a clone
+    /// @dev data cannot exceed 65535 bytes, since 2 bytes are used to store the data length
+    /// @param implementation The implementation contract to clone
+    /// @param data Encoded immutable args
+    /// @return ret Creation bytecode for the clone contract
+    function getCreationBytecode(address implementation, bytes memory data) internal pure returns (bytes memory ret) {
         // unrealistic for memory ptr or data length to exceed 256 bits
         unchecked {
             uint256 extraLength = data.length + 2; // +2 bytes for telling how much data there is appended to the call
             uint256 creationSize = 0x41 + extraLength;
             uint256 runSize = creationSize - 10;
             uint256 dataPtr;
             uint256 ptr;
+
             // solhint-disable-next-line no-inline-assembly
             assembly {
-                ptr := mload(0x40)
+                ret := mload(0x40)
+                mstore(ret, creationSize)
+                mstore(0x40, add(ret, creationSize))
+                ptr := add(ret, 0x20)
 
                 // -------------------------------------------------------------------------------------------------------------
                 // CREATION (10 bytes)
@@ -145,13 +211,6 @@ library ClonesWithImmutableArgs {
             assembly {
                 mstore(copyPtr, shl(240, extraLength))
             }
-            // solhint-disable-next-line no-inline-assembly
-            assembly {
-                instance := create(0, ptr, creationSize)
-            }
-            if (instance == address(0)) {
-                revert CreateFail();
-            }
         }
     }
 

diff --git a/src/ExampleCloneFactory.sol b/src/ExampleCloneFactory.sol
@@ -22,6 +22,26 @@ contract ExampleCloneFactory {
         bytes memory data = abi.encodePacked(param1, param2, param3, param4);
         clone = ExampleClone(address(implementation).clone(data));
     }
+
+    function createClone2(
+        address param1,
+        uint256 param2,
+        uint64 param3,
+        uint8 param4
+    ) external returns (ExampleClone clone) {
+        bytes memory data = abi.encodePacked(param1, param2, param3, param4);
+        clone = ExampleClone(address(implementation).clone2(data));
+    }
+
+    function addressOfClone2(
+        address param1,
+        uint256 param2,
+        uint64 param3,
+        uint8 param4
+    ) external view returns (address clone) {
+        bytes memory data = abi.encodePacked(param1, param2, param3, param4);
+        clone = address(implementation).addressOfClone2(data);
+    }
 
     function createClone3(
         address param1,

diff --git a/src/test/ExampleCloneFactory.t.sol b/src/test/ExampleCloneFactory.t.sol
@@ -62,6 +62,24 @@ contract ExampleCloneFactoryTest is DSTest {
         assertEq(clone.param4(), param4);
     }
 
+    function testCorrectness_clone2(
+        address param1,
+        uint256 param2,
+        uint64 param3,
+        uint8 param4
+    ) public {
+        ExampleClone clone = factory.createClone2(
+            param1,
+            param2,
+            param3,
+            param4
+        );
+        assertEq(clone.param1(), param1);
+        assertEq(clone.param2(), param2);
+        assertEq(clone.param3(), param3);
+        assertEq(clone.param4(), param4);
+    }
+
     function testCorrectness_clone3(
         address param1,
         uint256 param2,
@@ -83,6 +101,27 @@ contract ExampleCloneFactoryTest is DSTest {
         assertEq(address(clone), factory.addressOfClone3(salt));
     }
 
+    function testCorrectness_addressOfClone2(
+        address param1,
+        uint256 param2,
+        uint64 param3,
+        uint8 param4
+    ) public {
+        address predicted = factory.addressOfClone2(
+            param1,
+            param2,
+            param3,
+            param4
+        );
+        ExampleClone clone = factory.createClone2(
+            param1,
+            param2,
+            param3,
+            param4
+        );
+        assertEq(predicted, address(clone));
+    }
+
     /// -----------------------------------------------------------------------
     /// Failure tests
     /// -----------------------------------------------------------------------