Example of Adding New State Variables in Diamond Upgrade

The purpose of this article is to show how to add new state variables to an already deployed diamond and use them.

The purpose of this article is to show how to add new state variables to an already deployed diamond and use them.

Let’s say we deploy a diamond that uses the following AppStorage and PersonsFacet:

// AppStorage.sol
struct Person {
  string name;
}


struct AppStorage {
  uint256 personsCount;
  // personId to Person
  mapping(uint256 => Person) persons;
}

PersonsFacet.sol

// PersonsFacet.sol
import "./AppStorage.sol";

contract PersonsFacet {
  AppStorage storage s;

  function mintPerson(string calldata _name) external {
    uint256 count = s.personsCount;
    count++;
    s.personsCount = count;
    s.persons[count].name = _name;
  }

  function getPerson(uint256 _personId) 
    external 
    view 
    returns(Person memory person_) 
  {
    person_ = s.persons[_personId];
  }
}

We deploy the diamond and add the functions from PersonsFacet to the diamond.

After that we now want to add an `age` state variable to Person struct and we want the `getPerson` function to return the upgraded Person struct that includes the age.

And we also want to add two new friends functions. One function that adds friends to a person and a one function that returns the list of friends that a person has.

To do this we first modify the AppStorage source code like so:

// AppStorage.sol
struct Person {
  string name;
  uint256 age;
}

AppStorage {
  uint256 personsCount;
  // personId to Person
  mapping(uint256 => Person) persons;
  // personId to personIds
  mapping(uint256 => uint256[]) friends;
}

Note that it is fine to add new mappings and arrays and other types to the end of AppStorage. Also note that if you add a struct variable or struct array variable directly to AppStorage then you can’t extend the struct later. But if a struct is in a mapping then it can be extended later. In our example the `Person` struct is used in the `persons` mapping. Since the Person struct is in a mapping we can extend it by adding new state variables to it in an upgrade.

Now we edit the `mintPerson` function in`PersonsFacet.sol` so that it uses the new `age` state variable:

// PersonsFacet.sol
import "./AppStorage.sol";

contract PersonsFacet {
  AppStorage storage s;

  function mintPerson(string calldata _name, uint256 _age)
    external 
  {
    uint256 count = s.personsCount;
    count++;
    s.personsCount = count;
    s.persons[count] = Person(_name, _age);
  }

  function getPerson(uint256 _personId) 
    external 
    view 
    returns(Person memory person_) 
  {
    person_ = s.persons[_personId];
  }
}

Recompile the `PersonsFacet.sol`.

New let’s create a `FriendsFacet.sol` to create some friends functions. Note that this is not a complete implementation. It is used to demonstrate upgrades.

// FriendsFacet.sol
import "./AppStorage.sol";

contract PersonsFacet {
  AppStorage storage s;
 
  function addFriend(uint256 _personId, uint256 _friendId)
    external 
  {
    s.friends[_personId].push(_friendId);
  }

  function getFriendsIds(uint256 _personId) 
    external 
    view
    returns(uint256[] memory friendsIds_)
  {
    friendsIds_ = s.friends(_personId);
  }
}

Compile FriendsFacet.sol.

Deploy the new `PersonsFacet` and `FriendsFacet` contracts.

Then execute the `diamondCut` function to add the `mintPerson(string calldata _name, uint256 _age)`, `addFriend(uint256 _personId, uint256 _friendId)` and the `getFriendsIds(uint256 _personId)` functions, and replace the `getPerson(uint256 _personId) ` function because it already exists. You may also want to remove the old `mintPerson(string calldata _name)` function. Using the second and third parameters of `diamondCut` you can also execute an external function with delegatecall if you want to initialize any ages or friends for the upgrade. Note that adding/replacing/removing multiple functions and calling an initialization function can and should be done in a single transaction using `diamondCut`. A code example of using `diamondCut` is here: https://github.com/mudgen/diamond-1-hardhat/blob/main/scripts/deploy.js#L58

More information about upgrades on diamonds is here: Diamond Upgrades