Design and Implementation of ERC20 Rebase - 2 Share and Balance

Reprinted from personal blog https://www.hackdefi.xyz/posts/erc20-rebase-2/

We need to optimize the issues raised above. By observing the dividend mechanism, it can be seen that the balance corresponding to all accounts increases proportionally. So is there a need for this loop update? Can we globally record a coefficient, updating only this coefficient each time we call distribute, and when querying the balance, simply multiply by this coefficient?

Here we split into two concepts: balance and share, and derive sharePrice. For all users entering rebase, only share is recorded, and there is a global sharePrice. When querying a user's balance, it needs to be calculated based on share and sharePrice, with sharePrice increasing during dividends.

For the first version, let's sort out the parts that need to be modified and the accumulated issues:

Need to record the share of each address participating in rebase
How to map share to balance
How sharePrice accumulates during each dividend

The specific changes are as follows:

Expand the mapping value to a struct to record share
Change rebasingSupply to totalShares
Define initial sharePrice=1e30

    struct RebasingState {
        bool isRebasing;
        uint256 nShares;
    }
    mapping(address => RebasingState) internal rebasingAccount;
    uint256 public totalShares;
    uint256 public sharePrice = 1e30;

Use the following functions to handle the mapping relationship between share and balance:

    function rebasingSupply() public view returns (uint256) {
        return share2Balance(totalShares);
    }

    function nonRebasingSupply() public view returns (uint256) {
        return totalSupply() - rebasingSupply();
    }

    function share2Balance(uint256 shares) view public returns (uint256) {
        return shares * sharePrice / 1e30;
    }

    function balance2Share(uint256 balance) view public returns (uint256) {
        return balance * 1e30 / sharePrice ;
    }

Continue rewriting the enter/exit and distribute parts:

When entering, first convert balance to share and update the rebase record
When exiting, convert share to balance and remove the rebase record
During distribution, first burn the token, then update sharePrice

    function _enterRebase(address user) internal {
        uint256 balance = balanceOf(user);
        uint256 shares = balance2Share(balance);
        rebasingAccount[user].isRebasing = true;
        rebasingAccount[user].nShares = shares;
        totalShares += shares;
        emit RebaseEnter(user, shares, block.timestamp);
    }

    function _exitRebase(address user)  internal {
        uint256 shares = rebasingAccount[user].nShares;
        rebasingAccount[user].isRebasing = false;
        rebasingAccount[user].nShares = 0;
        totalShares -= shares;
        emit RebaseExit(user, shares, block.timestamp);
    }

    function distribute(uint256 amount) external {
        require(balanceOf(msg.sender)>=amount, "SimpleERC20Rebase: not enough");
        _burn(msg.sender, amount);
        sharePrice += amount*1e30 / totalShares;
    }

ERC20 Method Overriding#

At the same time, we need to override several methods of ERC20:

balanceOf: For users participating in rebase, it needs to convert share to balance
mint burn transfer transferFrom: For relevant users, if they previously participated in rebase, they need to exit first, then perform the operation and re-enter.

    function balanceOf(address account) public view override returns (uint256) {
        uint256 rawBalance = ERC20.balanceOf(account);
        if (rebasingAccount[account].isRebasing) {
            return share2Balance(rebasingAccount[account].nShares);
        } else {
            return rawBalance;
        }
    }

    function mint(address user, uint256 amount) external {
        bool isRebasing = rebasingAccount[user].isRebasing;
        if (isRebasing) {
            _exitRebase(user);
        }
        ERC20._mint(user, amount);
        if (isRebasing) {
            _enterRebase(user);
        }
    }
    
    function burn(address user, uint256 amount) external {
        bool isRebasing = rebasingAccount[user].isRebasing;
        if (isRebasing) {
            _exitRebase(user);
        }
        ERC20._burn(user, amount);
        if (isRebasing) {
            _enterRebase(user);
        }
    }
    
    function transfer(address to, uint256 amount)  public virtual override returns (bool) {
        bool isFromRebasing = rebasingAccount[msg.sender].isRebasing;
        bool isToRebasing = rebasingAccount[to].isRebasing;
        if (isFromRebasing) {
            _exitRebase(msg.sender);
        }
        if (isToRebasing && to != msg.sender) {
            _exitRebase(to);
        }
        bool result = ERC20.transfer(to, amount);
        if (isFromRebasing) {
            _enterRebase(msg.sender);
        }
        if (isToRebasing && to != msg.sender) {
            _enterRebase(to);
        }
        return result;
    }

    function transferFrom(
        address from,
        address to,
        uint256 amount
    ) public virtual override returns (bool) {
        bool isFromRebasing = rebasingAccount[from].isRebasing;
        bool isToRebasing = rebasingAccount[to].isRebasing;
        if (isFromRebasing) {
            _exitRebase(from);
        }
        if (isToRebasing && to != from) {
            _exitRebase(to);
        }
        bool result = ERC20.transfer(to, amount);
        if (isFromRebasing) {
            _enterRebase(from);
        }
        if (isToRebasing && to != from) {
            _enterRebase(to);
        }
        return result;
    }

Remaining Issues#

The general logic has been implemented, and the user's balance can also be correctly calculated through share and sharePrice. Now let's see if it satisfies the previous identity:

totalSupply() == nonRebasingSupply() + rebasingSupply()
- rebasingSupply can be calculated using the current share and sharePrice, but totalSupply has not been rewritten, totalSupply=nonRebasingSupply. For example, nonRebasingSupply has decreased, but share remains unchanged, and sharePrice has increased, so obviously totalSupply is less than rebasingSupply.
sum of balanceOf(x) == totalSupply()
- Similarly, the user's share remains unchanged but sharePrice has increased, so obviously totalSupply is also too small.

So what is missing in this part?

Code#

// SPDX-License-Identifier: GPL-3.0-or-later
pragma solidity 0.8.13;

import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";

contract SimpleERC20Rebase is ERC20 {
    event RebaseEnter(address indexed account, uint256 indexed shares, uint256 indexed timestamp);
    event RebaseExit(address indexed account, uint256 indexed shares, uint256 indexed timestamp);
    
    struct RebasingState {
        bool isRebasing;
        uint256 nShares;
    }
    mapping(address => RebasingState) internal rebasingAccount;
    uint256 public totalShares;
    uint256 public sharePrice = 1e30;

    constructor(
        string memory _name,
        string memory _symbol
    ) ERC20(_name, _symbol) {}

    function rebasingSupply() public view returns (uint256) {
        return share2Balance(totalShares);
    }

    function nonRebasingSupply() public view returns (uint256) {
        return totalSupply() - rebasingSupply();
    }

    function share2Balance(uint256 shares) view public returns (uint256) {
        return shares * sharePrice / 1e30;
    }

    function balance2Share(uint256 balance) view public returns (uint256) {
        return balance * 1e30 / sharePrice ;
    }

    function enterRebase() external {
        require(!rebasingAccount[msg.sender].isRebasing, "SimpleERC20Rebase: already rebasing");
        _enterRebase(msg.sender);
    }

    function _enterRebase(address user) internal {
        uint256 balance = balanceOf(user);
        uint256 shares = balance2Share(balance);
        rebasingAccount[user].isRebasing = true;
        rebasingAccount[user].nShares = shares;
        totalShares += shares;
        emit RebaseEnter(user, shares, block.timestamp);
    }

    function exitRebase()  external {
        require(rebasingAccount[msg.sender].isRebasing, "SimpleERC20Rebase: not rebasing");
        _exitRebase(msg.sender);
    }

    function _exitRebase(address user)  internal {
        uint256 shares = rebasingAccount[user].nShares;
        rebasingAccount[user].isRebasing = false;
        rebasingAccount[user].nShares = 0;
        totalShares -= shares;
        emit RebaseExit(user, shares, block.timestamp);
    }

    function distribute(uint256 amount) external {
        require(balanceOf(msg.sender)>=amount, "SimpleERC20Rebase: not enough");
        _burn(msg.sender, amount);
        sharePrice += amount*1e30 / totalShares;
    }

    function balanceOf(address account) public view override returns (uint256) {
        uint256 rawBalance = ERC20.balanceOf(account);
        if (rebasingAccount[account].isRebasing) {
            return share2Balance(rebasingAccount[account].nShares);
        } else {
            return rawBalance;
        }
    }

    function mint(address user, uint256 amount) external {
        bool isRebasing = rebasingAccount[user].isRebasing;
        if (isRebasing) {
            _exitRebase(user);
        }
        ERC20._mint(user, amount);
        if (isRebasing) {
            _enterRebase(user);
        }
    }
    
    function burn(address user, uint256 amount) external {
        bool isRebasing = rebasingAccount[user].isRebasing;
        if (isRebasing) {
            _exitRebase(user);
        }
        ERC20._burn(user, amount);
        if (isRebasing) {
            _enterRebase(user);
        }
    }
    
    function transfer(address to, uint256 amount)  public virtual override returns (bool) {
        bool isFromRebasing = rebasingAccount[msg.sender].isRebasing;
        bool isToRebasing = rebasingAccount[to].isRebasing;
        if (isFromRebasing) {
            _exitRebase(msg.sender);
        }
        if (isToRebasing && to != msg.sender) {
            _exitRebase(to);
        }
        bool result = ERC20.transfer(to, amount);
        if (isFromRebasing) {
            _enterRebase(msg.sender);
        }
        if (isToRebasing && to != msg.sender) {
            _enterRebase(to);
        }
        return result;
    }

    function transferFrom(
        address from,
        address to,
        uint256 amount
    ) public virtual override returns (bool) {
        bool isFromRebasing = rebasingAccount[from].isRebasing;
        bool isToRebasing = rebasingAccount[to].isRebasing;
        if (isFromRebasing) {
            _exitRebase(from);
        }
        if (isToRebasing && to != from) {
            _exitRebase(to);
        }
        bool result = ERC20.transfer(to, amount);
        if (isFromRebasing) {
            _enterRebase(from);
        }
        if (isToRebasing && to != from) {
            _enterRebase(to);
        }
        return result;
    }
}