A guide to extending ERC2535 Diamonds.

In this article, I introduce the Diamond standard for writing upgrade-able smart contracts, and give a walk through of how it can be extended,to add support for roll-back upgrades.

Pre-requisites:

• understanding ethereum and smart contracts.
• experience with solidity.
• basic terminal skills, hardhat and foundry installation.

For foundry installation, follow this link.

Brief introduction to Diamonds.

EIP2535 introduces the diamond standard as:

A modular smart contract systems that can be upgraded/extended after deployment, and have virtually no size limit. More technically, a diamond is a contract with external functions that are supplied by contracts called facets. Facets are separate, independent contracts that can share internal functions, libraries, and state variables.

In my opinion, what makes diamond standard shines is the clear separation of concern(a.k.a logic) baked into it from the ground up. This is made possible through what is called facets, which are smart contracts that groups logic that are related in one place. The default Diamond implementation used in this article includes three such facets described below:

1. DiamondCutFacet - This facet is mainly concerned with performing upgrade related logic on the Diamond.
2. DiamondLoupeFacet - This facet is mainly concerned with introspection into the diamond.
3. OwnershipFacet - Like the name implies, this facet is mainly concerned with matters related to ownership/admin/access-control on the diamond.

Now, onto the most important component, The Diamond itself, which is the proxy smart contract that enshrines the state of the whole system, and exposes the methods of the facets using delegatecall.
There are other components built into this default implementation, which includes libraries, interfaces and so on, but we'd discuss those later on as they're encountered.

To follow along in your preferred IDE, you can clone the repository and follow the ReadME to install the necessary dependencies.

The adventure begins... 🚗

Tracking the state of the diamond.

The Diamond for this demonstration, utilizes the diamond storage pattern, which declares the state of the Diamond as a storage struct, this is made possible via a feature of libraries in solidity, which allows for a library function to return a struct in storage, see here.

function diamondStorage() internal pure returns (DiamondStorage storage ds) {
        bytes32 position = DIAMOND_STORAGE_POSITION;
        assembly {
            ds.slot := position
        }
    }

This struct is then set in storage using the hash of a unique string, effectively preventing storage slot collision.
The library defines important data structure and logic that are shared amongst facets, thereby upholding the justifications of this standard (i.e Facets can share internal functions, libraries).

How upgrades are accomplished?

In diamond lingo, upgrade is synonymous to cut, the DiamondCutFacet implements the diamondCut to achieve that. However the struct FacetCut models the data unique to performing the said upgrade, i.e cutting the Diamond.

 struct FacetCut {
        address facetAddress;
        FacetCutAction action;
        bytes4[] functionSelectors;
    }

How can we roll-back an upgrade? 😐

First we define a custom RollBack struct, to model the data for performing a rollback

  struct RollBackCuts {
        address[] facetAddress;
        IDiamondCut.FacetCutAction action;
        bytes4[] functionSelectors;
    }

Next we append an array of this struct to the storage struct of the diamond, and subsequently populate this array with every upgrade to the diamond, see here, here and here, these functions adds, replaces and removes functions from the Diamond as part of an upgrade.

struct DiamondStorage {
        // maps function selector to the facet address and
        // the position of the selector in the facetFunctionSelectors.selectors array
        mapping(bytes4 => FacetAddressAndPosition) selectorToFacetAndPosition;
        // maps facet addresses to function selectors
        mapping(address => FacetFunctionSelectors) facetFunctionSelectors;
        // facet addresses
        address[] facetAddresses;
        // Used to query if a contract implements an interface.
        // Used to implement ERC-165.
        mapping(bytes4 => bool) supportedInterfaces;
        // owner of the contract
        address contractOwner;
        // history for rollback
        RollBackCuts[] rollbackCuts;
    }

A constant address, labeled as ROLLBACK_ADDRESS is defined in the library this:

address constant ROLLBACK_ADDRESS = 0x000000000000000000000000000000000000bEEF;

This constant address distinguishes a roll-back from an upgrade, hence preventing the duplication bug from multiple roll-backs.
Further more, the ROLLBACK_ADDRESS also serves to deter an initialization call on roll-back.

Then we proceed to define the logic for the rollback as part of the DiamondCutFacet:

function rollback() external override {
        // enforce is contract owner
        LibDiamond.enforceIsContractOwner();
        // LibDiamond.DiamondStorage storage ds = LibDiamond.diamondStorage();
        LibDiamond.RollBackCuts[] storage rollbackCuts = LibDiamond.diamondStorage().rollbackCuts;
        uint256 rollbackCutsLength = rollbackCuts.length;
        if(rollbackCutsLength == 0) {
            revert NoRollBackAction();
        }
        LibDiamond.RollBackCuts memory rollbackCut = rollbackCuts[rollbackCutsLength - 1];
        // remove last rollback action
        rollbackCuts.pop();
        uint256 cutLength = rollbackCut.facetAddress.length;
        IDiamondCut.FacetCut[] memory facetCut = new IDiamondCut.FacetCut[](
            cutLength
        );
        for (uint i = 0; i < cutLength; ++i) {
            bytes4[] memory rollbackSelector = new bytes4[](1);
            rollbackSelector[0] = rollbackCut.functionSelectors[i];

            facetCut[i] = IDiamondCut.FacetCut(
                rollbackCut.facetAddress[i],
                rollbackCut.action,
                rollbackSelector
            );
        }
        LibDiamond.diamondCut(
            facetCut,
            LibDiamond.ROLLBACK_ADDRESS,
            ""
        );
    }

The logic for the rollback effectively reads the storage struct, checks if a rollback is valid, constructs the data for rollback and proceeds to cut the diamond with the said data, using our beloved ROLLBACK_ADDRESS once again.

The Diamond once again. 💎

This time we dive into the logic of the Diamond constructor, and the main change here is registering the selector of our rollback implementation as seen here.

constructor(address _contractOwner, address _diamondCutFacet) payable {
        LibDiamond.setContractOwner(_contractOwner);

        // Add the diamondCut external function from the diamondCutFacet
        IDiamondCut.FacetCut[] memory cut = new IDiamondCut.FacetCut[](1);
        bytes4[] memory functionSelectors = new bytes4[](2);
        functionSelectors[0] = IDiamondCut.diamondCut.selector;
        // Add the rollback extxernal function from the diamondCutFacet
        functionSelectors[1] = IDiamondCut.rollback.selector;

        cut[0] = IDiamondCut.FacetCut({
            facetAddress: _diamondCutFacet,
            action: IDiamondCut.FacetCutAction.Add,
            functionSelectors: functionSelectors
        });

        LibDiamond.diamondCut(cut, address(0), "");
    }

With all major changes out of the way, we proceed to write test to validate our code:
The tests can be found here.
Now in your terminal excute the commands to compile and run the tests:

npx hardhat compile

forge test

The test essentially deploys, upgrade, and execute a roll-back on the diamond.
Verify that it passes, and viola!!! 🎉 We have basically extended an ERC2535 Diamond to add support for roll-back, which brings us to the end of this article.

If you enjoy reading code, i encourage an in-depth look at the Repo.

Once again, congrats on making it to the end, and I hope you found this article useful and can proceed to implementing Diamonds in the wild. Reach me on twitter @nuel