Feat/diamond refactor to eip2535

contracts/diamond/Diamond.sol

@@ -1,46 +1,63 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.4;
 
-import {Address} from "@openzeppelin/contracts/utils/Address.sol";
-import {EnumerableSet} from "@openzeppelin/contracts/utils/structs/EnumerableSet.sol";
-
-import {DiamondStorage} from "./DiamondStorage.sol";
-
-/**
- * @notice The Diamond standard module
- *
- * This is a custom implementation of a Diamond Proxy standard (https://eips.ethereum.org/EIPS/eip-2535).
- * This contract acts as a highest level contract of that standard. What is different from the EIP2535,
- * in order to use the DiamondStorage, storage is defined in a separate contract that the facets have to inherit from,
- * not an internal library.
- *
- * As a convention, view and pure function should be defined in the storage contract while function that modify state, in
- * the facet itself.
- *
- * If you wish to add a receive() function, you can attach a "0x00000000" selector to a facet that has such function.
- */
-contract Diamond is DiamondStorage {
-    using Address for address;
-    using EnumerableSet for EnumerableSet.Bytes32Set;
-    using EnumerableSet for EnumerableSet.AddressSet;
-
-    /**
-     * @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);
+/******************************************************************************\
+* Author: Nick Mudge <[email protected]> (https://twitter.com/mudgen)
+* EIP-2535 Diamonds: https://eips.ethereum.org/EIPS/eip-2535
+*
+* Implementation of a diamond.
+/******************************************************************************/
+
+import {LibDiamond} from "./libraries/LibDiamond.sol";
+import {IDiamondCut} from "./interfaces/IDiamondCut.sol";
+import {IDiamondLoupe} from "./interfaces/IDiamondLoupe.sol";
+import {IERC173} from "./interfaces/IERC173.sol";
+import {IERC165} from "./interfaces/IERC165.sol";
+
+// When no function exists for function called
+error FunctionNotFound(bytes4 _functionSelector);
+
+// This is used in diamond constructor
+// more arguments are added to this struct
+// this avoids stack too deep errors
+struct DiamondArgs {
+    address owner;
+    address init;
+    bytes initCalldata;
+}
 
-        require(facet_ != address(0), "Diamond: selector is not registered");
+contract Diamond {
+    constructor(IDiamondCut.FacetCut[] memory _diamondCut, DiamondArgs memory _args) payable {
+        LibDiamond.setContractOwner(_args.owner);
+        LibDiamond.diamondCut(_diamondCut, _args.init, _args.initCalldata);
 
-        _beforeFallback(facet_, msg.sig);
+        // Code can be added here to perform actions and set state variables.
+    }
 
+    // Find facet for function that is called and execute the
+    // function if a facet is found and return any value.
+    fallback() external payable {
+        LibDiamond.DiamondStorage storage ds;
+        bytes32 position = LibDiamond.DIAMOND_STORAGE_POSITION;
+        // get diamond storage
         assembly {
+            ds.slot := position
+        }
+        // get facet from function selector
+        address facet = ds.facetAddressAndSelectorPosition[msg.sig].facetAddress;
+        if (facet == address(0)) {
+            revert FunctionNotFound(msg.sig);
+        }
+        // Execute external function from facet using delegatecall and return any value.
+        assembly {
+            // copy function selector and any arguments
             calldatacopy(0, 0, calldatasize())
-            let result_ := delegatecall(gas(), facet_, 0, calldatasize(), 0, 0)
+            // execute function call using the facet
+            let result := delegatecall(gas(), facet, 0, calldatasize(), 0, 0)
+            // get any return value
             returndatacopy(0, 0, returndatasize())
-
-            switch result_
+            // return any return value or error back to the caller
+            switch result
             case 0 {
                 revert(0, returndatasize())
             }
@@ -50,69 +67,5 @@ contract Diamond is DiamondStorage {
         }
     }
 
-    /**
-     * @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 {
-        require(facet_.isContract(), "Diamond: facet is not a contract");
-        require(selectors_.length > 0, "Diamond: no selectors provided");
-
-        DStorage storage _ds = _getDiamondStorage();
-
-        for (uint256 i = 0; i < selectors_.length; i++) {
-            require(
-                _ds.selectorToFacet[selectors_[i]] == address(0),
-                "Diamond: selector already added"
-            );
-
-            _ds.selectorToFacet[selectors_[i]] = facet_;
-            _ds.facetToSelectors[facet_].add(bytes32(selectors_[i]));
-        }
-
-        _ds.facets.add(facet_);
-    }
-
-    /**
-     * @notice The internal function to remove facets from the diamond
-     * @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");
-
-        DStorage storage _ds = _getDiamondStorage();
-
-        for (uint256 i = 0; i < selectors_.length; i++) {
-            require(
-                _ds.selectorToFacet[selectors_[i]] == facet_,
-                "Diamond: selector from another facet"
-            );
-
-            _ds.selectorToFacet[selectors_[i]] = address(0);
-            _ds.facetToSelectors[facet_].remove(bytes32(selectors_[i]));
-        }
-
-        if (_ds.facetToSelectors[facet_].length() == 0) {
-            _ds.facets.remove(facet_);
-        }
-    }
-
-    /**
-     * @notice The internal function to update the facets 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
-     */
-    function _updateFacet(
-        address facet_,
-        bytes4[] memory fromSelectors_,
-        bytes4[] memory toSelectors_
-    ) internal {
-        _addFacet(facet_, toSelectors_);
-        _removeFacet(facet_, fromSelectors_);
-    }
-
-    function _beforeFallback(address facet_, bytes4 selector_) internal virtual {}
+    receive() external payable {}
 }

contracts/diamond/facets/DiamondCutFacet.sol

@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: CC0-1.0
+pragma solidity ^0.8.0;
+
+/******************************************************************************\
+* Author: Nick Mudge <[email protected]>, Twitter/Github: @mudgen
+* EIP-2535 Diamonds
+/******************************************************************************/
+
+import {IDiamondCut} from "../interfaces/IDiamondCut.sol";
+import {LibDiamond} from "../libraries/LibDiamond.sol";
+
+// Remember to add the loupe functions from DiamondLoupeFacet to the diamond.
+// The loupe functions are required by the EIP2535 Diamonds standard
+
+contract DiamondCutFacet is IDiamondCut {
+    /// @notice Add/replace/remove any number of functions and optionally execute
+    ///         a function with delegatecall
+    /// @param _diamondCut 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(
+        FacetCut[] calldata _diamondCut,
+        address _init,
+        bytes calldata _calldata
+    ) external override {
+        LibDiamond.enforceIsContractOwner();
+        LibDiamond.diamondCut(_diamondCut, _init, _calldata);
+    }
+}