Diamond refactor

contracts/diamond/Diamond.sol

@@ -24,12 +24,27 @@ contract Diamond is DiamondStorage {
     using EnumerableSet for EnumerableSet.Bytes32Set;
     using EnumerableSet for EnumerableSet.AddressSet;
 
+    enum FacetAction {
+        Add,
+        Replace,
+        Remove
+    }
+    // Add=0, Replace=1, Remove=2
+
+    struct Facet {
+        address facetAddress;
+        FacetAction action;
+        bytes4[] functionSelectors;
+    }
+
+    event DiamondCut(Facet[] facets_, address init_, bytes calldata_);
+
     /**
      * @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,6 +65,53 @@ contract Diamond is DiamondStorage {
         }
     }
 
+    receive() external payable virtual {}
+
+    /**
+     * @notice Add/replace/remove any number of functions and optionally execute a function with delegatecall
+     * @param facets_ Contains the facet addresses and function selectors
+     */
+    function _diamondCut(Facet[] calldata facets_) internal {
+        _diamondCut(facets_, address(0), "");
+    }
+
+    /**
+     * @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 init_ The address of the contract or facet to execute calldata_
+     * @param calldata_ A function call, including function selector and arguments calldata_ is executed with delegatecall on init_
+     */
+    function _diamondCut(
+        Facet[] calldata facets_,
+        address init_,
+        bytes memory calldata_
+    ) internal {
+        for (uint256 i; i < facets_.length; i++) {
+            bytes4[] memory _functionSelectors = facets_[i].functionSelectors;
+            address _facetAddress = facets_[i].facetAddress;
+
+            require(
+                _functionSelectors.length != 0,
+                "Diamond: no selectors provided for facet for cut"
+            );
+
+            FacetAction _action = facets_[i].action;
+            if (_action == FacetAction.Add) {
+                _addFacet(_facetAddress, _functionSelectors);
+            } else if (_action == FacetAction.Remove) {
+                _removeFacet(_facetAddress, _functionSelectors);
+            } else if (_action == FacetAction.Replace) {
+                _updateFacet(_facetAddress, _functionSelectors);
+            } else {
+                revert("Diamond: incorrect facet action");
+            }
+        }
+
+        emit DiamondCut(facets_, init_, calldata_);
+
+        _initializeDiamondCut(init_, calldata_);
+    }
+
     /**
      * @notice The internal function to add facets to a diamond (aka diamondCut())
      * @param facet_ the implementation address
@@ -100,18 +162,56 @@ 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 {
+        require(facet_ != address(0), "Diamond:  facet cannot be zero address");
+        require(facet_.isContract(), "Diamond: replace facet has no code");
+
+        DStorage storage _ds = _getDiamondStorage();
+
+        for (uint256 i; i < selectors_.length; i++) {
+            bytes4 _selector = selectors_[i];
+            address _oldFacet = facetAddress(_selector);
+
+            // can't replace immutable functions -- functions defined directly in the diamond in this case
+            require(_oldFacet == address(this), "Diamond: cannot replace function of this");
+            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);
+            }
+        }
+    }
+
+    function _initializeDiamondCut(address init_, bytes memory calldata_) internal {
+        if (init_ == address(0)) {
+            return;
+        }
+
+        require(init_.isContract(), "Diamond: init_ address has no code");
+
+        (bool success, bytes memory err) = init_.delegatecall(calldata_);
+
+        if (!success) {
+            require(err.length > 0, "Diamond: initialization function reverted");
+            // bubble up error
+            /// @solidity memory-safe-assembly
+            assembly {
+                let returndata_size := mload(err)
+                revert(add(32, err), returndata_size)
+            }
+        }
     }
 
     function _beforeFallback(address facet_, bytes4 selector_) internal virtual {}

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_];
     }
 }

contracts/diamond/init/DiamondInit.sol

@@ -0,0 +1,27 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+// It is expected that this contract is customized if you want to deploy your diamond
+// with data from a deployment script. Use the init function to initialize state variables
+// of your diamond. Add parameters to the init funciton if you need to.
+
+// Adding parameters to the `init` or other functions you add here can make a single deployed
+// DiamondInit contract reusable accross upgrades, and can be used for multiple diamonds.
+
+import "../introspection/DiamondERC165.sol";
+
+contract DiamondInit is DiamondERC165 {
+    // You can add parameters to this function in order to pass in
+    // data to set your own state variables
+    function init() external {
+        // adding ERC165 data
+        registerInterface(type(IERC165).interfaceId);
+
+        // add your own state variables
+        // EIP-2535 specifies that the `diamondCut` function takes two optional
+        // arguments: address _init and bytes calldata _calldata
+        // These arguments are used to execute an arbitrary function using delegatecall
+        // in order to set state variables in the diamond during deployment or an upgrade
+        // More info here: https://eips.ethereum.org/EIPS/eip-2535#diamond-interface
+    }
+}

contracts/diamond/introspection/DiamondERC165.sol

@@ -0,0 +1,11 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import "./DiamondERC165Storage.sol";
+
+contract DiamondERC165 is DiamondERC165Storage {
+    function registerInterface(bytes4 interfaceId_) public virtual {
+        require(interfaceId_ != 0xffffffff, "ERC165: invalid interface id");
+        _getErc165Storage().supportedInterfaces[interfaceId_] = true;
+    }
+}

contracts/diamond/introspection/DiamondERC165Storage.sol

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import "@openzeppelin/contracts/utils/introspection/IERC165.sol";
+
+abstract contract DiamondERC165Storage is IERC165 {
+    bytes32 public constant DIAMOND_ERC165_STORAGE_SLOT =
+        keccak256("diamond.standard.diamond.erc165.storage");
+
+    struct DERC165Storage {
+        mapping(bytes4 => bool) supportedInterfaces;
+    }
+
+    function _getErc165Storage() internal pure returns (DERC165Storage storage ds_) {
+        bytes32 slot_ = DIAMOND_ERC165_STORAGE_SLOT;
+
+        assembly {
+            ds_.slot := slot_
+        }
+    }
+
+    function supportsInterface(bytes4 interfaceId_) public view virtual override returns (bool) {
+        return _getErc165Storage().supportedInterfaces[interfaceId_];
+    }
+}

contracts/diamond/presets/OwnableDiamond/OwnableDiamond.sol

@@ -18,6 +18,18 @@ contract OwnableDiamond is Diamond, OwnableDiamondStorage {
         _getOwnableDiamondStorage().owner = newOwner_;
     }
 
+    function diamondCut(Facet[] calldata facets_) public virtual onlyOwner {
+        _diamondCut(facets_);
+    }
+
+    function diamondCut(
+        Facet[] calldata facets_,
+        address init_,
+        bytes memory calldata_
+    ) public virtual onlyOwner {
+        _diamondCut(facets_, init_, calldata_);
+    }
+
     function addFacet(address facet_, bytes4[] memory selectors_) public virtual onlyOwner {
         _addFacet(facet_, selectors_);
     }
@@ -26,11 +38,7 @@ contract OwnableDiamond is Diamond, OwnableDiamondStorage {
         _removeFacet(facet_, selectors_);
     }
 
-    function updateFacet(
-        address facet_,
-        bytes4[] memory fromSelectors_,
-        bytes4[] memory toSelectors_
-    ) public virtual onlyOwner {
-        _updateFacet(facet_, fromSelectors_, toSelectors_);
+    function updateFacet(address facet_, bytes4[] memory selectors_) public virtual onlyOwner {
+        _updateFacet(facet_, selectors_);
     }
 }