diff --git a/src/WrappedMToken.sol b/src/WrappedMToken.sol index 87ff58b..4e500ad 100644 --- a/src/WrappedMToken.sol +++ b/src/WrappedMToken.sol @@ -48,6 +48,18 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { /* ============ Interactive Functions ============ */ + function wrap(address recipient_, uint256 amount_) external { + _mint(recipient_, UIntMath.safe240(amount_)); + + IMTokenLike(mToken).transferFrom(msg.sender, address(this), amount_); + } + + function unwrap(address recipient_, uint256 amount_) external { + _burn(msg.sender, UIntMath.safe240(amount_)); + + IMTokenLike(mToken).transfer(recipient_, amount_); + } + function claimFor(address account_) external returns (uint240 yield_) { return _claim(account_, currentIndex()); } @@ -58,22 +70,10 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { IMTokenLike(mToken).transfer(vault, yield_); } - function deposit(address recipient_, uint256 amount_) external { - _revertIfInsufficientAmount(amount_); - _revertIfInvalidRecipient(recipient_); - - emit Transfer(address(0), recipient_, amount_); - - // TODO: Technically, might want to `_revertIfInsufficientAmount` if earning principal is 0. - _addAmount(recipient_, UIntMath.safe240(amount_)); - - IMTokenLike(mToken).transferFrom(msg.sender, address(this), amount_); - } - 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; @@ -81,69 +81,45 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { uint128 currentIndex_ = currentIndex(); - _setBalanceInfo( - account_, - true, - currentIndex_, - IndexingMath.getPrincipalAmountRoundedDown(rawBalance_, currentIndex_) - ); + _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_); + + if (!isEarning_) return; - uint240 amount_ = IndexingMath.getPresentAmountRoundedDown(uint112(rawBalance_), index_); + emit StoppedEarning(account_); - _setBalanceInfo(account_, false, 0, amount_); + _setBalanceInfo(account_, false, 0, balance_); + _subtractTotalEarningSupply(balance_, currentIndex_); unchecked { - totalNonEarningSupply += amount_; + totalNonEarningSupply += balance_; } - - _subtractTotalEarningSupply(amount_, currentIndex_); - } - - function withdraw(address recipient_, uint256 amount_) external { - _revertIfInsufficientAmount(amount_); - - emit Transfer(msg.sender, address(0), amount_); - - // TODO: Technically, might want to `_revertIfInsufficientAmount` if earning principal is 0. - _subtractAmount(msg.sender, UIntMath.safe240(amount_)); - - IMTokenLike(mToken).transfer(recipient_, amount_); } /* ============ 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_) { @@ -151,14 +127,13 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { } 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; } } @@ -177,109 +152,128 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { /* ============ Internal Interactive Functions ============ */ - function _addAmount(address recipient_, uint240 amount_) internal { - (bool isEarning_, , ) = _getBalanceInfo(recipient_); + function _mint(address recipient_, uint240 amount_) internal { + _revertIfInsufficientAmount(amount_); + _revertIfInvalidRecipient(recipient_); + + emit Transfer(address(0), recipient_, amount_); + + (bool isEarning_, , , ) = _getBalanceInfo(recipient_); if (!isEarning_) return _addNonEarningAmount(recipient_, amount_); uint128 currentIndex_ = currentIndex(); _claim(recipient_, currentIndex_); + + // TODO: Technically, might want to `_revertIfInsufficientAmount` if earning principal is 0. _addEarningAmount(recipient_, amount_, currentIndex_); } - function _addNonEarningAmount(address recipient_, uint240 amount_) internal { - (, , uint240 rawBalance_) = _getBalanceInfo(recipient_); + function _burn(address account_, uint240 amount_) internal { + _revertIfInsufficientAmount(amount_); + + emit Transfer(msg.sender, address(0), amount_); + + (bool isEarning_, , , ) = _getBalanceInfo(account_); + + if (!isEarning_) return _subtractNonEarningAmount(account_, amount_); + + uint128 currentIndex_ = currentIndex(); + + _claim(account_, currentIndex_); + + // TODO: Technically, might want to `_revertIfInsufficientAmount` if earning principal is 0. + _subtractEarningAmount(account_, amount_, currentIndex_); + } + function _addNonEarningAmount(address recipient_, uint240 amount_) internal { unchecked { - _setBalanceInfo(recipient_, false, 0, rawBalance_ + amount_); + (, , , uint240 balance_) = _getBalanceInfo(recipient_); + _setBalanceInfo(recipient_, false, 0, balance_ + amount_); totalNonEarningSupply += amount_; } } - function _addEarningAmount(address recipient_, uint240 amount_, uint128 currentIndex_) internal { - (, , uint240 rawBalance_) = _getBalanceInfo(recipient_); + function _subtractNonEarningAmount(address account_, uint240 amount_) internal { + unchecked { + (, , , uint240 balance_) = _getBalanceInfo(account_); + + if (balance_ < amount_) revert InsufficientBalance(account_, balance_, amount_); + + _setBalanceInfo(account_, false, 0, balance_ - amount_); + totalNonEarningSupply -= amount_; + } + } + function _addEarningAmount(address recipient_, uint240 amount_, uint128 currentIndex_) internal { unchecked { - _setBalanceInfo( - recipient_, - true, - currentIndex_, - rawBalance_ + IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) - ); + (, , , uint240 balance_) = _getBalanceInfo(recipient_); + + _setBalanceInfo(recipient_, true, currentIndex_, balance_ + amount_); + _addTotalEarningSupply(amount_, currentIndex_); } + } + + function _subtractEarningAmount(address account_, uint240 amount_, uint128 currentIndex_) internal { + unchecked { + (, , , uint240 balance_) = _getBalanceInfo(account_); - _addTotalEarningSupply(amount_, currentIndex_); + if (balance_ < amount_) revert InsufficientBalance(account_, balance_, amount_); + + _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; - if (yield_ == 0) return 0; + _updateIndex(account_, currentIndex_); - emit Claimed(account_, yield_); + (, , , uint240 endingBalance_) = _getBalanceInfo(account_); unchecked { + yield_ = endingBalance_ - startingBalance_; + + if (yield_ == 0) return 0; + _setTotalEarningSupply(totalEarningSupply() + yield_, _principalOfTotalEarningSupply); } address claimOverrideRecipient_ = _getClaimOverrideRecipient(account_); if (claimOverrideRecipient_ == address(0)) { - // NOTE: The `Transfer` event for a non-zero `claimOverrideRecipient_` + emit Claimed(account_, account_, yield_); emit Transfer(address(0), account_, yield_); - return 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. - _transfer(account_, claimOverrideRecipient_, yield_, currentIndex_); - } - - 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(amount_)); - } - - function _subtractAmount(address account_, uint240 amount_) internal { - (bool isEarning_, , ) = _getBalanceInfo(account_); - - if (!isEarning_) return _subtractNonEarningAmount(account_, amount_); - - uint128 currentIndex_ = currentIndex(); - - _claim(account_, currentIndex_); - _subtractEarningAmount(account_, amount_, currentIndex_); - } - - function _subtractNonEarningAmount(address account_, uint240 amount_) internal { - (, , uint240 rawBalance_) = _getBalanceInfo(account_); - - if (rawBalance_ < amount_) revert InsufficientBalance(account_, rawBalance_, amount_); + } else { + emit Claimed(account_, claimOverrideRecipient_, yield_); - unchecked { - _setBalanceInfo(account_, false, 0, rawBalance_ - amount_); - totalNonEarningSupply -= amount_; + // NOTE: Watch out for a long chain of earning claim override recipients. + _transfer(account_, claimOverrideRecipient_, yield_, currentIndex_); } } - function _subtractEarningAmount(address account_, uint240 amount_, uint128 currentIndex_) internal { - (, , uint240 rawBalance_) = _getBalanceInfo(account_); + function _updateIndex(address account_, uint128 index_) internal { + uint256 unwrapped_ = BalanceInfo.unwrap(_balances[account_]); - uint112 principalAmount_ = IndexingMath.getPrincipalAmountRoundedUp(amount_, currentIndex_); + unwrapped_ &= ~(uint256(type(uint112).max) << 128); - if (rawBalance_ < principalAmount_) revert InsufficientBalance(account_, rawBalance_, principalAmount_); + _balances[account_] = BalanceInfo.wrap(unwrapped_ | (uint256(index_) << 112)); + } - unchecked { - _setBalanceInfo(account_, true, currentIndex_, rawBalance_ - principalAmount_); - _subtractTotalEarningSupply(amount_, currentIndex_); - } + function _setBalanceInfo(address account_, bool isEarning_, uint128 index_, uint240 amount_) internal { + _balances[account_] = isEarning_ + ? BalanceInfo.wrap( + (uint256(1) << 248) | + (uint256(index_) << 112) | + uint256(IndexingMath.getPrincipalAmountRoundedDown(amount_, index_)) + ) + : BalanceInfo.wrap(uint256(amount_)); } function _transfer(address sender_, address recipient_, uint240 amount_, uint128 currentIndex_) internal { @@ -290,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_) @@ -308,55 +302,38 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { function _addTotalEarningSupply(uint240 amount_, uint128 currentIndex_) internal { unchecked { - _setTotalEarningSupply( - totalEarningSupply() + amount_, - _principalOfTotalEarningSupply + IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) - ); + uint112 principal_ = IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_); + _setTotalEarningSupply(totalEarningSupply() + amount_, _principalOfTotalEarningSupply + principal_); } } function _subtractTotalEarningSupply(uint240 amount_, uint128 currentIndex_) internal { - // TODO: Consider `getPrincipalAmountRoundedUp` . unchecked { - _setTotalEarningSupply( - totalEarningSupply() - amount_, - _principalOfTotalEarningSupply - IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) - ); + // 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) { @@ -380,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_; } } diff --git a/src/interfaces/IWrappedMToken.sol b/src/interfaces/IWrappedMToken.sol index 0387af2..2be2631 100644 --- a/src/interfaces/IWrappedMToken.sol +++ b/src/interfaces/IWrappedMToken.sol @@ -9,7 +9,7 @@ import { IMigratable } from "./IMigratable.sol"; interface IWrappedMToken is IMigratable, IERC20Extended { /* ============ Events ============ */ - event Claimed(address indexed account, uint256 yield); + event Claimed(address indexed account, address indexed recipient, uint256 yield); event ExcessClaimed(uint256 yield); @@ -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(); @@ -46,18 +46,18 @@ interface IWrappedMToken is IMigratable, IERC20Extended { /* ============ Interactive Functions ============ */ + function wrap(address recipient, uint256 amount) external; + + function unwrap(address recipient, uint256 amount) external; + function claimFor(address account) external returns (uint240 yield); function claimExcess() external returns (uint240 yield); - function deposit(address recipient, uint256 amount) external; - function startEarningFor(address account) external; function stopEarningFor(address account) external; - function withdraw(address recipient, uint256 amount) external; - /* ============ View/Pure Functions ============ */ function accruedYieldOf(address account) external view returns (uint240 yield); diff --git a/test/Test.t.sol b/test/Test.t.sol index 49a8eb3..166ecf7 100644 --- a/test/Test.t.sol +++ b/test/Test.t.sol @@ -77,7 +77,7 @@ contract Tests is Test { _wrappedMToken.startEarningFor(_bob); vm.prank(_alice); - _wrappedMToken.deposit(_alice, 100_000000); + _wrappedMToken.wrap(_alice, 100_000000); _mToken.setBalanceOf(address(_wrappedMToken), 100_000000); @@ -93,7 +93,7 @@ contract Tests is Test { assertEq(_wrappedMToken.excess(), 0); vm.prank(_carol); - _wrappedMToken.deposit(_carol, 100_000000); + _wrappedMToken.wrap(_carol, 100_000000); _mToken.setBalanceOf(address(_wrappedMToken), 200_000000); @@ -127,7 +127,7 @@ contract Tests is Test { assertEq(_wrappedMToken.excess(), 100_000000); vm.prank(_bob); - _wrappedMToken.deposit(_bob, 100_000000); + _wrappedMToken.wrap(_bob, 100_000000); _mToken.setBalanceOf(address(_wrappedMToken), 500_000000); @@ -144,7 +144,7 @@ contract Tests is Test { assertEq(_wrappedMToken.excess(), 100_000000); vm.prank(_dave); - _wrappedMToken.deposit(_dave, 100_000000); + _wrappedMToken.wrap(_dave, 100_000000); _mToken.setBalanceOf(address(_wrappedMToken), 600_000000); @@ -321,7 +321,7 @@ contract Tests is Test { assertEq(_wrappedMToken.excess(), 600_000001); vm.prank(_alice); - _wrappedMToken.withdraw(_alice, 266_666664); + _wrappedMToken.unwrap(_alice, 266_666664); _mToken.setBalanceOf(address(_wrappedMToken), 1_233_333336); @@ -337,7 +337,7 @@ contract Tests is Test { assertEq(_wrappedMToken.excess(), 600_000001); vm.prank(_bob); - _wrappedMToken.withdraw(_bob, 333_333330); + _wrappedMToken.unwrap(_bob, 333_333330); _mToken.setBalanceOf(address(_wrappedMToken), 900_000006); @@ -353,7 +353,7 @@ contract Tests is Test { assertEq(_wrappedMToken.excess(), 600_000001); vm.prank(_carol); - _wrappedMToken.withdraw(_carol, 250_000000); + _wrappedMToken.unwrap(_carol, 250_000000); _mToken.setBalanceOf(address(_wrappedMToken), 650_000006); @@ -369,7 +369,7 @@ contract Tests is Test { assertEq(_wrappedMToken.excess(), 600_000001); vm.prank(_dave); - _wrappedMToken.withdraw(_dave, 50_000000); + _wrappedMToken.unwrap(_dave, 50_000000); _mToken.setBalanceOf(address(_wrappedMToken), 600_000006); diff --git a/test/WrappedMToken.t.sol b/test/WrappedMToken.t.sol index 7d91270..56eec8f 100644 --- a/test/WrappedMToken.t.sol +++ b/test/WrappedMToken.t.sol @@ -7,7 +7,6 @@ import { IERC20Extended } from "../lib/common/src/interfaces/IERC20Extended.sol" import { UIntMath } from "../lib/common/src/libs/UIntMath.sol"; import { IWrappedMToken } from "../src/interfaces/IWrappedMToken.sol"; -import { IRegistrarLike } from "../src/interfaces/IRegistrarLike.sol"; import { IndexingMath } from "../src/libs/IndexingMath.sol"; @@ -35,7 +34,7 @@ contract WrappedMTokenTests is Test { address internal _vault = makeAddr("vault"); - uint128 internal _expectedCurrentIndex; + uint128 internal _currentIndex; MockM internal _mToken; MockRegistrar internal _registrar; @@ -54,7 +53,7 @@ contract WrappedMTokenTests is Test { _wrappedMToken = WrappedMTokenHarness(address(new Proxy(address(_implementation)))); - _mToken.setCurrentIndex(_expectedCurrentIndex = 1_100000068703); + _mToken.setCurrentIndex(_currentIndex = 1_100000068703); } /* ============ constructor ============ */ @@ -71,30 +70,30 @@ contract WrappedMTokenTests is Test { new WrappedMTokenHarness(address(0)); } - /* ============ deposit ============ */ + /* ============ wrap ============ */ function test_deposit_insufficientAmount() external { vm.expectRevert(abi.encodeWithSelector(IERC20Extended.InsufficientAmount.selector, 0)); - _wrappedMToken.deposit(_alice, 0); + _wrappedMToken.wrap(_alice, 0); } function test_deposit_invalidRecipient() external { vm.expectRevert(abi.encodeWithSelector(IERC20Extended.InvalidRecipient.selector, address(0))); - _wrappedMToken.deposit(address(0), 1_000); + _wrappedMToken.wrap(address(0), 1_000); } function test_deposit_invalidAmount() external { vm.expectRevert(UIntMath.InvalidUInt240.selector); vm.prank(_alice); - _wrappedMToken.deposit(_alice, uint256(type(uint240).max) + 1); + _wrappedMToken.wrap(_alice, uint256(type(uint240).max) + 1); } function test_deposit_toNonEarner() external { vm.prank(_alice); - _wrappedMToken.deposit(_alice, 1_000); + _wrappedMToken.wrap(_alice, 1_000); assertEq(_wrappedMToken.internalBalanceOf(_alice), 1_000); assertEq(_wrappedMToken.totalNonEarningSupply(), 1_000); @@ -103,66 +102,71 @@ contract WrappedMTokenTests is Test { } function test_deposit_toEarner() external { - _wrappedMToken.setIsEarningOf(_alice, true); + _wrappedMToken.setAccountOf(_alice, true, _EXP_SCALED_ONE, 0); vm.prank(_alice); - _wrappedMToken.deposit(_alice, 999); + _wrappedMToken.wrap(_alice, 999); - assertEq(_wrappedMToken.internalBalanceOf(_alice), 908); + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 908); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_alice), 998); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 908); assertEq(_wrappedMToken.totalEarningSupply(), 999); vm.prank(_alice); - _wrappedMToken.deposit(_alice, 1); + _wrappedMToken.wrap(_alice, 1); - // No change due to principal round down on deposit. - assertEq(_wrappedMToken.internalBalanceOf(_alice), 908); + // No change due to principal round down on wrap. + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 908); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_alice), 998); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 908); assertEq(_wrappedMToken.totalEarningSupply(), 1000); vm.prank(_alice); - _wrappedMToken.deposit(_alice, 2); + _wrappedMToken.wrap(_alice, 2); - assertEq(_wrappedMToken.internalBalanceOf(_alice), 909); + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 909); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_alice), 999); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 909); assertEq(_wrappedMToken.totalEarningSupply(), 1002); } - /* ============ withdraw ============ */ + /* ============ unwrap ============ */ function test_withdraw_insufficientAmount() external { vm.expectRevert(abi.encodeWithSelector(IERC20Extended.InsufficientAmount.selector, 0)); - _wrappedMToken.withdraw(_alice, 0); + _wrappedMToken.unwrap(_alice, 0); } function test_withdraw_insufficientBalance_fromNonEarner() external { - _wrappedMToken.setRawBalanceOf(_alice, 999); + _wrappedMToken.setBalanceOf(_alice, 999); vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); vm.prank(_alice); - _wrappedMToken.withdraw(_alice, 1_000); + _wrappedMToken.unwrap(_alice, 1_000); } function test_withdraw_insufficientBalance_fromEarner() external { - _wrappedMToken.setIsEarningOf(_alice, true); - _wrappedMToken.setRawBalanceOf(_alice, 908); + _wrappedMToken.setAccountOf(_alice, true, _currentIndex, 1_000); - vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 908, 910)); + vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); vm.prank(_alice); - _wrappedMToken.withdraw(_alice, 1_000); + _wrappedMToken.unwrap(_alice, 1_000); } function test_withdraw_fromNonEarner() external { _wrappedMToken.setTotalNonEarningSupply(1_000); - _wrappedMToken.setRawBalanceOf(_alice, 1_000); + _wrappedMToken.setBalanceOf(_alice, 1_000); vm.prank(_alice); - _wrappedMToken.withdraw(_alice, 500); + _wrappedMToken.unwrap(_alice, 500); assertEq(_wrappedMToken.internalBalanceOf(_alice), 500); assertEq(_wrappedMToken.totalNonEarningSupply(), 500); @@ -170,7 +174,7 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.indexOfTotalEarningSupply(), 0); vm.prank(_alice); - _wrappedMToken.withdraw(_alice, 500); + _wrappedMToken.unwrap(_alice, 500); assertEq(_wrappedMToken.internalBalanceOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); @@ -180,31 +184,35 @@ contract WrappedMTokenTests is Test { function test_withdraw_fromEarner() external { _wrappedMToken.setPrincipalOfTotalEarningSupply(909); - _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setIndexOfTotalEarningSupply(_currentIndex); + + _wrappedMToken.setAccountOf(_alice, true, _currentIndex, 1_000); - _wrappedMToken.setIsEarningOf(_alice, true); - _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); - _wrappedMToken.setRawBalanceOf(_alice, 909); + assertEq(_wrappedMToken.balanceOf(_alice), 999); vm.prank(_alice); - _wrappedMToken.withdraw(_alice, 1); + _wrappedMToken.unwrap(_alice, 1); - // Change due to principal round up on withdraw. - assertEq(_wrappedMToken.internalBalanceOf(_alice), 908); + // Change due to principal round up on unwrap. + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 907); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_alice), 997); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 999); vm.prank(_alice); - _wrappedMToken.withdraw(_alice, 998); + _wrappedMToken.unwrap(_alice, 997); + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 0); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); assertEq(_wrappedMToken.internalBalanceOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.totalEarningSupply(), 1); + assertEq(_wrappedMToken.totalEarningSupply(), 2); // TODO: Fix? } /* ============ transfer ============ */ function test_transfer_invalidRecipient() external { - _wrappedMToken.setRawBalanceOf(_alice, 1_000); + _wrappedMToken.setBalanceOf(_alice, 1_000); vm.expectRevert(abi.encodeWithSelector(IERC20Extended.InvalidRecipient.selector, address(0))); @@ -213,7 +221,7 @@ contract WrappedMTokenTests is Test { } function test_transfer_insufficientBalance_fromNonEarner_toNonEarner() external { - _wrappedMToken.setRawBalanceOf(_alice, 999); + _wrappedMToken.setBalanceOf(_alice, 999); vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); vm.prank(_alice); @@ -221,10 +229,9 @@ contract WrappedMTokenTests is Test { } function test_transfer_insufficientBalance_fromEarner_toNonEarner() external { - _wrappedMToken.setIsEarningOf(_alice, true); - _wrappedMToken.setRawBalanceOf(_alice, 908); + _wrappedMToken.setAccountOf(_alice, true, _currentIndex, 1_000); - vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 908, 910)); + vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); vm.prank(_alice); _wrappedMToken.transfer(_bob, 1_000); } @@ -232,8 +239,8 @@ contract WrappedMTokenTests is Test { function test_transfer_fromNonEarner_toNonEarner() external { _wrappedMToken.setTotalNonEarningSupply(1_500); - _wrappedMToken.setRawBalanceOf(_alice, 1_000); - _wrappedMToken.setRawBalanceOf(_bob, 500); + _wrappedMToken.setBalanceOf(_alice, 1_000); + _wrappedMToken.setBalanceOf(_bob, 500); vm.prank(_alice); _wrappedMToken.transfer(_bob, 500); @@ -259,8 +266,8 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setTotalNonEarningSupply(supply_); - _wrappedMToken.setRawBalanceOf(_alice, aliceBalance_); - _wrappedMToken.setRawBalanceOf(_bob, bobBalance); + _wrappedMToken.setBalanceOf(_alice, aliceBalance_); + _wrappedMToken.setBalanceOf(_bob, bobBalance); vm.prank(_alice); _wrappedMToken.transfer(_bob, transferAmount_); @@ -275,19 +282,19 @@ contract WrappedMTokenTests is Test { function test_transfer_fromEarner_toNonEarner() external { _wrappedMToken.setPrincipalOfTotalEarningSupply(909); - _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setIndexOfTotalEarningSupply(_currentIndex); _wrappedMToken.setTotalNonEarningSupply(500); - _wrappedMToken.setIsEarningOf(_alice, true); - _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); - _wrappedMToken.setRawBalanceOf(_alice, 909); + _wrappedMToken.setAccountOf(_alice, true, _currentIndex, 1_000); - _wrappedMToken.setRawBalanceOf(_bob, 500); + _wrappedMToken.setBalanceOf(_bob, 500); vm.prank(_alice); _wrappedMToken.transfer(_bob, 500); - assertEq(_wrappedMToken.internalBalanceOf(_alice), 454); + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 453); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_alice), 498); assertEq(_wrappedMToken.internalBalanceOf(_bob), 1_000); @@ -298,7 +305,9 @@ contract WrappedMTokenTests is Test { _wrappedMToken.transfer(_bob, 1); // Change due to principal round up on burn. - assertEq(_wrappedMToken.internalBalanceOf(_alice), 453); + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 451); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_alice), 496); assertEq(_wrappedMToken.internalBalanceOf(_bob), 1_001); @@ -308,21 +317,21 @@ contract WrappedMTokenTests is Test { function test_transfer_fromNonEarner_toEarner() external { _wrappedMToken.setPrincipalOfTotalEarningSupply(455); - _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setIndexOfTotalEarningSupply(_currentIndex); _wrappedMToken.setTotalNonEarningSupply(1_000); - _wrappedMToken.setRawBalanceOf(_alice, 1_000); + _wrappedMToken.setBalanceOf(_alice, 1_000); - _wrappedMToken.setIsEarningOf(_bob, true); - _wrappedMToken.setIndexOf(_bob, _expectedCurrentIndex); - _wrappedMToken.setRawBalanceOf(_bob, 455); + _wrappedMToken.setAccountOf(_bob, true, _currentIndex, 500); vm.prank(_alice); _wrappedMToken.transfer(_bob, 500); assertEq(_wrappedMToken.internalBalanceOf(_alice), 500); - assertEq(_wrappedMToken.internalBalanceOf(_bob), 909); + assertEq(_wrappedMToken.internalPrincipalOf(_bob), 908); + assertEq(_wrappedMToken.internalIndexOf(_bob), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_bob), 998); assertEq(_wrappedMToken.totalNonEarningSupply(), 500); assertEq(_wrappedMToken.totalEarningSupply(), 1_001); @@ -330,22 +339,22 @@ contract WrappedMTokenTests is Test { function test_transfer_fromEarner_toEarner() external { _wrappedMToken.setPrincipalOfTotalEarningSupply(1_364); - _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setIndexOfTotalEarningSupply(_currentIndex); - _wrappedMToken.setIsEarningOf(_alice, true); - _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); - _wrappedMToken.setRawBalanceOf(_alice, 909); + _wrappedMToken.setAccountOf(_alice, true, _currentIndex, 1_000); - _wrappedMToken.setIsEarningOf(_bob, true); - _wrappedMToken.setIndexOf(_bob, _expectedCurrentIndex); - _wrappedMToken.setRawBalanceOf(_bob, 454); + _wrappedMToken.setAccountOf(_bob, true, _currentIndex, 500); vm.prank(_alice); _wrappedMToken.transfer(_bob, 500); - assertEq(_wrappedMToken.internalBalanceOf(_alice), 454); + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 453); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_alice), 498); - assertEq(_wrappedMToken.internalBalanceOf(_bob), 908); + assertEq(_wrappedMToken.internalPrincipalOf(_bob), 908); + assertEq(_wrappedMToken.internalIndexOf(_bob), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_bob), 998); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 1_501); @@ -360,7 +369,7 @@ contract WrappedMTokenTests is Test { function test_startEarningFor() external { _wrappedMToken.setTotalNonEarningSupply(1_000); - _wrappedMToken.setRawBalanceOf(_alice, 1_000); + _wrappedMToken.setBalanceOf(_alice, 1_000); _registrar.setListContains(_EARNERS_LIST, _alice, true); @@ -369,8 +378,10 @@ contract WrappedMTokenTests is Test { _wrappedMToken.startEarningFor(_alice); - assertEq(_wrappedMToken.internalBalanceOf(_alice), 909); assertEq(_wrappedMToken.isEarning(_alice), true); + assertEq(_wrappedMToken.internalPrincipalOf(_alice), 909); + assertEq(_wrappedMToken.internalIndexOf(_alice), _currentIndex); + assertEq(_wrappedMToken.internalBalanceOf(_alice), 999); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 1_000); @@ -379,10 +390,11 @@ contract WrappedMTokenTests is Test { function test_startEarning_overflow() external { uint256 aliceBalance_ = uint256(type(uint112).max) + 20; - _mToken.setCurrentIndex(_expectedCurrentIndex = _EXP_SCALED_ONE); + _mToken.setCurrentIndex(_currentIndex = _EXP_SCALED_ONE); _wrappedMToken.setTotalNonEarningSupply(aliceBalance_); - _wrappedMToken.setRawBalanceOf(_alice, aliceBalance_); + + _wrappedMToken.setBalanceOf(_alice, aliceBalance_); _registrar.setListContains(_EARNERS_LIST, _alice, true); @@ -400,11 +412,9 @@ contract WrappedMTokenTests is Test { function test_stopEarningFor() external { _wrappedMToken.setPrincipalOfTotalEarningSupply(909); - _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setIndexOfTotalEarningSupply(_currentIndex); - _wrappedMToken.setIsEarningOf(_alice, true); - _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); - _wrappedMToken.setRawBalanceOf(_alice, 909); + _wrappedMToken.setAccountOf(_alice, true, _currentIndex, 1_000); _registrar.setListContains(_EARNERS_LIST, _alice, false); @@ -422,29 +432,27 @@ contract WrappedMTokenTests is Test { /* ============ balanceOf ============ */ function test_balanceOf_nonEarner() external { - _wrappedMToken.setRawBalanceOf(_alice, 500); + _wrappedMToken.setBalanceOf(_alice, 500); assertEq(_wrappedMToken.balanceOf(_alice), 500); - _wrappedMToken.setRawBalanceOf(_alice, 1_000); + _wrappedMToken.setBalanceOf(_alice, 1_000); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); } function test_balanceOf_earner() external { - _wrappedMToken.setIsEarningOf(_alice, true); - _wrappedMToken.setIndexOf(_alice, _EXP_SCALED_ONE); - _wrappedMToken.setRawBalanceOf(_alice, 454); + _wrappedMToken.setAccountOf(_alice, true, _EXP_SCALED_ONE, 500); - assertEq(_wrappedMToken.balanceOf(_alice), 454); + assertEq(_wrappedMToken.balanceOf(_alice), 500); - _wrappedMToken.setRawBalanceOf(_alice, 909); + _wrappedMToken.setBalanceOf(_alice, 1_000); - assertEq(_wrappedMToken.balanceOf(_alice), 909); + assertEq(_wrappedMToken.balanceOf(_alice), 1_000); - _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); + _wrappedMToken.setIndexOf(_alice, 2 * _EXP_SCALED_ONE); - assertEq(_wrappedMToken.balanceOf(_alice), 999); + assertEq(_wrappedMToken.balanceOf(_alice), 1_000); } /* ============ totalNonEarningSupply ============ */ @@ -461,7 +469,7 @@ contract WrappedMTokenTests is Test { function test_totalEarningSupply() external { // TODO: more variations _wrappedMToken.setPrincipalOfTotalEarningSupply(909); - _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setIndexOfTotalEarningSupply(_currentIndex); assertEq(_wrappedMToken.totalEarningSupply(), 1_000); } @@ -480,7 +488,7 @@ contract WrappedMTokenTests is Test { function test_totalSupply_onlyTotalEarningSupply() external { // TODO: more variations _wrappedMToken.setPrincipalOfTotalEarningSupply(909); - _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setIndexOfTotalEarningSupply(_currentIndex); assertEq(_wrappedMToken.totalSupply(), 1_000); } @@ -488,7 +496,7 @@ contract WrappedMTokenTests is Test { function test_totalSupply() external { // TODO: more variations _wrappedMToken.setPrincipalOfTotalEarningSupply(909); - _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setIndexOfTotalEarningSupply(_currentIndex); _wrappedMToken.setTotalNonEarningSupply(500); diff --git a/test/utils/WrappedMTokenHarness.sol b/test/utils/WrappedMTokenHarness.sol index a31c9b0..a77d861 100644 --- a/test/utils/WrappedMTokenHarness.sol +++ b/test/utils/WrappedMTokenHarness.sol @@ -8,18 +8,22 @@ contract WrappedMTokenHarness is WrappedMToken { constructor(address mToken_) WrappedMToken(mToken_) {} function setIsEarningOf(address account_, bool isEarning_) external { - (, uint128 index_, uint240 rawBalance_) = _getBalanceInfo(account_); - _setBalanceInfo(account_, isEarning_, index_, rawBalance_); + (, uint128 index_, , uint240 balance_) = _getBalanceInfo(account_); + _setBalanceInfo(account_, isEarning_, index_, balance_); } function setIndexOf(address account_, uint256 index_) external { - (bool isEarning_, , uint240 rawBalance_) = _getBalanceInfo(account_); - _setBalanceInfo(account_, isEarning_, uint128(index_), rawBalance_); + (bool isEarning_, , , uint240 balance_) = _getBalanceInfo(account_); + _setBalanceInfo(account_, isEarning_, uint128(index_), balance_); } - function setRawBalanceOf(address account_, uint256 rawBalance_) external { - (bool isEarning_, uint128 index_, ) = _getBalanceInfo(account_); - _setBalanceInfo(account_, isEarning_, index_, uint240(rawBalance_)); + function setBalanceOf(address account_, uint256 balance_) external { + (bool isEarning_, uint128 index_, , ) = _getBalanceInfo(account_); + _setBalanceInfo(account_, isEarning_, index_, uint240(balance_)); + } + + function setAccountOf(address account_, bool isEarning_, uint256 index_, uint256 balance_) external { + _setBalanceInfo(account_, isEarning_, uint128(index_), uint240(balance_)); } function setTotalNonEarningSupply(uint256 totalNonEarningSupply_) external { @@ -34,15 +38,23 @@ contract WrappedMTokenHarness is WrappedMToken { _indexOfTotalEarningSupply = uint128(indexOfTotalEarningSupply_); } - function internalBalanceOf(address account_) external view returns (uint256 balance_) { - (, , balance_) = _getBalanceInfo(account_); + function internalBalanceOf(address account_) external view returns (uint240 balance_) { + (, , , balance_) = _getBalanceInfo(account_); + } + + function internalIndexOf(address account_) external view returns (uint128 index_) { + (, index_, , ) = _getBalanceInfo(account_); + } + + function internalPrincipalOf(address account_) external view returns (uint112 principal_) { + (, , principal_, ) = _getBalanceInfo(account_); } - function principalOfTotalEarningSupply() external view returns (uint256 principalOfTotalEarningSupply_) { + function principalOfTotalEarningSupply() external view returns (uint240 principalOfTotalEarningSupply_) { principalOfTotalEarningSupply_ = _principalOfTotalEarningSupply; } - function indexOfTotalEarningSupply() external view returns (uint256 indexOfTotalEarningSupply_) { + function indexOfTotalEarningSupply() external view returns (uint128 indexOfTotalEarningSupply_) { indexOfTotalEarningSupply_ = _indexOfTotalEarningSupply; } }