Diamond refactor (#65)

* Diamond refactor * Diamond init example * Diamond minor fixes, add natspec, delete unnecessary * Delete recieve * Minor fixes * Add virtual + move require to the upper function * Diamond light customization, add tests, add/update dummy contracts * Delete unncessary revert * Delete only * dimond test js -> ts, add ownableDiamondMock * Minor fixes * Delete only * Add Virtual, delete imports, rename param --------- Co-authored-by: Artjom Galaktionov <[email protected]>
dl-solarity · Oct 9, 2023 · fe21d37 · fe21d37
1 parent 28e7045
commit fe21d37
Show file tree

Hide file tree

Showing 13 changed files with 563 additions and 135 deletions.
diff --git a/contracts/diamond/Diamond.sol b/contracts/diamond/Diamond.sol
@@ -24,12 +24,26 @@ contract Diamond is DiamondStorage {
     using EnumerableSet for EnumerableSet.Bytes32Set;
     using EnumerableSet for EnumerableSet.AddressSet;
 
+    enum FacetAction {
+        Add,
+        Replace,
+        Remove
+    }
+
+    struct Facet {
+        address facetAddress;
+        FacetAction action;
+        bytes4[] functionSelectors;
+    }
+
+    event DiamondCut(Facet[] facets, address initFacet, bytes initData);
+
     /**
      * @notice The payable fallback function that delegatecall's the facet with associated selector
      */
     // solhint-disable-next-line
     fallback() external payable virtual {
-        address facet_ = getFacetBySelector(msg.sig);
+        address facet_ = facetAddress(msg.sig);
 
         require(facet_ != address(0), "Diamond: selector is not registered");
 
@@ -50,14 +64,46 @@ contract Diamond is DiamondStorage {
         }
     }
 
+    /**
+     * @notice Add/replace/remove any number of functions and optionally execute a function with delegatecall
+     * @param facets_ Contains the facet addresses and function selectors
+     * @param initFacet_ The address of the contract or facet to execute initData_
+     * @param initData_ A function call, including function selector and arguments initData_ is executed with delegatecall on initFacet_
+     */
+    function _diamondCut(
+        Facet[] memory facets_,
+        address initFacet_,
+        bytes memory initData_
+    ) internal virtual {
+        for (uint256 i; i < facets_.length; i++) {
+            bytes4[] memory _functionSelectors = facets_[i].functionSelectors;
+            address _facetAddress = facets_[i].facetAddress;
+
+            FacetAction _action = facets_[i].action;
+
+            if (_action == FacetAction.Add) {
+                _addFacet(_facetAddress, _functionSelectors);
+            } else if (_action == FacetAction.Remove) {
+                _removeFacet(_facetAddress, _functionSelectors);
+            } else {
+                _updateFacet(_facetAddress, _functionSelectors);
+            }
+        }
+
+        emit DiamondCut(facets_, initFacet_, initData_);
+
+        _initializeDiamondCut(initFacet_, initData_);
+    }
+
     /**
      * @notice The internal function to add facets to a diamond (aka diamondCut())
      * @param facet_ the implementation address
      * @param selectors_ the function selectors the implementation has
      */
-    function _addFacet(address facet_, bytes4[] memory selectors_) internal {
+    function _addFacet(address facet_, bytes4[] memory selectors_) internal virtual {
+        require(facet_ != address(0), "Diamond: facet cannot be zero address");
         require(facet_.isContract(), "Diamond: facet is not a contract");
-        require(selectors_.length > 0, "Diamond: no selectors provided");
+        require(selectors_.length != 0, "Diamond: no selectors provided");
 
         DStorage storage _ds = _getDiamondStorage();
 
@@ -79,8 +125,9 @@ contract Diamond is DiamondStorage {
      * @param facet_ the implementation to be removed. The facet itself will be removed only if there are no selectors left
      * @param selectors_ the selectors of that implementation to be removed
      */
-    function _removeFacet(address facet_, bytes4[] memory selectors_) internal {
-        require(selectors_.length > 0, "Diamond: no selectors provided");
+    function _removeFacet(address facet_, bytes4[] memory selectors_) internal virtual {
+        require(facet_ != address(0), "Diamond: facet cannot be zero address");
+        require(selectors_.length != 0, "Diamond: no selectors provided");
 
         DStorage storage _ds = _getDiamondStorage();
 
@@ -100,18 +147,61 @@ contract Diamond is DiamondStorage {
     }
 
     /**
-     * @notice The internal function to update the facets of the diamond
+     * @notice The internal function to update the facet selectors of the diamond
      * @param facet_ the facet to update
-     * @param fromSelectors_ the selectors to remove from the facet
-     * @param toSelectors_ the selectors to add to the facet
+     * @param selectors_ the selectors of the facet
      */
-    function _updateFacet(
-        address facet_,
-        bytes4[] memory fromSelectors_,
-        bytes4[] memory toSelectors_
-    ) internal {
-        _addFacet(facet_, toSelectors_);
-        _removeFacet(facet_, fromSelectors_);
+    function _updateFacet(address facet_, bytes4[] memory selectors_) internal virtual {
+        require(facet_ != address(0), "Diamond: facet cannot be zero address");
+        require(facet_.isContract(), "Diamond: facet is not a contract");
+        require(selectors_.length != 0, "Diamond: no selectors provided");
+
+        DStorage storage _ds = _getDiamondStorage();
+
+        for (uint256 i; i < selectors_.length; i++) {
+            bytes4 selector_ = selectors_[i];
+            address oldFacet_ = facetAddress(selector_);
+
+            require(oldFacet_ != facet_, "Diamond: cannot replace to the same facet");
+            require(oldFacet_ != address(0), "Diamond: no facet found for selector");
+
+            // replace old facet address
+            _ds.selectorToFacet[selector_] = facet_;
+            _ds.facetToSelectors[facet_].add(bytes32(selector_));
+
+            // remove old facet address
+            _ds.facetToSelectors[oldFacet_].remove(bytes32(selector_));
+
+            if (_ds.facetToSelectors[oldFacet_].length() == 0) {
+                _ds.facets.remove(oldFacet_);
+            }
+        }
+
+        _ds.facets.add(facet_);
+    }
+
+    /**
+     * @notice The internal function to initialize the diamond cut.
+     * @param initFacet_ the address of the contract or facet to execute initData_
+     * @param initData_ a function call, including function selector and arguments, to be executed with delegatecall on initFacet_
+     */
+    function _initializeDiamondCut(address initFacet_, bytes memory initData_) internal virtual {
+        if (initFacet_ == address(0)) {
+            return;
+        }
+
+        require(initFacet_.isContract(), "Diamond: init_ address has no code");
+
+        (bool success_, bytes memory err_) = initFacet_.delegatecall(initData_);
+
+        if (!success_) {
+            require(err_.length > 0, "Diamond: initialization function reverted");
+            // bubble up error
+            // @solidity memory-safe-assembly
+            assembly {
+                revert(add(32, err_), mload(err_))
+            }
+        }
     }
 
     function _beforeFallback(address facet_, bytes4 selector_) internal virtual {}

diff --git a/contracts/diamond/DiamondStorage.sol b/contracts/diamond/DiamondStorage.sol
@@ -26,6 +26,11 @@ abstract contract DiamondStorage {
         EnumerableSet.AddressSet facets;
     }
 
+    struct FacetInfo {
+        address facetAddress;
+        bytes4[] functionSelectors;
+    }
+
     /**
      * @notice The internal function to get the diamond proxy storage
      * @return _ds the struct from the DIAMOND_STORAGE_SLOT
@@ -39,19 +44,30 @@ abstract contract DiamondStorage {
     }
 
     /**
-     * @notice The function to get all the facets of this diamond
-     * @return facets_ the array of facets' addresses
+     * @notice The function to get all the facets and their selectors
+     * @return facets_ the array of FacetInfo
      */
-    function getFacets() public view returns (address[] memory facets_) {
-        return _getDiamondStorage().facets.values();
+    function facets() public view returns (FacetInfo[] memory facets_) {
+        EnumerableSet.AddressSet storage _facets = _getDiamondStorage().facets;
+
+        facets_ = new FacetInfo[](_facets.length());
+
+        for (uint256 i = 0; i < facets_.length; i++) {
+            address facet_ = _facets.at(i);
+
+            facets_[i].facetAddress = facet_;
+            facets_[i].functionSelectors = facetFunctionSelectors(facet_);
+        }
     }
 
     /**
      * @notice The function to get all the selectors assigned to the facet
      * @param facet_ the facet to get assigned selectors of
      * @return selectors_ the array of assigned selectors
      */
-    function getFacetSelectors(address facet_) public view returns (bytes4[] memory selectors_) {
+    function facetFunctionSelectors(
+        address facet_
+    ) public view returns (bytes4[] memory selectors_) {
         EnumerableSet.Bytes32Set storage _f2s = _getDiamondStorage().facetToSelectors[facet_];
 
         selectors_ = new bytes4[](_f2s.length());
@@ -61,12 +77,20 @@ abstract contract DiamondStorage {
         }
     }
 
+    /**
+     * @notice The function to get all the facets of this diamond
+     * @return facets_ the array of facets' addresses
+     */
+    function facetAddresses() public view returns (address[] memory facets_) {
+        return _getDiamondStorage().facets.values();
+    }
+
     /**
      * @notice The function to get associated facet by the selector
      * @param selector_ the selector
      * @return facet_ the associated facet address
      */
-    function getFacetBySelector(bytes4 selector_) public view returns (address facet_) {
+    function facetAddress(bytes4 selector_) public view returns (address facet_) {
         return _getDiamondStorage().selectorToFacet[selector_];
     }
 }
diff --git a/contracts/diamond/introspection/DiamondERC165.sol b/contracts/diamond/introspection/DiamondERC165.sol
@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import {ERC165} from "@openzeppelin/contracts/utils/introspection/ERC165.sol";
+
+/**
+ * @notice DiamondERC165 - Contract implementing ERC165 interface for Diamonds
+ */
+contract DiamondERC165 is ERC165 {
+    function supportsInterface(bytes4 interfaceId) public view virtual override returns (bool) {
+        // This section of code provides support for the Diamond Loupe and Diamond Cut interfaces.
+        // Diamond Loupe interface is defined as: 0x48e2b093
+        // Diamond Cut interface is defined as: 0x1f931c1c
+        return
+            interfaceId == 0x1f931c1c ||
+            interfaceId == 0x48e2b093 ||
+            super.supportsInterface(interfaceId);
+    }
+}
diff --git a/contracts/diamond/presets/OwnableDiamond/OwnableDiamond.sol b/contracts/diamond/presets/OwnableDiamond/OwnableDiamond.sol
@@ -18,19 +18,15 @@ contract OwnableDiamond is Diamond, OwnableDiamondStorage {
         _getOwnableDiamondStorage().owner = newOwner_;
     }
 
-    function addFacet(address facet_, bytes4[] memory selectors_) public virtual onlyOwner {
-        _addFacet(facet_, selectors_);
+    function diamondCut(Facet[] memory facets_) public onlyOwner {
+        diamondCut(facets_, address(0), "");
     }
 
-    function removeFacet(address facet_, bytes4[] memory selectors_) public virtual onlyOwner {
-        _removeFacet(facet_, selectors_);
-    }
-
-    function updateFacet(
-        address facet_,
-        bytes4[] memory fromSelectors_,
-        bytes4[] memory toSelectors_
-    ) public virtual onlyOwner {
-        _updateFacet(facet_, fromSelectors_, toSelectors_);
+    function diamondCut(
+        Facet[] memory facets_,
+        address init_,
+        bytes memory initData_
+    ) public onlyOwner {
+        _diamondCut(facets_, init_, initData_);
     }
 }
diff --git a/contracts/mock/diamond/DummyFacet.sol b/contracts/mock/diamond/DummyFacet.sol
@@ -4,7 +4,7 @@ pragma solidity ^0.8.4;
 import {DummyStorage} from "./DummyStorage.sol";
 
 contract DummyFacet is DummyStorage {
-    function setDummyString(string calldata dummyString_) external {
+    function setDummyString(string memory dummyString_) public {
         getDummyFacetStorage().dummyString = dummyString_;
     }
 

diff --git a/contracts/mock/diamond/DummyInit.sol b/contracts/mock/diamond/DummyInit.sol
@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import {DummyFacet} from "./DummyFacet.sol";
+
+contract DummyInit is DummyFacet {
+    event Initialized();
+
+    function init() external {
+        setDummyString("dummy facet initialized");
+        emit Initialized();
+    }
+
+    function initWithError() external pure {
+        revert();
+    }
+
+    function initWithErrorMsg() external pure {
+        revert("DiamondInit: init error");
+    }
+}
diff --git a/contracts/mock/diamond/OwnableDiamondMock.sol b/contracts/mock/diamond/OwnableDiamondMock.sol
@@ -0,0 +1,14 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import {OwnableDiamond} from "../../diamond/presets/OwnableDiamond/OwnableDiamond.sol";
+
+contract OwnableDiamondMock is OwnableDiamond {
+    function diamondCutShort(Facet[] memory facets_) public {
+        diamondCut(facets_);
+    }
+
+    function diamondCutLong(Facet[] memory facets_, address init_, bytes memory calldata_) public {
+        diamondCut(facets_, init_, calldata_);
+    }
+}
diff --git a/contracts/package.json b/contracts/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@solarity/solidity-lib",
-  "version": "2.5.11",
+  "version": "2.5.12",
   "license": "MIT",
   "author": "Distributed Lab",
   "readme": "README.md",

diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@solarity/solidity-lib",
-  "version": "2.5.11",
+  "version": "2.5.12",
   "license": "MIT",
   "author": "Distributed Lab",
   "description": "Solidity Library by Distributed Lab",