Improved functions for adding/replacing/removing functions

mudgen · Dec 19, 2020 · b009cd0 · b009cd0
1 parent bea5be1
commit b009cd0
Show file tree

Hide file tree

Showing 14 changed files with 90 additions and 94 deletions.
diff --git a/contracts/Diamond.sol b/contracts/Diamond.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 /******************************************************************************\

diff --git a/contracts/Migrations.sol b/contracts/Migrations.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 contract Migrations {

diff --git a/contracts/facets/DiamondCutFacet.sol b/contracts/facets/DiamondCutFacet.sol
@@ -1,13 +1,12 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 /******************************************************************************\
 * Author: Nick Mudge <[email protected]> (https://twitter.com/mudgen)
 * EIP-2535 Diamond Standard: https://eips.ethereum.org/EIPS/eip-2535
 /******************************************************************************/
 
-
 import "../interfaces/IDiamondCut.sol";
 import "../libraries/LibDiamond.sol";
 
@@ -24,14 +23,6 @@ contract DiamondCutFacet is IDiamondCut {
         bytes calldata _calldata
     ) external override {
         LibDiamond.enforceIsContractOwner();
-        for (uint256 facetIndex; facetIndex < _diamondCut.length; facetIndex++) {
-            LibDiamond.addReplaceRemoveFacetSelectors(
-                _diamondCut[facetIndex].facetAddress,
-                _diamondCut[facetIndex].action,
-                _diamondCut[facetIndex].functionSelectors
-            );
-        }
-        emit DiamondCut(_diamondCut, _init, _calldata);
-        LibDiamond.initializeDiamondCut(_init, _calldata);
+        LibDiamond.diamondCut(_diamondCut, _init, _calldata);
     }
 }
diff --git a/contracts/facets/DiamondLoupeFacet.sol b/contracts/facets/DiamondLoupeFacet.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 /******************************************************************************\

diff --git a/contracts/facets/OwnershipFacet.sol b/contracts/facets/OwnershipFacet.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 
 import "../libraries/LibDiamond.sol";
 import "../interfaces/IERC173.sol";

diff --git a/contracts/facets/Test1Facet.sol b/contracts/facets/Test1Facet.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 contract Test1Facet {

diff --git a/contracts/facets/Test2Facet.sol b/contracts/facets/Test2Facet.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 contract Test2Facet {

diff --git a/contracts/interfaces/IDiamondCut.sol b/contracts/interfaces/IDiamondCut.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 /******************************************************************************\

diff --git a/contracts/interfaces/IDiamondLoupe.sol b/contracts/interfaces/IDiamondLoupe.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 /******************************************************************************\

diff --git a/contracts/interfaces/IERC165.sol b/contracts/interfaces/IERC165.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 interface IERC165 {

diff --git a/contracts/interfaces/IERC173.sol b/contracts/interfaces/IERC173.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 
 /// @title ERC-173 Contract Ownership Standard
 ///  Note: the ERC-165 identifier for this interface is 0x7f5828d0

diff --git a/contracts/libraries/LibDiamond.sol b/contracts/libraries/LibDiamond.sol
@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-pragma solidity ^0.7.1;
+pragma solidity ^0.7.6;
 pragma experimental ABIEncoderV2;
 
 /******************************************************************************\
@@ -61,117 +61,119 @@ library LibDiamond {
         require(msg.sender == diamondStorage().contractOwner, "LibDiamond: Must be contract owner");
     }
 
-    modifier onlyOwner {
-        require(msg.sender == diamondStorage().contractOwner, "LibDiamond: Must be contract owner");
-        _;
-    }
-
     event DiamondCut(IDiamondCut.FacetCut[] _diamondCut, address _init, bytes _calldata);
 
     // Internal function version of diamondCut
-    // This code is almost the same as the external diamondCut,
-    // except it is using 'FacetCut[] memory _diamondCut' instead of
-    // 'FacetCut[] calldata _diamondCut'.
-    // The code is duplicated to prevent copying calldata to memory which
-    // causes an error for a two dimensional array.
     function diamondCut(
         IDiamondCut.FacetCut[] memory _diamondCut,
         address _init,
         bytes memory _calldata
     ) internal {
         for (uint256 facetIndex; facetIndex < _diamondCut.length; facetIndex++) {
-            addReplaceRemoveFacetSelectors(
-                _diamondCut[facetIndex].facetAddress,
-                _diamondCut[facetIndex].action,
-                _diamondCut[facetIndex].functionSelectors
-            );
+            IDiamondCut.FacetCutAction action = _diamondCut[facetIndex].action;
+            if (action == IDiamondCut.FacetCutAction.Add) {
+                addFunctions(_diamondCut[facetIndex].facetAddress, _diamondCut[facetIndex].functionSelectors);
+            } else if (action == IDiamondCut.FacetCutAction.Replace) {
+                replaceFunctions(_diamondCut[facetIndex].facetAddress, _diamondCut[facetIndex].functionSelectors);
+            } else if (action == IDiamondCut.FacetCutAction.Remove) {
+                removeFunctions(_diamondCut[facetIndex].facetAddress, _diamondCut[facetIndex].functionSelectors);
+            } else {
+                revert("LibDiamondCut: Incorrect FacetCutAction");
+            }
         }
         emit DiamondCut(_diamondCut, _init, _calldata);
         initializeDiamondCut(_init, _calldata);
     }
 
-    function addReplaceRemoveFacetSelectors(
-        address _newFacetAddress,
-        IDiamondCut.FacetCutAction _action,
-        bytes4[] memory _selectors
-    ) internal {
+    function addFunctions(address _facetAddress, bytes4[] memory _functionSelectors) internal {
+        require(_functionSelectors.length > 0, "LibDiamondCut: No selectors in facet to cut");
         DiamondStorage storage ds = diamondStorage();
-        require(_selectors.length > 0, "LibDiamondCut: No selectors in facet to cut");
-        // add or replace functions
-        if (_newFacetAddress != address(0)) {
-            uint256 facetAddressPosition = ds.facetFunctionSelectors[_newFacetAddress].facetAddressPosition;
-            // add new facet address if it does not exist
-            if (facetAddressPosition == 0 && ds.facetFunctionSelectors[_newFacetAddress].functionSelectors.length == 0) {
-                enforceHasContractCode(_newFacetAddress, "LibDiamondCut: New facet has no code");
-                facetAddressPosition = ds.facetAddresses.length;
-                ds.facetAddresses.push(_newFacetAddress);
-                ds.facetFunctionSelectors[_newFacetAddress].facetAddressPosition = uint16(facetAddressPosition);
-            }
-            // add or replace selectors
-            for (uint256 selectorIndex; selectorIndex < _selectors.length; selectorIndex++) {
-                bytes4 selector = _selectors[selectorIndex];
-                address oldFacetAddress = ds.selectorToFacetAndPosition[selector].facetAddress;
-                // add
-                if (_action == IDiamondCut.FacetCutAction.Add) {
-                    require(oldFacetAddress == address(0), "LibDiamondCut: Can't add function that already exists");
-                    addSelector(_newFacetAddress, selector);
-                } else if (_action == IDiamondCut.FacetCutAction.Replace) {
-                    // replace
-                    require(oldFacetAddress != _newFacetAddress, "LibDiamondCut: Can't replace function with same function");
-                    removeSelector(oldFacetAddress, selector);
-                    addSelector(_newFacetAddress, selector);
-                } else {
-                    revert("LibDiamondCut: Incorrect FacetCutAction");
-                }
-            }
-        } else {
-            require(_action == IDiamondCut.FacetCutAction.Remove, "LibDiamondCut: action not set to FacetCutAction.Remove");
-            // remove selectors
-            for (uint256 selectorIndex; selectorIndex < _selectors.length; selectorIndex++) {
-                bytes4 selector = _selectors[selectorIndex];
-                removeSelector(ds.selectorToFacetAndPosition[selector].facetAddress, selector);
-            }
+        // uint16 selectorCount = uint16(diamondStorage().selectors.length);
+        require(_facetAddress != address(0), "LibDiamondCut: Add facet can't be address(0)");
+        uint16 selectorPosition = uint16(ds.facetFunctionSelectors[_facetAddress].functionSelectors.length);
+        // add new facet address if it does not exist
+        if (selectorPosition == 0) {
+            enforceHasContractCode(_facetAddress, "LibDiamondCut: New facet has no code");
+            ds.facetFunctionSelectors[_facetAddress].facetAddressPosition = uint16(ds.facetAddresses.length);
+            ds.facetAddresses.push(_facetAddress);
+        }
+        for (uint256 selectorIndex; selectorIndex < _functionSelectors.length; selectorIndex++) {
+            bytes4 selector = _functionSelectors[selectorIndex];
+            address oldFacetAddress = ds.selectorToFacetAndPosition[selector].facetAddress;
+            require(oldFacetAddress == address(0), "LibDiamondCut: Can't add function that already exists");
+            ds.facetFunctionSelectors[_facetAddress].functionSelectors.push(selector);
+            ds.selectorToFacetAndPosition[selector].facetAddress = _facetAddress;
+            ds.selectorToFacetAndPosition[selector].functionSelectorPosition = selectorPosition;
+            selectorPosition++;
         }
     }
 
-    function addSelector(address _newFacet, bytes4 _selector) internal {
+    function replaceFunctions(address _facetAddress, bytes4[] memory _functionSelectors) internal {
+        require(_functionSelectors.length > 0, "LibDiamondCut: No selectors in facet to cut");
         DiamondStorage storage ds = diamondStorage();
-        uint256 selectorPosition = ds.facetFunctionSelectors[_newFacet].functionSelectors.length;
-        ds.facetFunctionSelectors[_newFacet].functionSelectors.push(_selector);
-        ds.selectorToFacetAndPosition[_selector].facetAddress = _newFacet;
-        ds.selectorToFacetAndPosition[_selector].functionSelectorPosition = uint16(selectorPosition);
+        require(_facetAddress != address(0), "LibDiamondCut: Add facet can't be address(0)");
+        uint16 selectorPosition = uint16(ds.facetFunctionSelectors[_facetAddress].functionSelectors.length);
+        // add new facet address if it does not exist
+        if (selectorPosition == 0) {
+            enforceHasContractCode(_facetAddress, "LibDiamondCut: New facet has no code");
+            ds.facetFunctionSelectors[_facetAddress].facetAddressPosition = uint16(ds.facetAddresses.length);
+            ds.facetAddresses.push(_facetAddress);
+        }
+        for (uint256 selectorIndex; selectorIndex < _functionSelectors.length; selectorIndex++) {
+            bytes4 selector = _functionSelectors[selectorIndex];
+            address oldFacetAddress = ds.selectorToFacetAndPosition[selector].facetAddress;
+            require(oldFacetAddress != _facetAddress, "LibDiamondCut: Can't replace function with same function");
+            removeFunction(oldFacetAddress, selector);
+            // add function
+            ds.selectorToFacetAndPosition[selector].functionSelectorPosition = selectorPosition;
+            ds.facetFunctionSelectors[_facetAddress].functionSelectors.push(selector);
+            ds.selectorToFacetAndPosition[selector].facetAddress = _facetAddress;
+            selectorPosition++;
+        }
     }
 
-    function removeSelector(address _oldFacetAddress, bytes4 _selector) internal {
+    function removeFunctions(address _facetAddress, bytes4[] memory _functionSelectors) internal {
+        require(_functionSelectors.length > 0, "LibDiamondCut: No selectors in facet to cut");
         DiamondStorage storage ds = diamondStorage();
         // if function does not exist then do nothing and return
-        require(_oldFacetAddress != address(0), "LibDiamondCut: Can't remove or replace function that doesn't exist");
-        require(_oldFacetAddress != address(this), "LibDiamondCut: Can't remove or replace immutable function");
+        require(_facetAddress == address(0), "LibDiamondCut: Remove facet address must be address(0)");
+        for (uint256 selectorIndex; selectorIndex < _functionSelectors.length; selectorIndex++) {
+            bytes4 selector = _functionSelectors[selectorIndex];
+            address oldFacetAddress = ds.selectorToFacetAndPosition[selector].facetAddress;
+            removeFunction(oldFacetAddress, selector);
+        }
+    }
+
+    function removeFunction(address _facetAddress, bytes4 _selector) internal {
+        DiamondStorage storage ds = diamondStorage();
+        require(_facetAddress != address(0), "LibDiamondCut: Can't remove function that doesn't exist");
+        // an immutable function is a function defined directly in a diamond
+        require(_facetAddress != address(this), "LibDiamondCut: Can't remove immutable function");
         // replace selector with last selector, then delete last selector
         uint256 selectorPosition = ds.selectorToFacetAndPosition[_selector].functionSelectorPosition;
-        uint256 lastSelectorPosition = ds.facetFunctionSelectors[_oldFacetAddress].functionSelectors.length - 1;
+        uint256 lastSelectorPosition = ds.facetFunctionSelectors[_facetAddress].functionSelectors.length - 1;
         // if not the same then replace _selector with lastSelector
         if (selectorPosition != lastSelectorPosition) {
-            bytes4 lastSelector = ds.facetFunctionSelectors[_oldFacetAddress].functionSelectors[lastSelectorPosition];
-            ds.facetFunctionSelectors[_oldFacetAddress].functionSelectors[selectorPosition] = lastSelector;
+            bytes4 lastSelector = ds.facetFunctionSelectors[_facetAddress].functionSelectors[lastSelectorPosition];
+            ds.facetFunctionSelectors[_facetAddress].functionSelectors[selectorPosition] = lastSelector;
             ds.selectorToFacetAndPosition[lastSelector].functionSelectorPosition = uint16(selectorPosition);
         }
         // delete the last selector
-        ds.facetFunctionSelectors[_oldFacetAddress].functionSelectors.pop();
+        ds.facetFunctionSelectors[_facetAddress].functionSelectors.pop();
         delete ds.selectorToFacetAndPosition[_selector];
 
         // if no more selectors for facet address then delete the facet address
         if (lastSelectorPosition == 0) {
             // replace facet address with last facet address and delete last facet address
             uint256 lastFacetAddressPosition = ds.facetAddresses.length - 1;
-            uint256 facetAddressPosition = ds.facetFunctionSelectors[_oldFacetAddress].facetAddressPosition;
+            uint256 facetAddressPosition = ds.facetFunctionSelectors[_facetAddress].facetAddressPosition;
             if (facetAddressPosition != lastFacetAddressPosition) {
                 address lastFacetAddress = ds.facetAddresses[lastFacetAddressPosition];
                 ds.facetAddresses[facetAddressPosition] = lastFacetAddress;
                 ds.facetFunctionSelectors[lastFacetAddress].facetAddressPosition = uint16(facetAddressPosition);
             }
             ds.facetAddresses.pop();
-            delete ds.facetFunctionSelectors[_oldFacetAddress];
+            delete ds.facetFunctionSelectors[_facetAddress].facetAddressPosition;
         }
     }
 

diff --git a/test/diamondTest.js b/test/diamondTest.js
@@ -47,7 +47,7 @@ contract('DiamondTest', async (accounts) => {
   let test1Facet
   let test2Facet
   let result
-  let addresses
+  let addresses = []
 
   const zeroAddress = '0x0000000000000000000000000000000000000000'
 
@@ -65,7 +65,10 @@ contract('DiamondTest', async (accounts) => {
   })
 
   it('should have three facets -- call to facetAddresses function', async () => {
-    addresses = await diamondLoupeFacet.methods.facetAddresses().call()
+    for (const address of await diamondLoupeFacet.methods.facetAddresses().call()) {
+      addresses.push(address)
+    }
+
     // console.log(addresses)
     assert.equal(addresses.length, 3)
   })

diff --git a/truffle-config.js b/truffle-config.js
@@ -86,7 +86,7 @@ module.exports = {
   // Configure your compilers
   compilers: {
     solc: {
-      version: '0.7.1', // Fetch exact version from solc-bin (default: truffle's version)
+      version: '0.7.6', // Fetch exact version from solc-bin (default: truffle's version)
       // docker: true,        // Use "0.5.1" you've installed locally with docker (default: false)
       settings: { // See the solidity docs for advice about optimization and evmVersion
         optimizer: {