From 76f82b0fab3e946ee48e51a015d9081cf23ed1e6 Mon Sep 17 00:00:00 2001 From: Michael De Luca Date: Fri, 10 Jan 2025 14:58:17 -0500 Subject: [PATCH] feat: principal-based account - principal-based earner account, instead of last index - auto and manual migration for existing earners --- src/WrappedMToken.sol | 230 ++++++++++++++++---------- src/interfaces/IWrappedMToken.sol | 27 +++- test/integration/MorphoBlue.t.sol | 8 +- test/integration/Protocol.t.sol | 42 ++--- test/integration/UniswapV3.t.sol | 4 +- test/unit/Stories.t.sol | 32 ++-- test/unit/WrappedMToken.t.sol | 242 ++++++++++++++++------------ test/utils/Invariants.sol | 15 +- test/utils/WrappedMTokenHarness.sol | 23 ++- 9 files changed, 363 insertions(+), 260 deletions(-) diff --git a/src/WrappedMToken.sol b/src/WrappedMToken.sol index 1f6c562..7a01507 100644 --- a/src/WrappedMToken.sol +++ b/src/WrappedMToken.sol @@ -33,12 +33,12 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { /* ============ Structs ============ */ /** - * @dev Struct to represent an account's balance and yield earning details + * @dev Struct to represent an account's balance and yield earning details with last index (prior version). * @param isEarning Whether the account is actively earning yield. * @param balance The present amount of tokens held by the account. * @param lastIndex The index of the last interaction for the account (0 for non-earning accounts). */ - struct Account { + struct IndexBasedAccount { // First Slot bool isEarning; uint240 balance; @@ -46,6 +46,26 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { uint128 lastIndex; } + enum EarningState { + NOT_EARNING, + INDEX_BASED, + PRINCIPAL_BASED + } + + /** + * @dev Struct to represent an account's balance and yield earning details. + * @param earningState How the account is actively earning yield. + * @param balance The present amount of tokens held by the account. + * @param earningPrincipal The earning principal for the account (0 for non-earning accounts). + */ + struct Account { + // First Slot + EarningState earningState; + uint240 balance; + // Second slot + uint112 earningPrincipal; + } + /* ============ Variables ============ */ /// @inheritdoc IWrappedMToken @@ -73,7 +93,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { address public immutable excessDestination; /// @inheritdoc IWrappedMToken - uint112 public principalOfTotalEarningSupply; + uint112 public totalEarningPrincipal; /// @inheritdoc IWrappedMToken uint240 public totalEarningSupply; @@ -131,6 +151,8 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { /// @inheritdoc IWrappedMToken function unwrap(address recipient_) external returns (uint240 unwrapped_) { + _migrateEarner(msg.sender); // NOTE: Need to migrate before calling `balanceWithYieldOf`. + return _unwrap(msg.sender, recipient_, uint240(balanceWithYieldOf(msg.sender))); } @@ -198,6 +220,18 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { } } + /// @inheritdoc IWrappedMToken + function migrateAccount(address account_) external { + _migrateEarner(account_); + } + + /// @inheritdoc IWrappedMToken + function migrateAccounts(address[] calldata accounts_) external { + for (uint256 index_; index_ < accounts_.length; ++index_) { + _migrateEarner(accounts_[index_]); + } + } + /* ============ Temporary Admin Migration ============ */ /// @inheritdoc IWrappedMToken @@ -213,8 +247,13 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { function accruedYieldOf(address account_) public view returns (uint240 yield_) { Account storage accountInfo_ = _accounts[account_]; + // TODO: Add function to compute accrued yield for an account given a last index. + if (accountInfo_.earningState == EarningState.INDEX_BASED) revert AccountNotMigrated(); + return - accountInfo_.isEarning ? _getAccruedYield(accountInfo_.balance, accountInfo_.lastIndex, currentIndex()) : 0; + accountInfo_.earningState == EarningState.PRINCIPAL_BASED + ? _getAccruedYield(accountInfo_.balance, accountInfo_.earningPrincipal, currentIndex()) + : 0; } /// @inheritdoc IERC20 @@ -230,8 +269,8 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { } /// @inheritdoc IWrappedMToken - function lastIndexOf(address account_) external view returns (uint128 lastIndex_) { - return _accounts[account_].lastIndex; + function earningPrincipalOf(address account_) external view returns (uint112 earningPrincipal_) { + return _accounts[account_].earningPrincipal; } /// @inheritdoc IWrappedMToken @@ -253,7 +292,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { /// @inheritdoc IWrappedMToken function isEarning(address account_) external view returns (bool isEarning_) { - return _accounts[account_].isEarning; + return _accounts[account_].earningState != EarningState.NOT_EARNING; } /// @inheritdoc IWrappedMToken @@ -301,15 +340,15 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { _revertIfInsufficientAmount(amount_); _revertIfZeroAccount(recipient_); - if (_accounts[recipient_].isEarning) { + if (_accounts[recipient_].earningState == EarningState.NOT_EARNING) { + _addNonEarningAmount(recipient_, amount_); + } else { uint128 currentIndex_ = currentIndex(); _claim(recipient_, currentIndex_); // NOTE: Additional principal may end up being rounded to 0 and this will not `_revertIfInsufficientAmount`. _addEarningAmount(recipient_, amount_, currentIndex_); - } else { - _addNonEarningAmount(recipient_, amount_); } emit Transfer(address(0), recipient_, amount_); @@ -323,15 +362,15 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { function _burn(address account_, uint240 amount_) internal { _revertIfInsufficientAmount(amount_); - if (_accounts[account_].isEarning) { + if (_accounts[account_].earningState == EarningState.NOT_EARNING) { + _subtractNonEarningAmount(account_, amount_); + } else { uint128 currentIndex_ = currentIndex(); _claim(account_, currentIndex_); // NOTE: Subtracted principal may end up being rounded to 0 and this will not `_revertIfInsufficientAmount`. _subtractEarningAmount(account_, amount_, currentIndex_); - } else { - _subtractNonEarningAmount(account_, amount_); } emit Transfer(account_, address(0), amount_); @@ -369,21 +408,26 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { } /** - * @dev Increments the token balance of `account_` by `amount_`, assuming earning status and updated index. + * @dev Increments the token balance of `account_` by `amount_`, assuming earning status. * @param account_ The address whose account balance will be incremented. * @param amount_ The present amount of tokens to increment by. * @param currentIndex_ The current index to use to compute the principal amount. */ function _addEarningAmount(address account_, uint240 amount_, uint128 currentIndex_) internal { + Account storage accountInfo_ = _accounts[account_]; + uint112 principal_ = IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_); + // NOTE: Can be `unchecked` because the max amount of wrappable M is never greater than `type(uint240).max`. unchecked { - _accounts[account_].balance += amount_; - _addTotalEarningSupply(amount_, currentIndex_); + accountInfo_.balance += amount_; + accountInfo_.earningPrincipal = UIntMath.safe112(uint256(accountInfo_.earningPrincipal) + principal_); } + + _addTotalEarningSupply(amount_, principal_); } /** - * @dev Decrements the token balance of `account_` by `amount_`, assuming earning status and updated index. + * @dev Decrements the token balance of `account_` by `amount_`, assuming earning status. * @param account_ The address whose account balance will be decremented. * @param amount_ The present amount of tokens to decrement by. * @param currentIndex_ The current index to use to compute the principal amount. @@ -395,10 +439,19 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { if (balance_ < amount_) revert InsufficientBalance(account_, balance_, amount_); + uint112 earningPrincipal_ = accountInfo_.earningPrincipal; + + uint112 principal_ = UIntMath.min112( + IndexingMath.getPrincipalAmountRoundedUp(amount_, currentIndex_), + earningPrincipal_ + ); + unchecked { accountInfo_.balance = balance_ - amount_; - _subtractTotalEarningSupply(amount_, currentIndex_); + accountInfo_.earningPrincipal = earningPrincipal_ - principal_; } + + _subtractTotalEarningSupply(amount_, principal_); } /** @@ -410,24 +463,20 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { function _claim(address account_, uint128 currentIndex_) internal returns (uint240 yield_) { Account storage accountInfo_ = _accounts[account_]; - if (!accountInfo_.isEarning) return 0; + if (accountInfo_.earningState == EarningState.NOT_EARNING) return 0; - uint128 index_ = accountInfo_.lastIndex; - - if (currentIndex_ == index_) return 0; + _migrateEarner(account_); uint240 startingBalance_ = accountInfo_.balance; - yield_ = _getAccruedYield(startingBalance_, index_, currentIndex_); - - accountInfo_.lastIndex = currentIndex_; + // NOTE": Account must be migrated before entering this section. + yield_ = _getAccruedYield(startingBalance_, accountInfo_.earningPrincipal, currentIndex_); if (yield_ == 0) return 0; unchecked { + // Update balance and total earning supply to account for the yield, but the principals have not changed. accountInfo_.balance = startingBalance_ + yield_; - - // Update the total earning supply to account for the yield, but the principal has not changed. totalEarningSupply += yield_; } @@ -455,8 +504,6 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { _revertIfZeroAccount(sender_); _revertIfZeroAccount(recipient_); - // Claims for both the sender and recipient are required before transferring since add an subtract functions - // assume accounts' balances are up-to-date with the current index. _claim(sender_, currentIndex_); _claim(recipient_, currentIndex_); @@ -464,31 +511,23 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { if (amount_ == 0) return; - Account storage senderAccountInfo_ = _accounts[sender_]; - Account storage recipientAccountInfo_ = _accounts[recipient_]; - - // If the sender and recipient are both earning or both non-earning, update their balances without affecting - // the total earning and non-earning supply storage variables. - if (senderAccountInfo_.isEarning == recipientAccountInfo_.isEarning) { - uint240 senderBalance_ = senderAccountInfo_.balance; + if (sender_ == recipient_) { + uint240 balance_ = _accounts[sender_].balance; - if (senderBalance_ < amount_) revert InsufficientBalance(sender_, senderBalance_, amount_); - - unchecked { - senderAccountInfo_.balance = senderBalance_ - amount_; - recipientAccountInfo_.balance += amount_; - } + if (balance_ < amount_) revert InsufficientBalance(sender_, balance_, amount_); return; } - senderAccountInfo_.isEarning - ? _subtractEarningAmount(sender_, amount_, currentIndex_) - : _subtractNonEarningAmount(sender_, amount_); + // TODO: Don't touch globals if both ae earning or not earning. + + _accounts[sender_].earningState == EarningState.NOT_EARNING + ? _subtractNonEarningAmount(sender_, amount_) + : _subtractEarningAmount(sender_, amount_, currentIndex_); - recipientAccountInfo_.isEarning - ? _addEarningAmount(recipient_, amount_, currentIndex_) - : _addNonEarningAmount(recipient_, amount_); + _accounts[recipient_].earningState == EarningState.NOT_EARNING + ? _addNonEarningAmount(recipient_, amount_) + : _addEarningAmount(recipient_, amount_, currentIndex_); } /** @@ -503,38 +542,29 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { /** * @dev Increments total earning supply by `amount_` tokens. - * @param amount_ The present amount of tokens to increment total earning supply by. - * @param currentIndex_ The current index used to compute the principal amount. + * @param amount_ The present amount of tokens to increment total earning supply by. + * @param principal_ The principal amount of tokens to increment total earning principal by. */ - function _addTotalEarningSupply(uint240 amount_, uint128 currentIndex_) internal { + function _addTotalEarningSupply(uint240 amount_, uint112 principal_) internal { unchecked { // Increment the total earning supply and principal proportionally. totalEarningSupply += amount_; - principalOfTotalEarningSupply += IndexingMath.getPrincipalAmountRoundedUp(amount_, currentIndex_); + totalEarningPrincipal = UIntMath.safe112(uint256(totalEarningPrincipal) + principal_); } } /** * @dev Decrements total earning supply by `amount_` tokens. - * @param amount_ The present amount of tokens to decrement total earning supply by. - * @param currentIndex_ The current index used to compute the principal amount. + * @param amount_ The present amount of tokens to decrement total earning supply by. + * @param principal_ The principal amount of tokens to decrement total earning principal by. */ - function _subtractTotalEarningSupply(uint240 amount_, uint128 currentIndex_) internal { - if (amount_ >= totalEarningSupply) { - delete totalEarningSupply; - delete principalOfTotalEarningSupply; - - return; - } - - uint112 principal_ = IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_); + function _subtractTotalEarningSupply(uint240 amount_, uint112 principal_) internal { + uint240 totalEarningSupply_ = totalEarningSupply; + uint112 totalEarningPrincipal_ = totalEarningPrincipal; unchecked { - principalOfTotalEarningSupply -= ( - principal_ > principalOfTotalEarningSupply ? principalOfTotalEarningSupply : principal_ - ); - - totalEarningSupply -= amount_; + totalEarningSupply = totalEarningSupply_ - UIntMath.min240(amount_, totalEarningSupply_); + totalEarningPrincipal = totalEarningPrincipal_ - UIntMath.min112(principal_, totalEarningPrincipal_); } } @@ -590,14 +620,15 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { Account storage accountInfo_ = _accounts[account_]; - if (accountInfo_.isEarning) return; - - accountInfo_.isEarning = true; - accountInfo_.lastIndex = currentIndex_; + if (accountInfo_.earningState != EarningState.NOT_EARNING) return; uint240 balance_ = accountInfo_.balance; + uint112 earningPrincipal_ = IndexingMath.getPrincipalAmountRoundedDown(balance_, currentIndex_); - _addTotalEarningSupply(balance_, currentIndex_); + accountInfo_.earningState = EarningState.PRINCIPAL_BASED; + accountInfo_.earningPrincipal = earningPrincipal_; + + _addTotalEarningSupply(balance_, earningPrincipal_); unchecked { totalNonEarningSupply -= balance_; @@ -618,14 +649,15 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { Account storage accountInfo_ = _accounts[account_]; - if (!accountInfo_.isEarning) return; - - delete accountInfo_.isEarning; - delete accountInfo_.lastIndex; + if (accountInfo_.earningState == EarningState.NOT_EARNING) return; uint240 balance_ = accountInfo_.balance; + uint112 earningPrincipal_ = accountInfo_.earningPrincipal; + + delete accountInfo_.earningState; + delete accountInfo_.earningPrincipal; - _subtractTotalEarningSupply(balance_, currentIndex_); + _subtractTotalEarningSupply(balance_, earningPrincipal_); unchecked { totalNonEarningSupply += balance_; @@ -634,6 +666,29 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { emit StoppedEarning(account_); } + /** + * @dev Migrates the account struct for `account` from v1 to v2. + * @param account_ The account to migrate. + */ + function _migrateEarner(address account_) internal { + Account storage accountInfo_ = _accounts[account_]; + + if (accountInfo_.earningState != EarningState.INDEX_BASED) return; + + IndexBasedAccount storage accountInfoV1_; + + assembly { + accountInfoV1_.slot := accountInfo_.slot + } + + uint128 lastIndex_ = accountInfoV1_.lastIndex; + + delete accountInfoV1_.lastIndex; + + accountInfo_.earningPrincipal = IndexingMath.getPrincipalAmountRoundedDown(accountInfoV1_.balance, lastIndex_); + accountInfo_.earningState = EarningState.PRINCIPAL_BASED; + } + /* ============ Internal View/Pure Functions ============ */ /// @dev Returns the current index of the M Token. @@ -642,21 +697,18 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { } /** - * @dev Compute the yield given an account's balance, last index, and the current index. - * @param balance_ The token balance of an earning account. - * @param lastIndex_ The index of ast interaction for the account. - * @param currentIndex_ The current index. - * @return yield_ The yield accrued since the last interaction. + * @dev Compute the yield given an account's balance, earning principal, and the current index. + * @param balance_ The token balance of an earning account. + * @param earningPrincipal_ The earning principal of the account. + * @param currentIndex_ The current index. + * @return yield_ The yield accrued since the last interaction. */ function _getAccruedYield( uint240 balance_, - uint128 lastIndex_, + uint112 earningPrincipal_, uint128 currentIndex_ ) internal pure returns (uint240 yield_) { - uint240 balanceWithYield_ = IndexingMath.getPresentAmountRoundedDown( - IndexingMath.getPrincipalAmountRoundedDown(balance_, lastIndex_), - currentIndex_ - ); + uint240 balanceWithYield_ = IndexingMath.getPresentAmountRoundedDown(earningPrincipal_, currentIndex_); unchecked { return (balanceWithYield_ <= balance_) ? 0 : balanceWithYield_ - balance_; @@ -727,7 +779,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { * @return supply_ The projected total earning supply. */ function _projectedEarningSupply(uint128 currentIndex_) internal view returns (uint240 supply_) { - return IndexingMath.getPresentAmountRoundedDown(principalOfTotalEarningSupply, currentIndex_); + return IndexingMath.getPresentAmountRoundedDown(totalEarningPrincipal, currentIndex_); } /** diff --git a/src/interfaces/IWrappedMToken.sol b/src/interfaces/IWrappedMToken.sol index 366448f..9a5f573 100644 --- a/src/interfaces/IWrappedMToken.sol +++ b/src/interfaces/IWrappedMToken.sol @@ -52,6 +52,9 @@ interface IWrappedMToken is IMigratable, IERC20Extended { /* ============ Custom Errors ============ */ + /// @notice Emitted when performing an operation that is not allowed on an un-migrated account. + error AccountNotMigrated(); + /// @notice Emitted when performing an operation that is not allowed when earning is disabled. error EarningIsDisabled(); @@ -171,6 +174,18 @@ interface IWrappedMToken is IMigratable, IERC20Extended { */ function stopEarningFor(address[] calldata accounts) external; + /** + * @notice Migrates the account struct for `account` from v1 to v2. + * @param account The account to migrate. + */ + function migrateAccount(address account) external; + + /** + * @notice Migrates the account structs for `accounts` from v1 to v2. + * @param accounts The accounts to migrate. + */ + function migrateAccounts(address[] calldata accounts) external; + /* ============ Temporary Admin Migration ============ */ /** @@ -208,11 +223,11 @@ interface IWrappedMToken is IMigratable, IERC20Extended { function balanceWithYieldOf(address account) external view returns (uint256 balance); /** - * @notice Returns the last index of `account`. - * @param account The address of some account. - * @return lastIndex The last index of `account`, 0 if the account is not earning. + * @notice Returns the earning principal of `account`. + * @param account The address of some account. + * @return earningPrincipal The earning principal of `account`. */ - function lastIndexOf(address account) external view returns (uint128 lastIndex); + function earningPrincipalOf(address account) external view returns (uint112 earningPrincipal); /** * @notice Returns the recipient to override as the destination for an account's claim of yield. @@ -261,8 +276,8 @@ interface IWrappedMToken is IMigratable, IERC20Extended { /// @notice The portion of total supply that is earning yield. function totalEarningSupply() external view returns (uint240 totalSupply); - /// @notice The principal of totalEarningSupply to help compute totalAccruedYield(), and thus excess(). - function principalOfTotalEarningSupply() external view returns (uint112 principalOfTotalEarningSupply); + /// @notice The total earning principal to help compute totalAccruedYield(), and thus excess(). + function totalEarningPrincipal() external view returns (uint112 totalEarningPrincipal); /// @notice The address of the destination where excess is claimed to. function excessDestination() external view returns (address excessDestination); diff --git a/test/integration/MorphoBlue.t.sol b/test/integration/MorphoBlue.t.sol index 36ec030..0e1dd2c 100644 --- a/test/integration/MorphoBlue.t.sol +++ b/test/integration/MorphoBlue.t.sol @@ -29,6 +29,8 @@ contract MorphoBlueTests is MorphoTestBase { _deployV2Components(); _migrate(); + _wrappedMToken.migrateAccount(_MORPHO); + _oracle = _createOracle(); _morphoBalanceOfUSDC = IERC20(_USDC).balanceOf(_MORPHO); @@ -95,7 +97,7 @@ contract MorphoBlueTests is MorphoTestBase { vm.warp(vm.getBlockTimestamp() + 365 days); assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM); - assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 49_292101); + assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 49_292100); // USDC balance is unchanged. assertEq(IERC20(_USDC).balanceOf(_MORPHO), _morphoBalanceOfUSDC); @@ -188,7 +190,7 @@ contract MorphoBlueTests is MorphoTestBase { // `startEarningFor` has been called so wM yield has accrued in the pool. assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM); - assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 4_994258); + assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 4_994256); // USDC balance is unchanged. assertEq(IERC20(_USDC).balanceOf(_MORPHO), _morphoBalanceOfUSDC); @@ -222,7 +224,7 @@ contract MorphoBlueTests is MorphoTestBase { // `startEarningFor` has been called so wM yield has accrued in the pool. assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM); - assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 77193); + assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 77192); // USDC balance is unchanged. assertEq(IERC20(_USDC).balanceOf(_MORPHO), _morphoBalanceOfUSDC); diff --git a/test/integration/Protocol.t.sol b/test/integration/Protocol.t.sol index f8459e2..0efc1de 100644 --- a/test/integration/Protocol.t.sol +++ b/test/integration/Protocol.t.sol @@ -81,8 +81,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += 99_999999); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield); - assertEq(_wrappedMToken.excess(), _excess); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 1); + assertEq(_wrappedMToken.excess(), _excess += 1); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -148,8 +148,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += 199_999999); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield); - assertEq(_wrappedMToken.excess(), _excess); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 1); + assertEq(_wrappedMToken.excess(), _excess += 1); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -205,7 +205,7 @@ contract ProtocolIntegrationTests is TestBase { // Assert Alice (Earner) assertEq(_wrappedMToken.balanceOf(_alice), _aliceBalance); - assertEq(_wrappedMToken.accruedYieldOf(_alice), _aliceAccruedYield += 2_423880); + assertEq(_wrappedMToken.accruedYieldOf(_alice), _aliceAccruedYield += 2_423881); // Assert Bob (Earner) assertEq(_wrappedMToken.balanceOf(_bob), _bobBalance); @@ -248,8 +248,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += _aliceBalance); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply += 100_000000); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield); - assertEq(_wrappedMToken.excess(), _excess); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 1); + assertEq(_wrappedMToken.excess(), _excess += 1); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -303,8 +303,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += _bobBalance + 2_395361 - 100_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply += _daveBalance + 100_000000); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 2_395361 - 1); - assertEq(_wrappedMToken.excess(), _excess -= 1); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 2_395361 + 1); + assertEq(_wrappedMToken.excess(), _excess += 2); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -321,8 +321,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += 50_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply -= 50_000000); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield += 1); - assertEq(_wrappedMToken.excess(), _excess -= 1); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield); + assertEq(_wrappedMToken.excess(), _excess); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -374,8 +374,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += 99_999999); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply += 100_000000); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield); - assertEq(_wrappedMToken.excess(), _excess); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 1); + assertEq(_wrappedMToken.excess(), _excess += 1); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -403,8 +403,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += 100_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply += 100_000000); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield += 1); - assertEq(_wrappedMToken.excess(), _excess -= 1); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield); + assertEq(_wrappedMToken.excess(), _excess); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -432,8 +432,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply -= 99_999999); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply += _aliceBalance); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 3_614473); - assertEq(_wrappedMToken.excess(), _excess); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 3_614473 + 1); + assertEq(_wrappedMToken.excess(), _excess += 1); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -449,8 +449,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += _carolBalance); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply -= _carolBalance); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield += 1); - assertEq(_wrappedMToken.excess(), _excess -= 1); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield); + assertEq(_wrappedMToken.excess(), _excess); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); @@ -509,8 +509,8 @@ contract ProtocolIntegrationTests is TestBase { // Assert Globals assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply -= 100_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply); - assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 2_395361); - assertEq(_wrappedMToken.excess(), _excess); + assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 2_395361 + 1); + assertEq(_wrappedMToken.excess(), _excess += 1); assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess); diff --git a/test/integration/UniswapV3.t.sol b/test/integration/UniswapV3.t.sol index e06721d..7364fdd 100644 --- a/test/integration/UniswapV3.t.sol +++ b/test/integration/UniswapV3.t.sol @@ -56,6 +56,8 @@ contract UniswapV3IntegrationTests is TestBase { _deployV2Components(); _migrate(); + _wrappedMToken.migrateAccount(_pool); + _poolClaimRecipient = _wrappedMToken.claimOverrideRecipientFor(_pool); _wrapperBalanceOfM = _mToken.balanceOf(address(_wrappedMToken)); @@ -338,7 +340,7 @@ contract UniswapV3IntegrationTests is TestBase { // Move 5 days forward and check that yield has accrued. vm.warp(vm.getBlockTimestamp() + 5 days); - assertEq(_wrappedMToken.accruedYieldOf(_pool), _poolAccruedYield += 11_753_024234); + assertEq(_wrappedMToken.accruedYieldOf(_pool), _poolAccruedYield += 11_753_024235); /* ============ Eric (Earner) Swaps Exact wM for USDC ============ */ diff --git a/test/unit/Stories.t.sol b/test/unit/Stories.t.sol index 906a3a0..145d080 100644 --- a/test/unit/Stories.t.sol +++ b/test/unit/Stories.t.sol @@ -194,8 +194,8 @@ contract StoryTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 300_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), 300_000000); assertEq(_wrappedMToken.totalSupply(), 600_000000); - assertEq(_wrappedMToken.totalAccruedYield(), 50_000001); - assertEq(_wrappedMToken.excess(), 249_999999); + assertEq(_wrappedMToken.totalAccruedYield(), 49_999998); + assertEq(_wrappedMToken.excess(), 250_000002); vm.prank(_dave); _wrappedMToken.transfer(_bob, 50_000000); @@ -212,8 +212,8 @@ contract StoryTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 400_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), 250_000000); assertEq(_wrappedMToken.totalSupply(), 650_000000); - assertEq(_wrappedMToken.totalAccruedYield(), 2); - assertEq(_wrappedMToken.excess(), 249_999996); + assertEq(_wrappedMToken.totalAccruedYield(), 0); + assertEq(_wrappedMToken.excess(), 250_000002); _mToken.setCurrentIndex(4 * _EXP_SCALED_ONE); _mToken.setBalanceOf(address(_wrappedMToken), 1_200_000000); // was 900 @ 3.0, so 1200 @ 4.0 @@ -238,8 +238,8 @@ contract StoryTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 400_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), 250_000000); assertEq(_wrappedMToken.totalSupply(), 650_000000); - assertEq(_wrappedMToken.totalAccruedYield(), 133_333336); - assertEq(_wrappedMToken.excess(), 416_666664); + assertEq(_wrappedMToken.totalAccruedYield(), 133_333328); + assertEq(_wrappedMToken.excess(), 416_666672); _registrar.setListContains(_EARNERS_LIST_NAME, _alice, false); @@ -253,8 +253,8 @@ contract StoryTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 200_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), 516_666664); assertEq(_wrappedMToken.totalSupply(), 716_666664); - assertEq(_wrappedMToken.totalAccruedYield(), 66_666672); - assertEq(_wrappedMToken.excess(), 416_666664); + assertEq(_wrappedMToken.totalAccruedYield(), 66_666664); + assertEq(_wrappedMToken.excess(), 416_666672); _registrar.setListContains(_EARNERS_LIST_NAME, _carol, true); @@ -268,8 +268,8 @@ contract StoryTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 400_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), 316_666664); assertEq(_wrappedMToken.totalSupply(), 716_666664); - assertEq(_wrappedMToken.totalAccruedYield(), 66_666672); - assertEq(_wrappedMToken.excess(), 416_666664); + assertEq(_wrappedMToken.totalAccruedYield(), 66_666664); + assertEq(_wrappedMToken.excess(), 416_666672); _mToken.setCurrentIndex(5 * _EXP_SCALED_ONE); _mToken.setBalanceOf(address(_wrappedMToken), 1_500_000000); // was 1200 @ 4.0, so 1500 @ 5.0 @@ -294,8 +294,8 @@ contract StoryTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 400_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), 316_666664); assertEq(_wrappedMToken.totalSupply(), 716_666664); - assertEq(_wrappedMToken.totalAccruedYield(), 183_333340); - assertEq(_wrappedMToken.excess(), 599_999995); + assertEq(_wrappedMToken.totalAccruedYield(), 183_333330); + assertEq(_wrappedMToken.excess(), 600_000005); vm.prank(_alice); _wrappedMToken.unwrap(_alice, 266_666664); @@ -308,8 +308,8 @@ contract StoryTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 400_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), 50_000000); assertEq(_wrappedMToken.totalSupply(), 450_000000); - assertEq(_wrappedMToken.totalAccruedYield(), 183_333340); - assertEq(_wrappedMToken.excess(), 600_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 183_333330); + assertEq(_wrappedMToken.excess(), 600_000010); vm.prank(_bob); _wrappedMToken.unwrap(_bob, 333_333330); @@ -322,8 +322,8 @@ contract StoryTests is Test { assertEq(_wrappedMToken.totalEarningSupply(), 200_000000); assertEq(_wrappedMToken.totalNonEarningSupply(), 50_000000); assertEq(_wrappedMToken.totalSupply(), 250_000000); - assertEq(_wrappedMToken.totalAccruedYield(), 50_000010); - assertEq(_wrappedMToken.excess(), 600_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 50_000000); + assertEq(_wrappedMToken.excess(), 600_000010); vm.prank(_carol); _wrappedMToken.unwrap(_carol, 250_000000); diff --git a/test/unit/WrappedMToken.t.sol b/test/unit/WrappedMToken.t.sol index d29de51..3f33d31 100644 --- a/test/unit/WrappedMToken.t.sol +++ b/test/unit/WrappedMToken.t.sol @@ -128,22 +128,22 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setAccountOf(_alice, 1_000); - assertEq(_wrappedMToken.lastIndexOf(_alice), 0); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 1_000); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); assertEq(_wrappedMToken.totalAccruedYield(), 0); vm.prank(_alice); assertEq(_wrappedMToken.wrap(_alice, 1_000), 1_000); - assertEq(_wrappedMToken.lastIndexOf(_alice), 0); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_alice), 2_000); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 2_000); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -154,27 +154,27 @@ contract WrappedMTokenTests is Test { _mToken.setBalanceOf(_alice, 1_002); - _wrappedMToken.setPrincipalOfTotalEarningSupply(1_000); + _wrappedMToken.setTotalEarningPrincipal(1_000); _wrappedMToken.setTotalEarningSupply(1_000); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); // 1_100 balance with yield. + _wrappedMToken.setAccountOf(_alice, 1_000, 1_000); // 1_100 balance with yield. - assertEq(_wrappedMToken.lastIndexOf(_alice), _EXP_SCALED_ONE); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 1_000); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); assertEq(_wrappedMToken.accruedYieldOf(_alice), 100); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 1_000); + assertEq(_wrappedMToken.totalEarningPrincipal(), 1_000); assertEq(_wrappedMToken.totalEarningSupply(), 1_000); assertEq(_wrappedMToken.totalAccruedYield(), 100); vm.prank(_alice); assertEq(_wrappedMToken.wrap(_alice, 999), 999); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 1_000 + 908); assertEq(_wrappedMToken.balanceOf(_alice), 1_000 + 100 + 999); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 1_000 + 909); + assertEq(_wrappedMToken.totalEarningPrincipal(), 1_000 + 908); assertEq(_wrappedMToken.totalEarningSupply(), 1_000 + 100 + 999); assertEq(_wrappedMToken.totalAccruedYield(), 0); @@ -182,24 +182,24 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.wrap(_alice, 1), 1); // No change due to principal round down on wrap. - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 1_000 + 908 + 0); assertEq(_wrappedMToken.balanceOf(_alice), 1_000 + 100 + 999 + 1); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 1_000 + 909 + 1); + assertEq(_wrappedMToken.totalEarningPrincipal(), 1_000 + 908 + 0); assertEq(_wrappedMToken.totalEarningSupply(), 1_000 + 100 + 999 + 1); - assertEq(_wrappedMToken.totalAccruedYield(), 1); + assertEq(_wrappedMToken.totalAccruedYield(), 0); vm.prank(_alice); assertEq(_wrappedMToken.wrap(_alice, 2), 2); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 1_000 + 908 + 0 + 1); assertEq(_wrappedMToken.balanceOf(_alice), 1_000 + 100 + 999 + 1 + 2); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 1_000 + 909 + 1 + 2); + assertEq(_wrappedMToken.totalEarningPrincipal(), 1_000 + 908 + 0 + 1); assertEq(_wrappedMToken.totalEarningSupply(), 1_000 + 100 + 999 + 1 + 2); - assertEq(_wrappedMToken.totalAccruedYield(), 1); + assertEq(_wrappedMToken.totalAccruedYield(), 0); } function testFuzz_wrap( @@ -325,7 +325,7 @@ contract WrappedMTokenTests is Test { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setAccountOf(_alice, 909, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 999, 909); vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); vm.prank(_alice); @@ -343,44 +343,44 @@ contract WrappedMTokenTests is Test { _wrappedMToken.setAccountOf(_alice, 1_000); - assertEq(_wrappedMToken.lastIndexOf(_alice), 0); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 1_000); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); assertEq(_wrappedMToken.totalAccruedYield(), 0); vm.prank(_alice); assertEq(_wrappedMToken.unwrap(_alice, 1), 0); - assertEq(_wrappedMToken.lastIndexOf(_alice), 0); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_alice), 999); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 999); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); assertEq(_wrappedMToken.totalAccruedYield(), 0); vm.prank(_alice); assertEq(_wrappedMToken.unwrap(_alice, 499), 498); - assertEq(_wrappedMToken.lastIndexOf(_alice), 0); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_alice), 500); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 500); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); assertEq(_wrappedMToken.totalAccruedYield(), 0); vm.prank(_alice); assertEq(_wrappedMToken.unwrap(_alice, 500), 499); - assertEq(_wrappedMToken.lastIndexOf(_alice), 0); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_alice), 0); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -392,16 +392,16 @@ contract WrappedMTokenTests is Test { _mToken.setBalanceOf(address(_wrappedMToken), 1_000); - _wrappedMToken.setPrincipalOfTotalEarningSupply(909); + _wrappedMToken.setTotalEarningPrincipal(909); _wrappedMToken.setTotalEarningSupply(909); - _wrappedMToken.setAccountOf(_alice, 909, _EXP_SCALED_ONE); // 999 balance with yield. + _wrappedMToken.setAccountOf(_alice, 909, 909); // 999 balance with yield. - assertEq(_wrappedMToken.lastIndexOf(_alice), _EXP_SCALED_ONE); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 909); assertEq(_wrappedMToken.balanceOf(_alice), 909); assertEq(_wrappedMToken.accruedYieldOf(_alice), 90); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 909); + assertEq(_wrappedMToken.totalEarningPrincipal(), 909); assertEq(_wrappedMToken.totalEarningSupply(), 909); assertEq(_wrappedMToken.totalAccruedYield(), 90); @@ -409,33 +409,33 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.unwrap(_alice, 1), 0); // Change due to principal round up on unwrap. - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 909 - 1); assertEq(_wrappedMToken.balanceOf(_alice), 999 - 1); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 909 - 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 909 - 1); assertEq(_wrappedMToken.totalEarningSupply(), 999 - 1); - assertEq(_wrappedMToken.totalAccruedYield(), 1); + assertEq(_wrappedMToken.totalAccruedYield(), 0); vm.prank(_alice); assertEq(_wrappedMToken.unwrap(_alice, 498), 497); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 909 - 1 - 453); assertEq(_wrappedMToken.balanceOf(_alice), 999 - 1 - 498); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 909 - 0 - 452); + assertEq(_wrappedMToken.totalEarningPrincipal(), 909 - 1 - 453); assertEq(_wrappedMToken.totalEarningSupply(), 999 - 1 - 498); - assertEq(_wrappedMToken.totalAccruedYield(), 2); + assertEq(_wrappedMToken.totalAccruedYield(), 0); vm.prank(_alice); assertEq(_wrappedMToken.unwrap(_alice, 500), 499); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 909 - 1 - 453 - 455); // 0 assertEq(_wrappedMToken.balanceOf(_alice), 999 - 1 - 498 - 500); // 0 assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); // 0 due to 0 `totalEarningSupply`. + assertEq(_wrappedMToken.totalEarningPrincipal(), 909 - 1 - 453 - 455); // 0 assertEq(_wrappedMToken.totalEarningSupply(), 999 - 1 - 498 - 500); // 0 assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -561,10 +561,10 @@ contract WrappedMTokenTests is Test { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setPrincipalOfTotalEarningSupply(1_000); + _wrappedMToken.setTotalEarningPrincipal(1_000); _wrappedMToken.setTotalEarningSupply(1_000); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); // 1_100 balance with yield. + _wrappedMToken.setAccountOf(_alice, 1_000, 1_000); // 1_100 balance with yield. assertEq(_wrappedMToken.balanceOf(_alice), 1_000); assertEq(_wrappedMToken.accruedYieldOf(_alice), 100); @@ -577,10 +577,10 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.claimFor(_alice), 100); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 1_000); assertEq(_wrappedMToken.balanceOf(_alice), 1_100); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 1_000); + assertEq(_wrappedMToken.totalEarningPrincipal(), 1_000); assertEq(_wrappedMToken.totalEarningSupply(), 1_100); assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -594,10 +594,10 @@ contract WrappedMTokenTests is Test { bytes32(uint256(uint160(_bob))) ); - _wrappedMToken.setPrincipalOfTotalEarningSupply(1_000); + _wrappedMToken.setTotalEarningPrincipal(1_000); _wrappedMToken.setTotalEarningSupply(1_000); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); // 1_100 balance with yield. + _wrappedMToken.setAccountOf(_alice, 1_000, 1_000); // 1_100 balance with yield. assertEq(_wrappedMToken.balanceOf(_alice), 1_000); assertEq(_wrappedMToken.accruedYieldOf(_alice), 100); @@ -613,16 +613,16 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.claimFor(_alice), 100); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 909); assertEq(_wrappedMToken.balanceOf(_alice), 1_000); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_bob), 100); assertEq(_wrappedMToken.totalNonEarningSupply(), 100); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 910); + assertEq(_wrappedMToken.totalEarningPrincipal(), 909); assertEq(_wrappedMToken.totalEarningSupply(), 1_000); - assertEq(_wrappedMToken.totalAccruedYield(), 1); + assertEq(_wrappedMToken.totalAccruedYield(), 0); } function testFuzz_claimFor( @@ -680,7 +680,7 @@ contract WrappedMTokenTests is Test { uint128 enableMIndex_, uint128 disableIndex_, uint240 totalNonEarningSupply_, - uint240 totalProjectedEarningSupply_, + uint240 projectedTotalEarningSupply_, uint240 mBalance_ ) external { (currentMIndex_, enableMIndex_, disableIndex_) = _getFuzzedIndices( @@ -691,24 +691,25 @@ contract WrappedMTokenTests is Test { _setupIndexes(earningEnabled_, currentMIndex_, enableMIndex_, disableIndex_); - uint240 maxAmount_ = _getMaxAmount(_wrappedMToken.currentIndex()); + uint128 currentIndex_ = _wrappedMToken.currentIndex(); + uint240 maxAmount_ = _getMaxAmount(currentIndex_); totalNonEarningSupply_ = uint240(bound(totalNonEarningSupply_, 0, maxAmount_)); - totalProjectedEarningSupply_ = uint240( - bound(totalProjectedEarningSupply_, 0, maxAmount_ - totalNonEarningSupply_) + projectedTotalEarningSupply_ = uint240( + bound(projectedTotalEarningSupply_, 0, maxAmount_ - totalNonEarningSupply_) ); - uint112 principalOfTotalEarningSupply_ = IndexingMath.getPrincipalAmountRoundedUp( - totalProjectedEarningSupply_, - _wrappedMToken.currentIndex() + uint112 totalEarningPrincipal_ = IndexingMath.getPrincipalAmountRoundedUp( + projectedTotalEarningSupply_, + currentIndex_ ); mBalance_ = uint240(bound(mBalance_, 0, maxAmount_)); _mToken.setBalanceOf(address(_wrappedMToken), mBalance_); - _wrappedMToken.setPrincipalOfTotalEarningSupply(principalOfTotalEarningSupply_); + _wrappedMToken.setTotalEarningPrincipal(totalEarningPrincipal_); _wrappedMToken.setTotalNonEarningSupply(totalNonEarningSupply_); uint240 expectedExcess_ = _wrappedMToken.excess(); @@ -754,7 +755,7 @@ contract WrappedMTokenTests is Test { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setAccountOf(_alice, 909, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 909, 909); vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); vm.prank(_alice); @@ -778,7 +779,7 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.balanceOf(_bob), 1_000); assertEq(_wrappedMToken.totalNonEarningSupply(), 1_500); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); } @@ -807,7 +808,7 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.balanceOf(_bob), bobBalance + transferAmount_); assertEq(_wrappedMToken.totalNonEarningSupply(), supply_); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); } @@ -815,12 +816,12 @@ contract WrappedMTokenTests is Test { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setPrincipalOfTotalEarningSupply(1_000); + _wrappedMToken.setTotalEarningPrincipal(1_000); _wrappedMToken.setTotalEarningSupply(1_000); _wrappedMToken.setTotalNonEarningSupply(500); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); // 1_100 balance with yield. + _wrappedMToken.setAccountOf(_alice, 1_000, 1_000); // 1_100 balance with yield. _wrappedMToken.setAccountOf(_bob, 500); assertEq(_wrappedMToken.accruedYieldOf(_alice), 100); @@ -837,14 +838,14 @@ contract WrappedMTokenTests is Test { vm.prank(_alice); _wrappedMToken.transfer(_bob, 500); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 545); assertEq(_wrappedMToken.balanceOf(_alice), 600); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_bob), 1_000); assertEq(_wrappedMToken.totalNonEarningSupply(), 1_000); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 546); + assertEq(_wrappedMToken.totalEarningPrincipal(), 545); assertEq(_wrappedMToken.totalEarningSupply(), 600); assertEq(_wrappedMToken.totalAccruedYield(), 0); @@ -854,29 +855,29 @@ contract WrappedMTokenTests is Test { vm.prank(_alice); _wrappedMToken.transfer(_bob, 1); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 544); assertEq(_wrappedMToken.balanceOf(_alice), 599); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_bob), 1_001); assertEq(_wrappedMToken.totalNonEarningSupply(), 1_001); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 546); + assertEq(_wrappedMToken.totalEarningPrincipal(), 544); assertEq(_wrappedMToken.totalEarningSupply(), 599); - assertEq(_wrappedMToken.totalAccruedYield(), 1); + assertEq(_wrappedMToken.totalAccruedYield(), 0); } function test_transfer_fromNonEarner_toEarner() external { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setPrincipalOfTotalEarningSupply(500); + _wrappedMToken.setTotalEarningPrincipal(500); _wrappedMToken.setTotalEarningSupply(500); _wrappedMToken.setTotalNonEarningSupply(1_000); _wrappedMToken.setAccountOf(_alice, 1_000); - _wrappedMToken.setAccountOf(_bob, 500, _EXP_SCALED_ONE); // 550 balance with yield. + _wrappedMToken.setAccountOf(_bob, 500, 500); // 550 balance with yield. assertEq(_wrappedMToken.accruedYieldOf(_bob), 50); @@ -894,12 +895,12 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.balanceOf(_alice), 500); - assertEq(_wrappedMToken.lastIndexOf(_bob), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_bob), 954); assertEq(_wrappedMToken.balanceOf(_bob), 1_050); assertEq(_wrappedMToken.accruedYieldOf(_bob), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 500); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 955); + assertEq(_wrappedMToken.totalEarningPrincipal(), 954); assertEq(_wrappedMToken.totalEarningSupply(), 1_050); assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -908,11 +909,11 @@ contract WrappedMTokenTests is Test { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setPrincipalOfTotalEarningSupply(1_500); + _wrappedMToken.setTotalEarningPrincipal(1_500); _wrappedMToken.setTotalEarningSupply(1_500); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); // 1_100 balance with yield. - _wrappedMToken.setAccountOf(_bob, 500, _EXP_SCALED_ONE); // 550 balance with yield. + _wrappedMToken.setAccountOf(_alice, 1_000, 1_000); // 1_100 balance with yield. + _wrappedMToken.setAccountOf(_bob, 500, 500); // 550 balance with yield. assertEq(_wrappedMToken.accruedYieldOf(_alice), 100); assertEq(_wrappedMToken.accruedYieldOf(_bob), 50); @@ -935,16 +936,16 @@ contract WrappedMTokenTests is Test { vm.prank(_alice); _wrappedMToken.transfer(_bob, 500); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 545); assertEq(_wrappedMToken.balanceOf(_alice), 600); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); - assertEq(_wrappedMToken.lastIndexOf(_bob), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_bob), 954); assertEq(_wrappedMToken.balanceOf(_bob), 1_050); assertEq(_wrappedMToken.accruedYieldOf(_bob), 0); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 1_500); + assertEq(_wrappedMToken.totalEarningPrincipal(), 1_499); assertEq(_wrappedMToken.totalEarningSupply(), 1_650); assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -969,10 +970,10 @@ contract WrappedMTokenTests is Test { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setPrincipalOfTotalEarningSupply(1_000); + _wrappedMToken.setTotalEarningPrincipal(1_000); _wrappedMToken.setTotalEarningSupply(1_000); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); // 1_100 balance with yield. + _wrappedMToken.setAccountOf(_alice, 1_000, 1_000); // 1_100 balance with yield. assertEq(_wrappedMToken.balanceOf(_alice), 1_000); assertEq(_wrappedMToken.accruedYieldOf(_alice), 100); @@ -989,11 +990,11 @@ contract WrappedMTokenTests is Test { vm.prank(_alice); _wrappedMToken.transfer(_alice, 500); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 1_000); assertEq(_wrappedMToken.balanceOf(_alice), 1_100); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 1_000); + assertEq(_wrappedMToken.totalEarningPrincipal(), 1_000); assertEq(_wrappedMToken.totalEarningSupply(), 1_100); assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -1116,11 +1117,11 @@ contract WrappedMTokenTests is Test { _wrappedMToken.startEarningFor(_alice); assertEq(_wrappedMToken.isEarning(_alice), true); - assertEq(_wrappedMToken.lastIndexOf(_alice), 1_100000000000); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 909); assertEq(_wrappedMToken.balanceOf(_alice), 1000); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 910); + assertEq(_wrappedMToken.totalEarningPrincipal(), 909); assertEq(_wrappedMToken.totalEarningSupply(), 1_000); } @@ -1139,9 +1140,11 @@ contract WrappedMTokenTests is Test { _setupIndexes(earningEnabled_, currentMIndex_, enableMIndex_, disableIndex_); - balance_ = uint240(bound(balance_, 0, _getMaxAmount(_wrappedMToken.currentIndex()))); + uint128 currentIndex_ = _wrappedMToken.currentIndex(); + + balance_ = uint240(bound(balance_, 0, _getMaxAmount(currentIndex_))); - _setupAccount(_alice, false, 0, balance_); + _setupAccount(_alice, false, balance_, balance_); _registrar.setListContains(_EARNERS_LIST_NAME, _alice, true); @@ -1150,12 +1153,15 @@ contract WrappedMTokenTests is Test { _wrappedMToken.startEarningFor(_alice); + uint112 earningPrincipal_ = IndexingMath.getPrincipalAmountRoundedDown(balance_, currentIndex_); + assertEq(_wrappedMToken.isEarning(_alice), true); - assertEq(_wrappedMToken.lastIndexOf(_alice), _wrappedMToken.currentIndex()); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), earningPrincipal_); assertEq(_wrappedMToken.balanceOf(_alice), balance_); assertEq(_wrappedMToken.totalNonEarningSupply(), 0); assertEq(_wrappedMToken.totalEarningSupply(), balance_); + assertEq(_wrappedMToken.totalEarningPrincipal(), earningPrincipal_); } /* ============ startEarningFor batch ============ */ @@ -1199,10 +1205,10 @@ contract WrappedMTokenTests is Test { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setPrincipalOfTotalEarningSupply(1_000); + _wrappedMToken.setTotalEarningPrincipal(1_000); _wrappedMToken.setTotalEarningSupply(1_000); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); // 1_100 balance with yield. + _wrappedMToken.setAccountOf(_alice, 1_000, 1_000); // 1_100 balance with yield. assertEq(_wrappedMToken.accruedYieldOf(_alice), 100); @@ -1217,13 +1223,13 @@ contract WrappedMTokenTests is Test { _wrappedMToken.stopEarningFor(_alice); - assertEq(_wrappedMToken.lastIndexOf(_alice), 0); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_alice), 1_100); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.isEarning(_alice), false); assertEq(_wrappedMToken.totalNonEarningSupply(), 1_100); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -1267,13 +1273,13 @@ contract WrappedMTokenTests is Test { _wrappedMToken.stopEarningFor(_alice); - assertEq(_wrappedMToken.lastIndexOf(_alice), 0); + assertEq(_wrappedMToken.earningPrincipalOf(_alice), 0); assertEq(_wrappedMToken.balanceOf(_alice), balance_ + accruedYield_); assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); assertEq(_wrappedMToken.isEarning(_alice), false); assertEq(_wrappedMToken.totalNonEarningSupply(), balance_ + accruedYield_); - assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningPrincipal(), 0); assertEq(_wrappedMToken.totalEarningSupply(), 0); assertEq(_wrappedMToken.totalAccruedYield(), 0); } @@ -1291,8 +1297,8 @@ contract WrappedMTokenTests is Test { } function test_stopEarningFor_batch() external { - _wrappedMToken.setAccountOf(_alice, 0, _EXP_SCALED_ONE); - _wrappedMToken.setAccountOf(_bob, 0, _EXP_SCALED_ONE); + _wrappedMToken.setAccountOf(_alice, 0, 0); + _wrappedMToken.setAccountOf(_bob, 0, 0); address[] memory accounts_ = new address[](2); accounts_[0] = _alice; @@ -1313,9 +1319,26 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); } + function test_enableEarning_firstTime() external { + _registrar.setListContains(_EARNERS_LIST_NAME, address(_wrappedMToken), true); + + _mToken.setCurrentIndex(1_100000000000); + + vm.expectEmit(); + emit IWrappedMToken.EarningEnabled(1_100000000000); + + _wrappedMToken.enableEarning(); + + assertEq(_wrappedMToken.disableIndex(), 0); + assertEq(_wrappedMToken.enableMIndex(), 1_100000000000); + assertEq(_wrappedMToken.currentIndex(), _EXP_SCALED_ONE); + } + function test_enableEarning() external { _registrar.setListContains(_EARNERS_LIST_NAME, address(_wrappedMToken), true); + _wrappedMToken.setDisableIndex(1_100000000000); + _mToken.setCurrentIndex(1_210000000000); vm.expectEmit(); @@ -1323,7 +1346,9 @@ contract WrappedMTokenTests is Test { _wrappedMToken.enableEarning(); + assertEq(_wrappedMToken.disableIndex(), 1_100000000000); assertEq(_wrappedMToken.enableMIndex(), 1_210000000000); + assertEq(_wrappedMToken.currentIndex(), 1_100000000000); // 1.21 / 1.10 } /* ============ disableEarning ============ */ @@ -1347,6 +1372,10 @@ contract WrappedMTokenTests is Test { emit IWrappedMToken.EarningDisabled(1_100000000000); _wrappedMToken.disableEarning(); + + assertEq(_wrappedMToken.disableIndex(), 1_100000000000); + assertEq(_wrappedMToken.enableMIndex(), 0); + assertEq(_wrappedMToken.currentIndex(), 1_100000000000); } /* ============ balanceOf ============ */ @@ -1364,15 +1393,19 @@ contract WrappedMTokenTests is Test { _mToken.setCurrentIndex(1_210000000000); _wrappedMToken.setEnableMIndex(1_100000000000); - _wrappedMToken.setAccountOf(_alice, 500, _EXP_SCALED_ONE); // 550 balance with yield. + _wrappedMToken.setAccountOf(_alice, 500, 500); // 550 balance with yield. assertEq(_wrappedMToken.balanceOf(_alice), 500); - _wrappedMToken.setAccountOf(_alice, 1_000, _EXP_SCALED_ONE); // 1_100 balance with yield. + _wrappedMToken.setEarningPrincipalOf(_alice, 1_000); // Earning principal has no bearing on balance. + + assertEq(_wrappedMToken.balanceOf(_alice), 500); + + _wrappedMToken.setAccountOf(_alice, 1_000, 1_000); // 1_100 balance with yield. assertEq(_wrappedMToken.balanceOf(_alice), 1_000); - _wrappedMToken.setLastIndexOf(_alice, 1_100000000000); // Last index has no bearing on balance. + _wrappedMToken.setEarningPrincipalOf(_alice, 2_000); // Earning principal has no bearing on balance. assertEq(_wrappedMToken.balanceOf(_alice), 1_000); } @@ -1475,7 +1508,7 @@ contract WrappedMTokenTests is Test { assertEq(_wrappedMToken.excess(), 0); _wrappedMToken.setTotalNonEarningSupply(1_000); - _wrappedMToken.setPrincipalOfTotalEarningSupply(1_000); + _wrappedMToken.setTotalEarningPrincipal(1_000); _wrappedMToken.setTotalEarningSupply(1_000); _mToken.setBalanceOf(address(_wrappedMToken), 2_100); @@ -1524,7 +1557,7 @@ contract WrappedMTokenTests is Test { bound(totalProjectedEarningSupply_, 0, maxAmount_ - totalNonEarningSupply_) ); - uint112 principalOfTotalEarningSupply_ = IndexingMath.getPrincipalAmountRoundedUp( + uint112 totalEarningPrincipal_ = IndexingMath.getPrincipalAmountRoundedUp( totalProjectedEarningSupply_, _wrappedMToken.currentIndex() ); @@ -1533,7 +1566,7 @@ contract WrappedMTokenTests is Test { _mToken.setBalanceOf(address(_wrappedMToken), mBalance_); - _wrappedMToken.setPrincipalOfTotalEarningSupply(principalOfTotalEarningSupply_); + _wrappedMToken.setTotalEarningPrincipal(totalEarningPrincipal_); _wrappedMToken.setTotalNonEarningSupply(totalNonEarningSupply_); uint240 totalProjectedSupply_ = totalNonEarningSupply_ + totalProjectedEarningSupply_; @@ -1597,7 +1630,6 @@ contract WrappedMTokenTests is Test { balanceWithYield_ = uint240(bound(balanceWithYield_, 0, maxAmount_)); balance_ = uint240(bound(balance_, (balanceWithYield_ * _EXP_SCALED_ONE) / currentIndex_, balanceWithYield_)); - balance_ = balance_ == 0 ? 1 : balance_; return (balanceWithYield_, balance_); } @@ -1609,17 +1641,13 @@ contract WrappedMTokenTests is Test { uint240 balance_ ) internal { if (accountEarning_) { - uint128 lastIndex_ = balanceWithYield_ == 0 - ? _EXP_SCALED_ONE - : uint128((balance_ * _wrappedMToken.currentIndex()) / balanceWithYield_); - - _wrappedMToken.setAccountOf(account_, balance_, lastIndex_); - - _wrappedMToken.setPrincipalOfTotalEarningSupply( - _wrappedMToken.principalOfTotalEarningSupply() + - IndexingMath.getPrincipalAmountRoundedUp(balance_, lastIndex_) + uint112 principal_ = IndexingMath.getPrincipalAmountRoundedDown( + balanceWithYield_, + _wrappedMToken.currentIndex() ); + _wrappedMToken.setAccountOf(account_, balance_, principal_); + _wrappedMToken.setTotalEarningPrincipal(_wrappedMToken.totalEarningPrincipal() + principal_); _wrappedMToken.setTotalEarningSupply(_wrappedMToken.totalEarningSupply() + balance_); } else { _wrappedMToken.setAccountOf(account_, balance_); diff --git a/test/utils/Invariants.sol b/test/utils/Invariants.sol index 9f1eff6..3293c95 100644 --- a/test/utils/Invariants.sol +++ b/test/utils/Invariants.sol @@ -90,26 +90,23 @@ library Invariants { // Invariant 4: Sum of all earning accounts' principals is less than or equal to principal of total earning supply. function checkInvariant4(address wrappedMToken_, address[] memory accounts_) internal view returns (bool success_) { - uint256 principalOfTotalEarningSupply_; + uint256 totalEarningPrincipal_; for (uint256 i_; i_ < accounts_.length; ++i_) { address account_ = accounts_[i_]; if (!IWrappedMToken(wrappedMToken_).isEarning(account_)) continue; - principalOfTotalEarningSupply_ += IndexingMath.getPrincipalAmountRoundedDown( - uint240(IWrappedMToken(wrappedMToken_).balanceOf(account_)), - IWrappedMToken(wrappedMToken_).lastIndexOf(account_) - ); + totalEarningPrincipal_ += IWrappedMToken(wrappedMToken_).earningPrincipalOf(account_); } - // console2.log("Invariant 2: principalOfTotalEarningSupply_ = %d", principalOfTotalEarningSupply_); + // console2.log("Invariant 2: totalEarningPrincipal_ = %d", totalEarningPrincipal_); // console2.log( - // "Invariant 2: principalOfTotalEarningSupply() = %d", - // IWrappedMToken(wrappedMToken_).principalOfTotalEarningSupply() + // "Invariant 2: totalEarningPrincipal() = %d", + // IWrappedMToken(wrappedMToken_).totalEarningPrincipal() // ); - return IWrappedMToken(wrappedMToken_).principalOfTotalEarningSupply() >= principalOfTotalEarningSupply_; + return IWrappedMToken(wrappedMToken_).totalEarningPrincipal() >= totalEarningPrincipal_; } } diff --git a/test/utils/WrappedMTokenHarness.sol b/test/utils/WrappedMTokenHarness.sol index 185e0aa..973d64e 100644 --- a/test/utils/WrappedMTokenHarness.sol +++ b/test/utils/WrappedMTokenHarness.sol @@ -13,19 +13,19 @@ contract WrappedMTokenHarness is WrappedMToken { ) WrappedMToken(mToken_, registrar_, excessDestination_, migrationAdmin_) {} function setIsEarningOf(address account_, bool isEarning_) external { - _accounts[account_].isEarning = isEarning_; + _accounts[account_].earningState = isEarning_ ? EarningState.PRINCIPAL_BASED : EarningState.NOT_EARNING; } - function setLastIndexOf(address account_, uint256 index_) external { - _accounts[account_].lastIndex = uint128(index_); + function setEarningPrincipalOf(address account_, uint256 earningPrincipal_) external { + _accounts[account_].earningPrincipal = uint112(earningPrincipal_); } - function setAccountOf(address account_, uint256 balance_, uint256 index_) external { - _accounts[account_] = Account(true, uint240(balance_), uint128(index_)); + function setAccountOf(address account_, uint256 balance_, uint256 earningPrincipal_) external { + _accounts[account_] = Account(EarningState.PRINCIPAL_BASED, uint240(balance_), uint112(earningPrincipal_)); } function setAccountOf(address account_, uint256 balance_) external { - _accounts[account_] = Account(false, uint240(balance_), 0); + _accounts[account_] = Account(EarningState.NOT_EARNING, uint240(balance_), 0); } function setTotalNonEarningSupply(uint256 totalNonEarningSupply_) external { @@ -36,8 +36,8 @@ contract WrappedMTokenHarness is WrappedMToken { totalEarningSupply = uint240(totalEarningSupply_); } - function setPrincipalOfTotalEarningSupply(uint256 principalOfTotalEarningSupply_) external { - principalOfTotalEarningSupply = uint112(principalOfTotalEarningSupply_); + function setTotalEarningPrincipal(uint256 totalEarningPrincipal_) external { + totalEarningPrincipal = uint112(totalEarningPrincipal_); } function setEnableMIndex(uint256 enableMIndex_) external { @@ -47,4 +47,11 @@ contract WrappedMTokenHarness is WrappedMToken { function setDisableIndex(uint256 disableIndex_) external { disableIndex = uint128(disableIndex_); } + + function getAccountOf( + address account_ + ) external view returns (bool isEarning_, uint240 balance_, uint112 earningPrincipal_) { + Account storage account = _accounts[account_]; + return (account.earningState == EarningState.PRINCIPAL_BASED, account.balance, account.earningPrincipal); + } }