rewrite ZeroEx in Yul (#23)

rewrite ZeroEx in Yul

contracts/zero-ex/CHANGELOG.json

@@ -1,4 +1,13 @@
 [
+    {
+        "version": "0.9.0",
+        "changes": [
+            {
+                "note": "Rewrite the ZeroEx contract in Yul",
+                "pr": 23
+            }
+        ]
+    },
     {
         "version": "0.8.0",
         "changes": [

contracts/zero-ex/contracts/src/IZeroEx.sol

@@ -44,14 +44,4 @@ interface IZeroEx is
 
     /// @dev Fallback for just receiving ether.
     receive() external payable;
-
-    // solhint-enable state-visibility
-
-    /// @dev Get the implementation contract of a registered function.
-    /// @param selector The function selector.
-    /// @return impl The implementation contract address.
-    function getFunctionImplementation(bytes4 selector)
-        external
-        view
-        returns (address impl);
 }

contracts/zero-ex/contracts/src/ZeroEx.sol

@@ -19,19 +19,12 @@
 pragma solidity ^0.6.5;
 pragma experimental ABIEncoderV2;
 
-import "@0x/contracts-utils/contracts/src/v06/LibBytesV06.sol";
-import "./migrations/LibBootstrap.sol";
 import "./features/BootstrapFeature.sol";
 import "./storage/LibProxyStorage.sol";
-import "./errors/LibProxyRichErrors.sol";
-
 
 /// @dev An extensible proxy contract that serves as a universal entry point for
 ///      interacting with the 0x protocol.
 contract ZeroEx {
-    // solhint-disable separate-by-one-line-in-contract,indent,var-name-mixedcase
-    using LibBytesV06 for bytes;
-
     /// @dev Construct this contract and register the `BootstrapFeature` feature.
     ///      After constructing this contract, `bootstrap()` should be called
     ///      by `bootstrap()` to seed the initial feature set.
@@ -44,48 +37,55 @@ contract ZeroEx {
             address(bootstrap);
     }
 
+
     // solhint-disable state-visibility
 
     /// @dev Forwards calls to the appropriate implementation contract.
     fallback() external payable {
-        bytes4 selector = msg.data.readBytes4(0);
+        // This is used in assembly below as impls_slot.
-        address impl = getFunctionImplementation(selector);
+        mapping(bytes4 => address) storage impls =
-        if (impl == address(0)) {
+            LibProxyStorage.getStorage().impls;
-            _revertWithData(LibProxyRichErrors.NotImplementedError(selector));
+
+        assembly {
+            let cdlen := calldatasize()
+
+            // equivalent of receive() external payable {}
+            if iszero(cdlen) {
+                return(0, 0)
+            }
+
+            // Store at 0x40, to leave 0x00-0x3F for slot calculation below.
+            calldatacopy(0x40, 0, cdlen)
+            let selector := and(mload(0x40), 0xFFFFFFFF00000000000000000000000000000000000000000000000000000000)
+
+            // Slot for impls[selector] is keccak256(selector . impls_slot).
+            mstore(0, selector)
+            mstore(0x20, impls_slot)
+            let slot := keccak256(0, 0x40)
+
+            let delegate := sload(slot)
+            if iszero(delegate) {
+                // Revert with:
+                // abi.encodeWithSelector(
+                //   bytes4(keccak256("NotImplementedError(bytes4)")),
+                //   selector)
+                mstore(0, 0x734e6e1c00000000000000000000000000000000000000000000000000000000)
+                mstore(4, selector)
+                revert(0, 0x24)
+            }
+
+            let success := delegatecall(
+                gas(),
+                delegate,
+                0x40, cdlen,
+                0, 0
+            )
+            let rdlen := returndatasize()
+            returndatacopy(0, 0, rdlen)
+            if success {
+                return(0, rdlen)
+            }
+            revert(0, rdlen)
         }
-
-        (bool success, bytes memory resultData) = impl.delegatecall(msg.data);
-        if (!success) {
-            _revertWithData(resultData);
-        }
-        _returnWithData(resultData);
-    }
-
-    /// @dev Fallback for just receiving ether.
-    receive() external payable {}
-
-    // solhint-enable state-visibility
-
-    /// @dev Get the implementation contract of a registered function.
-    /// @param selector The function selector.
-    /// @return impl The implementation contract address.
-    function getFunctionImplementation(bytes4 selector)
-        public
-        view
-        returns (address impl)
-    {
-        return LibProxyStorage.getStorage().impls[selector];
-    }
-
-    /// @dev Revert with arbitrary bytes.
-    /// @param data Revert data.
-    function _revertWithData(bytes memory data) private pure {
-        assembly { revert(add(data, 32), mload(data)) }
-    }
-
-    /// @dev Return with arbitrary bytes.
-    /// @param data Return data.
-    function _returnWithData(bytes memory data) private pure {
-        assembly { return(add(data, 32), mload(data)) }
     }
 }

contracts/zero-ex/contracts/src/features/ISimpleFunctionRegistryFeature.sol

@@ -57,4 +57,12 @@ interface ISimpleFunctionRegistryFeature {
         external
         view
         returns (address impl);
+
+    /// @dev Get the implementation contract of a registered function.
+    /// @param selector The function selector.
+    /// @return impl The implementation contract address.
+    function getFunctionImplementation(bytes4 selector)
+        external
+        view
+        returns (address impl);
 }

contracts/zero-ex/contracts/src/features/SimpleFunctionRegistryFeature.sol

@@ -56,6 +56,7 @@ contract SimpleFunctionRegistryFeature is
         // Register getters.
         _extend(this.getRollbackLength.selector, _implementation);
         _extend(this.getRollbackEntryAtIndex.selector, _implementation);
+        _extend(this.getFunctionImplementation.selector, _implementation);
         return LibBootstrap.BOOTSTRAP_SUCCESS;
     }
 
@@ -151,6 +152,18 @@ contract SimpleFunctionRegistryFeature is
         return LibSimpleFunctionRegistryStorage.getStorage().implHistory[selector][idx];
     }
 
+    /// @dev Get the implementation contract of a registered function.
+    /// @param selector The function selector.
+    /// @return impl The implementation contract address.
+    function getFunctionImplementation(bytes4 selector)
+        external
+        override
+        view
+        returns (address impl)
+    {
+        return LibProxyStorage.getStorage().impls[selector];
+    }
+
     /// @dev Register or replace a function.
     /// @param selector The function selector.
     /// @param impl The implementation contract for the function.

contracts/zero-ex/contracts/test/TestInitialMigration.sol

@@ -22,6 +22,7 @@ pragma experimental ABIEncoderV2;
 import "../src/ZeroEx.sol";
 import "../src/features/IBootstrapFeature.sol";
 import "../src/migrations/InitialMigration.sol";
+import "../src/features/SimpleFunctionRegistryFeature.sol";
 
 
 contract TestInitialMigration is
@@ -44,7 +45,8 @@ contract TestInitialMigration is
     {
         success = InitialMigration.bootstrap(owner, features);
         // Snoop the bootstrap feature contract.
-        bootstrapFeature = ZeroEx(address(uint160(address(this))))
+        bootstrapFeature =
+            SimpleFunctionRegistryFeature(address(uint160(address(this))))
             .getFunctionImplementation(IBootstrapFeature.bootstrap.selector);
     }
 

contracts/zero-ex/test/features/simple_function_registry_test.ts

@@ -66,7 +66,7 @@ blockchainTests.resets('SimpleFunctionRegistry feature', env => {
 
     it('`rollback()` to zero impl succeeds for unregistered function', async () => {
         await registry.rollback(testFnSelector, NULL_ADDRESS).awaitTransactionSuccessAsync();
-        const impl = await zeroEx.getFunctionImplementation(testFnSelector).callAsync();
+        const impl = await registry.getFunctionImplementation(testFnSelector).callAsync();
         expect(impl).to.eq(NULL_ADDRESS);
     });
 

contracts/zero-ex/test/full_migration_test.ts

@@ -12,6 +12,7 @@ import {
     IMetaTransactionsFeatureContract,
     IOwnableFeatureContract,
     ISignatureValidatorFeatureContract,
+    ISimpleFunctionRegistryFeatureContract,
     ITokenSpenderFeatureContract,
     ITransformERC20FeatureContract,
     TestFullMigrationContract,
@@ -25,6 +26,7 @@ blockchainTests.resets('Full migration', env => {
     let zeroEx: ZeroExContract;
     let features: FullFeatures;
     let migrator: TestFullMigrationContract;
+    let registry: ISimpleFunctionRegistryFeatureContract;
     const transformerDeployer = randomAddress();
 
     before(async () => {
@@ -47,6 +49,7 @@ blockchainTests.resets('Full migration', env => {
         await migrator
             .initializeZeroEx(owner, zeroEx.address, features, { transformerDeployer })
             .awaitTransactionSuccessAsync();
+        registry = new ISimpleFunctionRegistryFeatureContract(zeroEx.address, env.provider, env.txDefaults);
     });
 
     it('ZeroEx has the correct owner', async () => {
@@ -157,7 +160,7 @@ blockchainTests.resets('Full migration', env => {
             for (const fn of featureInfo.fns) {
                 it(`${fn} is registered`, async () => {
                     const selector = contract.getSelector(fn);
-                    const impl = await zeroEx.getFunctionImplementation(selector).callAsync();
+                    const impl = await registry.getFunctionImplementation(selector).callAsync();
                     expect(impl).to.not.eq(NULL_ADDRESS);
                 });
 

contracts/zero-ex/test/zero_ex_test.ts

@@ -83,7 +83,7 @@ blockchainTests.resets('ZeroEx contract', env => {
                 // registry.getSelector('extendSelf'),
             ];
             const selectors = [...ownableSelectors, ...registrySelectors];
-            const impls = await Promise.all(selectors.map(s => zeroEx.getFunctionImplementation(s).callAsync()));
+            const impls = await Promise.all(selectors.map(s => registry.getFunctionImplementation(s).callAsync()));
             for (let i = 0; i < impls.length; ++i) {
                 const selector = selectors[i];
                 const impl = impls[i];