Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

refactor: Auto _setBalanceInfo #20

Merged
merged 2 commits into from
Jul 2, 2024
Merged

refactor: Auto _setBalanceInfo #20

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
64b6cd4
refactor: Auto `_setBalanceInfo`
deluca-mike Jun 28, 2024
67a75a7
fix: PR Review
deluca-mike Jul 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
169 changes: 77 additions & 92 deletions src/WrappedMToken.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,77 +73,67 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
function startEarningFor(address account_) external {
if (!_isApprovedEarner(account_)) revert NotApprovedEarner();

(bool isEarning_, , uint240 rawBalance_) = _getBalanceInfo(account_);
(bool isEarning_, , , uint240 balance_) = _getBalanceInfo(account_);

if (isEarning_) return;

emit StartedEarning(account_);

uint128 currentIndex_ = currentIndex();
uint112 principalAmount_ = IndexingMath.getPrincipalAmountRoundedDown(rawBalance_, currentIndex_);

_setBalanceInfo(account_, true, currentIndex_, principalAmount_);
_setBalanceInfo(account_, true, currentIndex_, balance_);
_addTotalEarningSupply(balance_, currentIndex_);

unchecked {
totalNonEarningSupply -= rawBalance_;
totalNonEarningSupply -= balance_;
}

_addTotalEarningSupply(rawBalance_, currentIndex_);
}

function stopEarningFor(address account_) external {
if (_isApprovedEarner(account_)) revert IsApprovedEarner();

(bool isEarning_, , ) = _getBalanceInfo(account_);

if (!isEarning_) return;

emit StoppedEarning(account_);

uint128 currentIndex_ = currentIndex();

_claim(account_, currentIndex_);

(, uint128 index_, uint256 rawBalance_) = _getBalanceInfo(account_);
(bool isEarning_, , , uint240 balance_) = _getBalanceInfo(account_);

uint240 amount_ = IndexingMath.getPresentAmountRoundedDown(uint112(rawBalance_), index_);
if (!isEarning_) return;

emit StoppedEarning(account_);

_setBalanceInfo(account_, false, 0, amount_);
_setBalanceInfo(account_, false, 0, balance_);
_subtractTotalEarningSupply(balance_, currentIndex_);

unchecked {
totalNonEarningSupply += amount_;
totalNonEarningSupply += balance_;
}

_subtractTotalEarningSupply(amount_, currentIndex_);
}

/* ============ View/Pure Functions ============ */

function accruedYieldOf(address account_) external view returns (uint240 yield_) {
(bool isEarning_, uint128 index_, uint240 rawBalance_) = _getBalanceInfo(account_);
(bool isEarning_, , uint112 principal_, uint240 balance_) = _getBalanceInfo(account_);

return isEarning_ ? _getAccruedYield(uint112(rawBalance_), index_, currentIndex()) : 0;
return isEarning_ ? (IndexingMath.getPresentAmountRoundedDown(principal_, currentIndex()) - balance_) : 0;
}

function balanceOf(address account_) external view returns (uint256 balance_) {
(bool isEarning_, uint128 index_, uint240 rawBalance_) = _getBalanceInfo(account_);

return isEarning_ ? IndexingMath.getPresentAmountRoundedDown(uint112(rawBalance_), index_) : rawBalance_;
(, , , balance_) = _getBalanceInfo(account_);
}

function currentIndex() public view returns (uint128 index_) {
return IMTokenLike(mToken).currentIndex();
}

function isEarning(address account_) external view returns (bool isEarning_) {
(isEarning_, , ) = _getBalanceInfo(account_);
(isEarning_, , , ) = _getBalanceInfo(account_);
}

function excess() public view returns (uint240 yield_) {
uint240 balance_ = uint240(IMTokenLike(mToken).balanceOf(address(this)));
uint240 earmarked_ = uint240(totalSupply()) + totalAccruedYield();

unchecked {
uint240 balance_ = uint240(IMTokenLike(mToken).balanceOf(address(this)));
uint240 earmarked_ = uint240(totalSupply()) + totalAccruedYield();
return balance_ > earmarked_ ? balance_ - earmarked_ : 0;
}
}
Expand All @@ -168,7 +158,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {

emit Transfer(address(0), recipient_, amount_);

(bool isEarning_, , ) = _getBalanceInfo(recipient_);
(bool isEarning_, , , ) = _getBalanceInfo(recipient_);

if (!isEarning_) return _addNonEarningAmount(recipient_, amount_);

Expand All @@ -185,7 +175,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {

emit Transfer(msg.sender, address(0), amount_);

(bool isEarning_, , ) = _getBalanceInfo(account_);
(bool isEarning_, , , ) = _getBalanceInfo(account_);

if (!isEarning_) return _subtractNonEarningAmount(account_, amount_);

Expand All @@ -198,61 +188,60 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
}

function _addNonEarningAmount(address recipient_, uint240 amount_) internal {
(, , uint240 rawBalance_) = _getBalanceInfo(recipient_);

unchecked {
_setBalanceInfo(recipient_, false, 0, rawBalance_ + amount_);
(, , , uint240 balance_) = _getBalanceInfo(recipient_);
_setBalanceInfo(recipient_, false, 0, balance_ + amount_);
totalNonEarningSupply += amount_;
}
}

function _subtractNonEarningAmount(address account_, uint240 amount_) internal {
(, , uint240 rawBalance_) = _getBalanceInfo(account_);
unchecked {
(, , , uint240 balance_) = _getBalanceInfo(account_);

if (rawBalance_ < amount_) revert InsufficientBalance(account_, rawBalance_, amount_);
if (balance_ < amount_) revert InsufficientBalance(account_, balance_, amount_);

unchecked {
_setBalanceInfo(account_, false, 0, rawBalance_ - amount_);
_setBalanceInfo(account_, false, 0, balance_ - amount_);
totalNonEarningSupply -= amount_;
}
}

function _addEarningAmount(address recipient_, uint240 amount_, uint128 currentIndex_) internal {
(, , uint240 rawBalance_) = _getBalanceInfo(recipient_);

uint112 principalAmount_ = IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_);

unchecked {
_setBalanceInfo(recipient_, true, currentIndex_, rawBalance_ + principalAmount_);
}
(, , , uint240 balance_) = _getBalanceInfo(recipient_);

_addTotalEarningSupply(amount_, currentIndex_);
_setBalanceInfo(recipient_, true, currentIndex_, balance_ + amount_);
_addTotalEarningSupply(amount_, currentIndex_);
}
}

function _subtractEarningAmount(address account_, uint240 amount_, uint128 currentIndex_) internal {
(, , uint240 rawBalance_) = _getBalanceInfo(account_);

uint112 principalAmount_ = IndexingMath.getPrincipalAmountRoundedUp(amount_, currentIndex_);
unchecked {
(, , , uint240 balance_) = _getBalanceInfo(account_);

if (rawBalance_ < principalAmount_) revert InsufficientBalance(account_, rawBalance_, principalAmount_);
if (balance_ < amount_) revert InsufficientBalance(account_, balance_, amount_);

unchecked {
_setBalanceInfo(account_, true, currentIndex_, rawBalance_ - principalAmount_);
_setBalanceInfo(account_, true, currentIndex_, balance_ - amount_);
_subtractTotalEarningSupply(amount_, currentIndex_);
}
}

function _claim(address account_, uint128 currentIndex_) internal returns (uint240 yield_) {
(bool isEarner_, uint128 index_, uint240 rawBalance_) = _getBalanceInfo(account_);
(bool isEarner_, uint128 index_, , uint240 startingBalance_) = _getBalanceInfo(account_);

if (!isEarner_) return 0;

yield_ = _getAccruedYield(uint112(rawBalance_), index_, currentIndex_);
_setBalanceInfo(account_, true, currentIndex_, rawBalance_);
if (currentIndex_ == index_) return 0;

_updateIndex(account_, currentIndex_);
deluca-mike marked this conversation as resolved.
Show resolved Hide resolved
Copy link
Contributor

@toninorair toninorair Jul 1, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

different suggestion: move yield calculation inside updateIndex. Rename updateIndex into updateEarnerIndex to improve comprehension and readability.

uint256 yield_ = updateEarnerIndex(account_, currentIndex_);

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I prefer not to since it will result in a cleaner change to the next PR that tackles start/stop (as well as more complex upgrades). Further, since _updateIndex never needs to compute yield, doing this creates unnecessary coupling to overcomplicate an otherwise straightforward and single-task _updateIndex.


if (yield_ == 0) return 0;
(, , , uint240 endingBalance_) = _getBalanceInfo(account_);

unchecked {
yield_ = endingBalance_ - startingBalance_;

if (yield_ == 0) return 0;

_setTotalEarningSupply(totalEarningSupply() + yield_, _principalOfTotalEarningSupply);
}

Expand All @@ -264,15 +253,26 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
} else {
emit Claimed(account_, claimOverrideRecipient_, yield_);

// NOTE: Watch out for a long chain of claim override recipients.
// TODO: Maybe can be optimized since we know `account_` is an earner and already claimed.
// NOTE: Watch out for a long chain of earning claim override recipients.
_transfer(account_, claimOverrideRecipient_, yield_, currentIndex_);
}
}

function _updateIndex(address account_, uint128 index_) internal {
uint256 unwrapped_ = BalanceInfo.unwrap(_balances[account_]);

unwrapped_ &= ~(uint256(type(uint112).max) << 128);

_balances[account_] = BalanceInfo.wrap(unwrapped_ | (uint256(index_) << 112));
}

function _setBalanceInfo(address account_, bool isEarning_, uint128 index_, uint240 amount_) internal {
_balances[account_] = isEarning_
? BalanceInfo.wrap((uint256(1) << 248) | (uint256(index_) << 112) | uint256(amount_))
? BalanceInfo.wrap(
(uint256(1) << 248) |
(uint256(index_) << 112) |
uint256(IndexingMath.getPrincipalAmountRoundedDown(amount_, index_))
)
: BalanceInfo.wrap(uint256(amount_));
}

Expand All @@ -284,8 +284,8 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {

emit Transfer(sender_, recipient_, amount_);

(bool senderIsEarning_, , ) = _getBalanceInfo(sender_);
(bool recipientIsEarning_, , ) = _getBalanceInfo(recipient_);
(bool senderIsEarning_, , , ) = _getBalanceInfo(sender_);
(bool recipientIsEarning_, , , ) = _getBalanceInfo(recipient_);

senderIsEarning_
? _subtractEarningAmount(sender_, amount_, currentIndex_)
Expand All @@ -301,54 +301,39 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
}

function _addTotalEarningSupply(uint240 amount_, uint128 currentIndex_) internal {
uint112 principalAmount_ = IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_);

unchecked {
_setTotalEarningSupply(totalEarningSupply() + amount_, _principalOfTotalEarningSupply + principalAmount_);
uint112 principal_ = IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_);
_setTotalEarningSupply(totalEarningSupply() + amount_, _principalOfTotalEarningSupply + principal_);
}
}

function _subtractTotalEarningSupply(uint240 amount_, uint128 currentIndex_) internal {
// TODO: Consider `getPrincipalAmountRoundedUp` .
uint112 principalAmount_ = IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_);

unchecked {
_setTotalEarningSupply(totalEarningSupply() - amount_, _principalOfTotalEarningSupply - principalAmount_);
// TODO: Consider `getPrincipalAmountRoundedUp` .
uint112 principal_ = IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_);
_setTotalEarningSupply(totalEarningSupply() - amount_, _principalOfTotalEarningSupply - principal_);
}
}

function _setTotalEarningSupply(uint240 amount_, uint112 principalAmount_) internal {
_indexOfTotalEarningSupply = principalAmount_ == 0
? 0
: IndexingMath.divide240by112Down(amount_, principalAmount_);

_principalOfTotalEarningSupply = principalAmount_;
function _setTotalEarningSupply(uint240 amount_, uint112 principal_) internal {
_indexOfTotalEarningSupply = (principal_ == 0) ? 0 : IndexingMath.divide240by112Down(amount_, principal_);
_principalOfTotalEarningSupply = principal_;
}

/* ============ Internal View/Pure Functions ============ */

function _getAccruedYield(
uint112 principalAmount_,
uint128 index_,
uint128 currentIndex_
) internal pure returns (uint240) {
unchecked {
return
currentIndex_ <= index_
? 0
: IndexingMath.getPresentAmountRoundedDown(principalAmount_, currentIndex_ - index_);
}
}

function _getBalanceInfo(
address account_
) internal view returns (bool isEarning_, uint128 index_, uint240 rawBalance_) {
) internal view returns (bool isEarning_, uint128 index_, uint112 principal_, uint240 balance_) {
uint256 unwrapped_ = BalanceInfo.unwrap(_balances[account_]);

return
(unwrapped_ >> 248) != 0
? (true, uint128((unwrapped_ << 8) >> 120), uint112(unwrapped_))
: (false, uint128(0), uint240(unwrapped_));
isEarning_ = (unwrapped_ >> 248) != 0;

if (!isEarning_) return (isEarning_, uint128(0), uint112(0), uint240(unwrapped_));

index_ = uint128((unwrapped_ << 8) >> 120);
principal_ = uint112(unwrapped_);
balance_ = IndexingMath.getPresentAmountRoundedDown(principal_, index_);
}

function _getClaimOverrideRecipient(address account_) internal view returns (address) {
Expand All @@ -372,15 +357,15 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
}

function _getTotalAccruedYield(uint128 currentIndex_) internal view returns (uint240 yield_) {
uint240 totalProjectedSupply_ = IndexingMath.getPresentAmountRoundedUp(
uint240 projectedEarningSupply_ = IndexingMath.getPresentAmountRoundedUp(
_principalOfTotalEarningSupply,
currentIndex_
);

uint240 totalEarningSupply_ = totalEarningSupply();
uint240 earningSupply_ = totalEarningSupply();

unchecked {
return totalProjectedSupply_ <= totalEarningSupply_ ? 0 : totalProjectedSupply_ - totalEarningSupply_;
return projectedEarningSupply_ <= earningSupply_ ? 0 : projectedEarningSupply_ - earningSupply_;
}
}

Expand Down
8 changes: 4 additions & 4 deletions src/interfaces/IWrappedMToken.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,11 +32,11 @@ interface IWrappedMToken is IMigratable, IERC20Extended {

/**
* @notice Emitted when there is insufficient balance to decrement from `account`.
* @param account The account with insufficient balance.
* @param rawBalance The raw balance of the account (can be present value or principal).
* @param amount The amount to decrement the `rawBalance` by (either present value or principal).
* @param account The account with insufficient balance.
* @param balance The balance of the account.
* @param amount The amount to decrement.
*/
error InsufficientBalance(address account, uint256 rawBalance, uint256 amount);
error InsufficientBalance(address account, uint256 balance, uint256 amount);

/// @notice Emitted when calling `startEarning` for an account not approved as earner by TTG.
error NotApprovedEarner();
Expand Down
Loading
Loading