Merge pull request #20 from dadadave80/feat/initializable-diamond

Add initializable Diamond pattern

Author: dadadave80

script/DeployDiamond.s.sol

@@ -6,7 +6,9 @@ import {MockDiamond} from "@diamond-test/mocks/MockDiamond.sol";
 import {DiamondCutFacet} from "@diamond/facets/DiamondCutFacet.sol";
 import {DiamondLoupeFacet} from "@diamond/facets/DiamondLoupeFacet.sol";
 import {OwnableRolesFacet} from "@diamond/facets/OwnableRolesFacet.sol";
-import {DiamondInit} from "@diamond/initializers/DiamondInit.sol";
+import {ERC165Init} from "@diamond/initializers/ERC165Init.sol";
+import {MultiInit} from "@diamond/initializers/MultiInit.sol";
+import {OwnableInit} from "@diamond/initializers/OwnableInit.sol";
 import {ContextLib} from "@diamond/libraries/ContextLib.sol";
 import {FacetCut, FacetCutAction} from "@diamond/libraries/DiamondLib.sol";
 import {Script} from "forge-std/Script.sol";
@@ -28,8 +30,10 @@ contract DeployDiamond is Script, GetSelectors {
         DiamondLoupeFacet diamondLoupeFacet = new DiamondLoupeFacet();
         OwnableRolesFacet ownableRolesFacet = new OwnableRolesFacet();
 
-        // Deploy ERC165 initializer contract
-        address diamondInit = address(new DiamondInit());
+        // Deploy initializer contracts
+        address multiInit = address(new MultiInit());
+        address ownableInit = address(new OwnableInit());
+        address erc165Init = address(new ERC165Init());
 
         // Create an array of FacetCut entries for standard facets
         FacetCut[] memory cut = new FacetCut[](3);
@@ -55,9 +59,21 @@ contract DeployDiamond is Script, GetSelectors {
             functionSelectors: _getSelectors("OwnableRolesFacet")
         });
 
+        // Build MultiInit arrays for granular initialization
+        address[] memory initAddresses = new address[](2);
+        bytes[] memory initData = new bytes[](2);
+
+        initAddresses[0] = ownableInit;
+        initData[0] = abi.encodeWithSignature("initOwner(address)", ContextLib.msgSender());
+
+        initAddresses[1] = erc165Init;
+        initData[1] = abi.encodeWithSignature("initERC165()");
+
         // Deploy the Diamond contract with the facets and initialization args
-        MockDiamond diamond =
-            new MockDiamond(cut, diamondInit, abi.encodeWithSignature("initDiamond(address)", ContextLib.msgSender()));
+        MockDiamond diamond = new MockDiamond();
+        diamond.initialize(
+            cut, multiInit, abi.encodeWithSignature("multiInit(address[],bytes[])", initAddresses, initData)
+        );
         diamond_ = address(diamond);
         vm.stopBroadcast();
     }

src/Diamond.sol

@@ -2,6 +2,7 @@
 pragma solidity ^0.8.20;
 
 import {DiamondLib, FacetCut} from "@diamond/libraries/DiamondLib.sol";
+import {InitializableLib} from "@diamond/libraries/InitializableLib.sol";
 
 /*
 ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀
@@ -32,23 +33,39 @@ import {DiamondLib, FacetCut} from "@diamond/libraries/DiamondLib.sol";
 
 /// @title Diamond
 /// @notice Implements ERC-2535 Diamond proxy pattern, allowing dynamic addition, replacement, and removal of facets
-/// @author Nick Mudge (https://github.com/mudgen/diamond-3-hardhat/blob/main/contracts/Diamond.sol)
-/// @author Modified by David Dada <daveproxy80@gmail.com> (https://github.com/dadadave80)
+/// @author David Dada <daveproxy80@gmail.com> (https://github.com/dadadave80)
+/// @author Modified from Nick Mudge (https://github.com/mudgen/diamond-3-hardhat/blob/main/contracts/Diamond.sol)
 abstract contract Diamond {
     /// @notice Initializes the Diamond proxy with the provided facets and initialization parameters
     /// @param _init Address of the initialization contract
     /// @param _calldata Calldata to be passed to the initialization contract
-    constructor(FacetCut[] memory _facetCuts, address _init, bytes memory _calldata) payable {
-        _diamondCut(_facetCuts, _init, _calldata);
+    function initialize(FacetCut[] calldata _facetCuts, address _init, bytes calldata _calldata)
+        external
+        payable
+        virtual
+    {
+        bytes32 s = InitializableLib.initializableSlot();
+        InitializableLib.preInitializer(s);
+
+        DiamondLib.diamondCutCalldata(_facetCuts, _init, _calldata);
+
+        InitializableLib.postInitializer(s);
     }
 
     /// @notice Fallback function that delegates calls to the appropriate facet based on function selector
     /// @dev Reads the facet address from diamond storage and performs a delegatecall; reverts if selector is not found
     fallback() external payable virtual {
-        _beforeDelegate();
         _delegate(_facet());
     }
 
+    receive() external payable virtual {}
+
+    /// @notice Retrieves the implementation address for the current function call
+    /// @dev A Facet is one of many implementations in a Diamond Proxy
+    function _facet() internal view virtual returns (address) {
+        return DiamondLib.selectorToFacet(msg.sig);
+    }
+
     /// @notice Internal function to perform a delegatecall to an implementation
     /// @param _implementation Address of the implementation to delegate to
     function _delegate(address _implementation) internal virtual {
@@ -72,19 +89,4 @@ abstract contract Diamond {
             }
         }
     }
-
-    /// @notice Internal function to perform a diamond cut
-    function _diamondCut(FacetCut[] memory _facetCuts, address _init, bytes memory _calldata) internal virtual {
-        DiamondLib.diamondCut(_facetCuts, _init, _calldata);
-    }
-
-    /// @notice Internal hook function to run before a delegatecall to the facet
-    /// @dev This function can be replaced to perform additional logic before the delegatecall
-    function _beforeDelegate() internal virtual {}
-
-    /// @notice Retrieves the implementation address for the current function call
-    /// @dev A Facet is one of many implementations in a Diamond Proxy
-    function _facet() internal view virtual returns (address) {
-        return DiamondLib.selectorToFacet(msg.sig);
-    }
 }

src/facets/DiamondCutFacet.sol

@@ -22,7 +22,7 @@ contract DiamondCutFacet is IDiamondCut {
     ///                  _calldata is executed with delegatecall on _init
     function diamondCut(FacetCut[] calldata _diamondCut, address _init, bytes calldata _calldata) external payable {
         // Check that the caller is the owner
-        OwnableRolesLib._checkOwner();
+        OwnableRolesLib.checkOwner();
         // Call the diamond cut function from the library
         DiamondLib.diamondCutCalldata(_diamondCut, _init, _calldata);
     }

src/facets/OwnableRolesFacet.sol

@@ -24,47 +24,47 @@ contract OwnableRolesFacet {
 
     /// @dev Allows the owner to transfer the ownership to `newOwner`.
     function transferOwnership(address _newOwner) public onlyOwner {
-        _newOwner._transferOwnership();
+        _newOwner.transferOwnership();
     }
 
     /// @dev Allows the owner to renounce their ownership.
     function renounceOwnership() public onlyOwner {
-        OwnableRolesLib._renounceOwnership();
+        OwnableRolesLib.renounceOwnership();
     }
 
     /// @dev Request a two-step ownership handover to the caller.
     /// The request will automatically expire in 48 hours (172800 seconds) by default.
     function requestOwnershipHandover() public {
-        OwnableRolesLib._requestOwnershipHandover();
+        OwnableRolesLib.requestOwnershipHandover();
     }
 
     /// @dev Cancels the two-step ownership handover to the caller, if any.
     function cancelOwnershipHandover() public {
-        OwnableRolesLib._cancelOwnershipHandover();
+        OwnableRolesLib.cancelOwnershipHandover();
     }
 
     /// @dev Allows the owner to complete the two-step ownership handover to `pendingOwner`.
     /// Reverts if there is no existing ownership handover requested by `pendingOwner`.
     function completeOwnershipHandover(address _pendingOwner) public onlyOwner {
-        _pendingOwner._completeOwnershipHandover();
+        _pendingOwner.completeOwnershipHandover();
     }
 
     /// @dev Allows the owner to grant `user` `roles`.
     /// If the `user` already has a role, then it will be an no-op for the role.
     function grantRoles(address _user, uint256 _roles) public onlyOwner {
-        _user._grantRoles(_roles);
+        _user.grantRoles(_roles);
     }
 
     /// @dev Allows the owner to remove `user` `roles`.
     /// If the `user` does not have a role, then it will be an no-op for the role.
     function revokeRoles(address _user, uint256 _roles) public onlyOwner {
-        _user._removeRoles(_roles);
+        _user.removeRoles(_roles);
     }
 
     /// @dev Allow the caller to remove their own roles.
     /// If the caller does not have a role, then it will be an no-op for the role.
     function renounceRoles(uint256 _roles) public {
-        msg.sender._removeRoles(_roles);
+        msg.sender.removeRoles(_roles);
     }
 
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
@@ -73,41 +73,41 @@ contract OwnableRolesFacet {
 
     /// @dev Returns the owner of the contract.
     function owner() public view returns (address) {
-        return OwnableRolesLib._owner();
+        return OwnableRolesLib.owner();
     }
 
     /// @dev Returns the expiry timestamp for the two-step ownership handover to `pendingOwner`.
     function ownershipHandoverExpiresAt(address _pendingOwner) public view returns (uint256) {
-        return _pendingOwner._ownershipHandoverExpiresAt();
+        return _pendingOwner.ownershipHandoverExpiresAt();
     }
 
     /// @dev Returns the roles of `user`.
     function rolesOf(address _user) public view returns (uint256) {
-        return _user._rolesOf();
+        return _user.rolesOf();
     }
 
     /// @dev Returns whether `user` has any of `roles`.
     function hasAnyRole(address _user, uint256 _roles) public view returns (bool) {
-        return _user._rolesOf() & _roles != 0;
+        return _user.rolesOf() & _roles != 0;
     }
 
     /// @dev Returns whether `user` has all of `roles`.
     function hasAllRoles(address _user, uint256 _roles) public view returns (bool) {
-        return _user._rolesOf() & _roles == _roles;
+        return _user.rolesOf() & _roles == _roles;
     }
 
     /// @dev Convenience function to return a `roles` bitmap from an array of `ordinals`.
     /// This is meant for frontends like Etherscan, and is therefore not fully optimized.
     /// Not recommended to be called on-chain.
     function rolesFromOrdinals(uint8[] memory ordinals) external pure returns (uint256) {
-        return ordinals._rolesFromOrdinals();
+        return ordinals.rolesFromOrdinals();
     }
 
     /// @dev Convenience function to return an array of `ordinals` from the `roles` bitmap.
     /// This is meant for frontends like Etherscan, and is therefore not fully optimized.
     /// Not recommended to be called on-chain.
     function ordinalsFromRoles(uint256 roles) external pure returns (uint8[] memory) {
-        return roles._ordinalsFromRoles();
+        return roles.ordinalsFromRoles();
     }
 
     /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
@@ -116,27 +116,27 @@ contract OwnableRolesFacet {
 
     /// @dev Marks a function as only callable by the owner.
     modifier onlyOwner() {
-        OwnableRolesLib._checkOwner();
+        OwnableRolesLib.checkOwner();
         _;
     }
 
     /// @dev Marks a function as only callable by an account with `roles`.
     modifier onlyRoles(uint256 _roles) {
-        OwnableRolesLib._checkRoles(_roles);
+        OwnableRolesLib.checkRoles(_roles);
         _;
     }
 
     /// @dev Marks a function as only callable by the owner or by an account
     ///      with `roles`. Checks for ownership first, then lazily checks for roles.
     modifier onlyOwnerOrRoles(uint256 _roles) {
-        OwnableRolesLib._checkOwnerOrRoles(_roles);
+        OwnableRolesLib.checkOwnerOrRoles(_roles);
         _;
     }
 
     /// @dev Marks a function as only callable by an account with `roles` or the owner.
     /// Checks for roles first, then lazily checks for ownership.
     modifier onlyRolesOrOwner(uint256 _roles) {
-        OwnableRolesLib._checkRolesOrOwner(_roles);
+        OwnableRolesLib.checkRolesOrOwner(_roles);
         _;
     }
 }

src/initializers/DiamondInit.sol

@@ -5,18 +5,20 @@ import {DiamondLib, DiamondStorage} from "@diamond/libraries/DiamondLib.sol";
 import {OwnableRolesLib} from "@diamond/libraries/OwnableRolesLib.sol";
 
 /// @title DiamondInit
-/// @notice Provides an initializer to register standard interface support (ERC-165, ERC-173, IDiamondCut, IDiamondLoupe)
+/// @notice Provides a combined initializer to set up ownership and register standard interface support
 /// @author Nick Mudge (https://github.com/mudgen/diamond-3-hardhat/blob/main/contracts/Diamond.sol)
 /// @author Modified by David Dada <daveproxy80@gmail.com> (https://github.com/dadadave80)
 ///
-/// @dev Intended to be called as the `initDiamond` function in a diamond cut to set up ERC-165 interface IDs
+/// @dev Intended to be called as the `initDiamond` function in a diamond cut.
+///      For granular initialization, use OwnableInit and ERC165Init separately via MultiInit.
 contract DiamondInit {
-    /// @notice Initialize the contract with the ERC165 interface support.
+    /// @notice Initialize the contract owner and register standard diamond ERC-165 interface IDs.
     /// @dev This function is called during the diamond cut process to set up
     ///      the initial state of the contract.
+    /// @param _owner The address to set as the contract owner.
     function initDiamond(address _owner) public {
         // Initialize the owner
-        OwnableRolesLib._initializeOwner(_owner);
+        OwnableRolesLib.initializeOwner(_owner);
 
         DiamondStorage storage ds = DiamondLib.diamondStorage();
         /// @dev type(ERC165).interfaceId

src/initializers/ERC165Init.sol

@@ -0,0 +1,25 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {DiamondLib, DiamondStorage} from "@diamond/libraries/DiamondLib.sol";
+
+/// @title ERC165Init
+/// @notice Provides an initializer to register standard diamond interface support (ERC-165, ERC-173, IDiamondCut, IDiamondLoupe)
+/// @author David Dada <daveproxy80@gmail.com> (https://github.com/dadadave80)
+///
+/// @dev Intended to be called as a standalone initializer via diamond cut or through MultiInit
+contract ERC165Init {
+    /// @notice Register standard diamond ERC-165 interface IDs.
+    /// @dev Sets support for ERC-165, ERC-173, IDiamondCut, and IDiamondLoupe.
+    function initERC165() public {
+        DiamondStorage storage ds = DiamondLib.diamondStorage();
+        /// @dev type(ERC165).interfaceId
+        ds.supportedInterfaces[0x01ffc9a7] = true;
+        /// @dev type(IERC173).interfaceId
+        ds.supportedInterfaces[0x7f5828d0] = true;
+        /// @dev type(IDiamondCut).interfaceId
+        ds.supportedInterfaces[0x1f931c1c] = true;
+        /// @dev type(IDiamondLoupe).interfaceId
+        ds.supportedInterfaces[0x48e2b093] = true;
+    }
+}

src/initializers/MultiInit.sol

@@ -1,11 +1,11 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.20;
 
-import {DiamondLib} from "@diamond/libraries/DiamondLib.sol";
-
 /// @notice Thrown when the length of the address array does not match the length of the calldata array.
 /// @dev Used in initializer logic to ensure one-to-one mapping between addresses and initialization calldata.
 error AddressAndCalldataLengthMismatch();
+error NoBytecodeAtAddress(address initAddress);
+error InitializeReverted(address initAddress, bytes initCalldata);
 
 contract MultiInit {
     function multiInit(address[] calldata _initAddresses, bytes[] calldata _initData) public {
@@ -15,7 +15,21 @@ contract MultiInit {
         }
 
         for (uint256 i; i < initAddressesLength; ++i) {
-            DiamondLib.initializeDiamondCut(_initAddresses[i], _initData[i]);
+            address initAddress = _initAddresses[i];
+            if (initAddress == address(0)) return;
+            if (initAddress.code.length == 0) revert NoBytecodeAtAddress(initAddress);
+            (bool success, bytes memory err) = initAddress.delegatecall(_initData[i]);
+            if (!success) {
+                if (err.length > 0) {
+                    // bubble up error
+                    assembly ("memory-safe") {
+                        let returndata_size := mload(err)
+                        revert(add(32, err), returndata_size)
+                    }
+                } else {
+                    revert InitializeReverted(initAddress, _initData[i]);
+                }
+            }
         }
     }
 }

src/initializers/OwnableInit.sol

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import {OwnableRolesLib} from "@diamond/libraries/OwnableRolesLib.sol";
+
+/// @title OwnableInit
+/// @notice Provides an initializer to set the contract owner
+/// @author David Dada <daveproxy80@gmail.com> (https://github.com/dadadave80)
+///
+/// @dev Intended to be called as a standalone initializer via diamond cut or through MultiInit
+contract OwnableInit {
+    /// @notice Initialize the contract owner.
+    /// @dev This function is called during the diamond cut process to set up the owner.
+    /// @param _owner The address to set as the contract owner.
+    function initOwner(address _owner) public {
+        OwnableRolesLib.initializeOwner(_owner);
+    }
+}