我们需要针对上面提出的问题进行优化,通过观察分红机制可以看出,所有account对应的balance都是按比例增加,那么这里的循环更新是否有存在的必要呢?是否可以全局记录一个系数,每次调用distribute时只更新这个系数,要查询balance时,乘上这个系数即可
那么这里就拆分出两个概念,balance和share,并且衍生出sharePrice,对于所有进入rebase的用户都只记录share,并且有一个全局的sharePrice,查询用户balance时需要根据share和sharePrice计算出balance,sharePrice在分红时增加
针对第一个版本,我们梳理一下需要修改的部分和积攒的问题:
具体的改动如下:
struct RebasingState {
bool isRebasing;
uint256 nShares;
}
mapping(address => RebasingState) internal rebasingAccount;
uint256 public totalShares;
uint256 public sharePrice = 1e30;
用下面的函数处理share和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 ;
}
继续改写enter/exit和distribute部分:
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的几个方法进行重写:
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;
}
大致逻辑已经实现了,用户的balance也能通过share和sharePrice来正确计算,我们再来看是否满足了前面的恒等式:
totalSupply() == nonRebasingSupply() + rebasingSupply()
sum of balanceOf(x) == totalSupply()
那么缺少的这部分是什么?
// 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;
}
}