From bda198a46354895450455ad9e4ea1dddabe83575 Mon Sep 17 00:00:00 2001 From: Michael De Luca Date: Wed, 26 Jun 2024 04:27:36 -0400 Subject: [PATCH 1/3] feat: tests and some relevant cleanup --- src/{WrappedM.sol => WrappedMToken.sol} | 165 ++++-- .../{IWrappedM.sol => IWrappedMToken.sol} | 33 +- test/Test.t.sol | 433 +++++++-------- test/WrappedMToken.t.sol | 518 ++++++++++++++++++ test/utils/Mocks.sol | 51 ++ test/utils/WrappedMTokenHarness.sol | 48 ++ 6 files changed, 952 insertions(+), 296 deletions(-) rename src/{WrappedM.sol => WrappedMToken.sol} (66%) rename src/interfaces/{IWrappedM.sol => IWrappedMToken.sol} (51%) create mode 100644 test/WrappedMToken.t.sol create mode 100644 test/utils/Mocks.sol create mode 100644 test/utils/WrappedMTokenHarness.sol diff --git a/src/WrappedM.sol b/src/WrappedMToken.sol similarity index 66% rename from src/WrappedM.sol rename to src/WrappedMToken.sol index 4f2009c..87ff58b 100644 --- a/src/WrappedM.sol +++ b/src/WrappedMToken.sol @@ -9,12 +9,12 @@ import { ERC20Extended } from "../lib/common/src/ERC20Extended.sol"; import { IndexingMath } from "./libs/IndexingMath.sol"; import { IMTokenLike } from "./interfaces/IMTokenLike.sol"; -import { IWrappedM } from "./interfaces/IWrappedM.sol"; +import { IWrappedMToken } from "./interfaces/IWrappedMToken.sol"; import { IRegistrarLike } from "./interfaces/IRegistrarLike.sol"; import { Migratable } from "./Migratable.sol"; -contract WrappedM is IWrappedM, Migratable, ERC20Extended { +contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended { type BalanceInfo is uint256; /* ============ Variables ============ */ @@ -23,7 +23,7 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { bytes32 internal constant _EARNERS_LIST_IGNORED = "earners_list_ignored"; bytes32 internal constant _EARNERS_LIST = "earners"; - bytes32 internal constant _CLAIM_DESTINATION_PREFIX = "wm_claim_destination"; + bytes32 internal constant _CLAIM_OVERRIDE_RECIPIENT_PREFIX = "wm_claim_override_recipient"; bytes32 internal constant _MIGRATOR_V1_PREFIX = "wm_migrator_v1"; address public immutable mToken; @@ -53,15 +53,19 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { } function claimExcess() external returns (uint240 yield_) { - emit ExcessClaim(yield_ = excess()); + emit ExcessClaimed(yield_ = excess()); IMTokenLike(mToken).transfer(vault, yield_); } - function deposit(address destination_, uint256 amount_) external { - emit Transfer(address(0), destination_, amount_); + function deposit(address recipient_, uint256 amount_) external { + _revertIfInsufficientAmount(amount_); + _revertIfInvalidRecipient(recipient_); - _addAmount(destination_, UIntMath.safe240(amount_)); + 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_); } @@ -84,13 +88,15 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { IndexingMath.getPrincipalAmountRoundedDown(rawBalance_, currentIndex_) ); - totalNonEarningSupply -= rawBalance_; + unchecked { + totalNonEarningSupply -= rawBalance_; + } _addTotalEarningSupply(rawBalance_, currentIndex_); } function stopEarningFor(address account_) external { - if (_isApprovedEarner(account_)) revert ApprovedEarner(); + if (_isApprovedEarner(account_)) revert IsApprovedEarner(); (bool isEarning_, , ) = _getBalanceInfo(account_); @@ -107,17 +113,23 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { uint240 amount_ = IndexingMath.getPresentAmountRoundedDown(uint112(rawBalance_), index_); _setBalanceInfo(account_, false, 0, amount_); - totalNonEarningSupply += amount_; + + unchecked { + totalNonEarningSupply += amount_; + } _subtractTotalEarningSupply(amount_, currentIndex_); } - function withdraw(address destination_, uint256 amount_) external { + 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(destination_, amount_); + IMTokenLike(mToken).transfer(recipient_, amount_); } /* ============ View/Pure Functions ============ */ @@ -138,11 +150,17 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { return IMTokenLike(mToken).currentIndex(); } + function isEarning(address account_) external view returns (bool isEarning_) { + (isEarning_, , ) = _getBalanceInfo(account_); + } + function excess() public view returns (uint240 yield_) { uint240 balance_ = uint240(IMTokenLike(mToken).balanceOf(address(this))); uint240 earmarked_ = uint240(totalSupply()) + totalAccruedYield(); - return balance_ > earmarked_ ? balance_ - earmarked_ : 0; + unchecked { + return balance_ > earmarked_ ? balance_ - earmarked_ : 0; + } } function totalAccruedYield() public view returns (uint240 yield_) { @@ -172,19 +190,24 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { function _addNonEarningAmount(address recipient_, uint240 amount_) internal { (, , uint240 rawBalance_) = _getBalanceInfo(recipient_); - _setBalanceInfo(recipient_, false, 0, rawBalance_ + amount_); - totalNonEarningSupply += amount_; + + unchecked { + _setBalanceInfo(recipient_, false, 0, rawBalance_ + amount_); + totalNonEarningSupply += amount_; + } } function _addEarningAmount(address recipient_, uint240 amount_, uint128 currentIndex_) internal { (, , uint240 rawBalance_) = _getBalanceInfo(recipient_); - _setBalanceInfo( - recipient_, - true, - currentIndex_, - rawBalance_ + IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) - ); + unchecked { + _setBalanceInfo( + recipient_, + true, + currentIndex_, + rawBalance_ + IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) + ); + } _addTotalEarningSupply(amount_, currentIndex_); } @@ -199,18 +222,23 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { if (yield_ == 0) return 0; - emit Claim(account_, yield_); - emit Transfer(address(0), account_, yield_); + emit Claimed(account_, yield_); - _setTotalEarningSupply(totalEarningSupply() + yield_, _principalOfTotalEarningSupply); + unchecked { + _setTotalEarningSupply(totalEarningSupply() + yield_, _principalOfTotalEarningSupply); + } - address claimOverrideDestination_ = _getClaimOverrideDestination(account_); + address claimOverrideRecipient_ = _getClaimOverrideRecipient(account_); - if (claimOverrideDestination_ == address(0)) return yield_; + if (claimOverrideRecipient_ == address(0)) { + // NOTE: The `Transfer` event for a non-zero `claimOverrideRecipient_` + emit Transfer(address(0), account_, yield_); + return yield_; + } - // // NOTE: Watch out for a long chain of claim override destinations. - // // TODO: Maybe can be optimized since we know `account_` is an earner and already claimed. - _transfer(account_, claimOverrideDestination_, yield_, currentIndex_); + // 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 { @@ -232,24 +260,31 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { function _subtractNonEarningAmount(address account_, uint240 amount_) internal { (, , uint240 rawBalance_) = _getBalanceInfo(account_); - _setBalanceInfo(account_, false, 0, rawBalance_ - amount_); - totalNonEarningSupply -= amount_; + + if (rawBalance_ < amount_) revert InsufficientBalance(account_, rawBalance_, amount_); + + unchecked { + _setBalanceInfo(account_, false, 0, rawBalance_ - amount_); + totalNonEarningSupply -= amount_; + } } function _subtractEarningAmount(address account_, uint240 amount_, uint128 currentIndex_) internal { (, , uint240 rawBalance_) = _getBalanceInfo(account_); - _setBalanceInfo( - account_, - true, - currentIndex_, - rawBalance_ - IndexingMath.getPrincipalAmountRoundedUp(amount_, currentIndex_) - ); + uint112 principalAmount_ = IndexingMath.getPrincipalAmountRoundedUp(amount_, currentIndex_); - _subtractTotalEarningSupply(amount_, currentIndex_); + if (rawBalance_ < principalAmount_) revert InsufficientBalance(account_, rawBalance_, principalAmount_); + + unchecked { + _setBalanceInfo(account_, true, currentIndex_, rawBalance_ - principalAmount_); + _subtractTotalEarningSupply(amount_, currentIndex_); + } } function _transfer(address sender_, address recipient_, uint240 amount_, uint128 currentIndex_) internal { + _revertIfInvalidRecipient(recipient_); + _claim(sender_, currentIndex_); _claim(recipient_, currentIndex_); @@ -272,17 +307,22 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { } function _addTotalEarningSupply(uint240 amount_, uint128 currentIndex_) internal { - _setTotalEarningSupply( - totalEarningSupply() + amount_, - _principalOfTotalEarningSupply + IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) - ); + unchecked { + _setTotalEarningSupply( + totalEarningSupply() + amount_, + _principalOfTotalEarningSupply + IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) + ); + } } function _subtractTotalEarningSupply(uint240 amount_, uint128 currentIndex_) internal { - _setTotalEarningSupply( - totalEarningSupply() - amount_, - _principalOfTotalEarningSupply - IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) - ); + // TODO: Consider `getPrincipalAmountRoundedUp` . + unchecked { + _setTotalEarningSupply( + totalEarningSupply() - amount_, + _principalOfTotalEarningSupply - IndexingMath.getPrincipalAmountRoundedDown(amount_, currentIndex_) + ); + } } function _setTotalEarningSupply(uint240 amount_, uint112 principalAmount_) internal { @@ -300,7 +340,12 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { uint128 index_, uint128 currentIndex_ ) internal pure returns (uint240) { - return IndexingMath.getPresentAmountRoundedDown(principalAmount_, currentIndex_ - index_); + unchecked { + return + currentIndex_ <= index_ + ? 0 + : IndexingMath.getPresentAmountRoundedDown(principalAmount_, currentIndex_ - index_); + } } function _getBalanceInfo( @@ -314,11 +359,13 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { : (false, uint128(0), uint240(unwrapped_)); } - function _getClaimOverrideDestination(address account_) internal view returns (address) { + function _getClaimOverrideRecipient(address account_) internal view returns (address) { return address( uint160( - uint256(IRegistrarLike(registrar).get(keccak256(abi.encode(_CLAIM_DESTINATION_PREFIX, account_)))) + uint256( + IRegistrarLike(registrar).get(keccak256(abi.encode(_CLAIM_OVERRIDE_RECIPIENT_PREFIX, account_))) + ) ) ); } @@ -340,7 +387,9 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { uint240 totalEarningSupply_ = totalEarningSupply(); - return totalProjectedSupply_ <= totalEarningSupply_ ? 0 : totalProjectedSupply_ - totalEarningSupply_; + unchecked { + return totalProjectedSupply_ <= totalEarningSupply_ ? 0 : totalProjectedSupply_ - totalEarningSupply_; + } } function _isApprovedEarner(address account_) internal view returns (bool) { @@ -348,4 +397,20 @@ contract WrappedM is IWrappedM, Migratable, ERC20Extended { IRegistrarLike(registrar).get(_EARNERS_LIST_IGNORED) != bytes32(0) || IRegistrarLike(registrar).listContains(_EARNERS_LIST, account_); } + + /** + * @dev Reverts if `amount_` is equal to 0. + * @param amount_ Amount of token. + */ + function _revertIfInsufficientAmount(uint256 amount_) internal pure { + if (amount_ == 0) revert InsufficientAmount(amount_); + } + + /** + * @dev Reverts if `recipient_` is address(0). + * @param recipient_ Address of a recipient. + */ + function _revertIfInvalidRecipient(address recipient_) internal pure { + if (recipient_ == address(0)) revert InvalidRecipient(recipient_); + } } diff --git a/src/interfaces/IWrappedM.sol b/src/interfaces/IWrappedMToken.sol similarity index 51% rename from src/interfaces/IWrappedM.sol rename to src/interfaces/IWrappedMToken.sol index 5e91221..4eeb4fd 100644 --- a/src/interfaces/IWrappedM.sol +++ b/src/interfaces/IWrappedMToken.sol @@ -6,23 +6,42 @@ import { IERC20Extended } from "../../lib/common/src/interfaces/IERC20Extended.s import { IMigratable } from "./IMigratable.sol"; -interface IWrappedM is IMigratable, IERC20Extended { +interface IWrappedMToken is IMigratable, IERC20Extended { /* ============ Events ============ */ - event Claim(address indexed account, uint256 yield); + event Claimed(address indexed account, uint256 yield); - event ExcessClaim(uint256 yield); + event ExcessClaimed(uint256 yield); + /** + * @notice Emitted when account starts being an M earner. + * @param account The account that started earning. + */ event StartedEarning(address indexed account); + /** + * @notice Emitted when account stops being an M earner. + * @param account The account that stopped earning. + */ event StoppedEarning(address indexed account); /* ============ Custom Errors ============ */ - error ApprovedEarner(); + /// @notice Emitted when calling `stopEarning` for an account approved as earner by TTG. + error IsApprovedEarner(); + /** + * @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). + */ + error InsufficientBalance(address account, uint256 rawBalance, uint256 amount); + + /// @notice Emitted when calling `startEarning` for an account not approved as earner by TTG. error NotApprovedEarner(); + /// @notice Emitted in constructor if M Token is 0x0. error ZeroMToken(); /* ============ Interactive Functions ============ */ @@ -31,13 +50,13 @@ interface IWrappedM is IMigratable, IERC20Extended { function claimExcess() external returns (uint240 yield); - function deposit(address destination, uint256 amount) external; + function deposit(address recipient, uint256 amount) external; function startEarningFor(address account) external; function stopEarningFor(address account) external; - function withdraw(address destination, uint256 amount) external; + function withdraw(address recipient, uint256 amount) external; /* ============ View/Pure Functions ============ */ @@ -45,6 +64,8 @@ interface IWrappedM is IMigratable, IERC20Extended { function currentIndex() external view returns (uint128 index); + function isEarning(address account) external view returns (bool isEarning); + function excess() external view returns (uint240 yield); function mToken() external view returns (address mToken); diff --git a/test/Test.t.sol b/test/Test.t.sol index ad8b920..49a8eb3 100644 --- a/test/Test.t.sol +++ b/test/Test.t.sol @@ -4,66 +4,20 @@ pragma solidity 0.8.23; import { Test, console2 } from "../lib/forge-std/src/Test.sol"; -import { IWrappedM } from "../src/interfaces/IWrappedM.sol"; +import { IWrappedMToken } from "../src/interfaces/IWrappedMToken.sol"; -import { WrappedM } from "../src/WrappedM.sol"; +import { WrappedMToken } from "../src/WrappedMToken.sol"; import { Proxy } from "../src/Proxy.sol"; -contract MockM { - address public ttgRegistrar; +import { MockM, MockRegistrar } from "./utils/Mocks.sol"; - uint128 public currentIndex; - - mapping(address account => uint256 balance) public balanceOf; - - function transfer(address, uint256) external returns (bool success_) { - return true; - } - - function transferFrom(address, address, uint256) external returns (bool success_) { - return true; - } - - function setBalanceOf(address account_, uint256 balance_) external { - balanceOf[account_] = balance_; - } - - function setCurrentIndex(uint128 currentIndex_) external { - currentIndex = currentIndex_; - } - - function setTtgRegistrar(address ttgRegistrar_) external { - ttgRegistrar = ttgRegistrar_; - } -} - -contract MockRegistrar { - address public vault; - - mapping(bytes32 key => bytes32 value) public get; - - mapping(bytes32 list => mapping(address account => bool contains)) public listContains; - - function set(bytes32 key_, bytes32 value_) external { - get[key_] = value_; - } - - function setListContains(bytes32 list_, address account_, bool contains_) external { - listContains[list_][account_] = contains_; - } - - function setVault(address vault_) external { - vault = vault_; - } -} - -contract WrappedMV2 { +contract WrappedMTokenV2 { function foo() external pure returns (uint256) { return 1; } } -contract WrappedMMigratorV1 { +contract WrappedMTokenMigratorV1 { bytes32 private constant _IMPLEMENTATION_SLOT = bytes32(0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc); @@ -87,7 +41,6 @@ contract Tests is Test { uint56 internal constant _EXP_SCALED_ONE = 1e12; bytes32 internal constant _EARNERS_LIST = "earners"; - bytes32 internal constant _CLAIM_DESTINATION_PREFIX = "wm_claim_destination"; bytes32 internal constant _MIGRATOR_V1_PREFIX = "wm_migrator_v1"; address internal _alice = makeAddr("alice"); @@ -99,8 +52,8 @@ contract Tests is Test { MockM internal _mToken; MockRegistrar internal _registrar; - WrappedM internal _implementation; - IWrappedM internal _wrappedM; + WrappedMToken internal _implementation; + IWrappedMToken internal _wrappedMToken; function setUp() external { _registrar = new MockRegistrar(); @@ -110,353 +63,353 @@ contract Tests is Test { _mToken.setCurrentIndex(_EXP_SCALED_ONE); _mToken.setTtgRegistrar(address(_registrar)); - _implementation = new WrappedM(address(_mToken)); + _implementation = new WrappedMToken(address(_mToken)); - _wrappedM = IWrappedM(address(new Proxy(address(_implementation)))); + _wrappedMToken = IWrappedMToken(address(new Proxy(address(_implementation)))); } function test_story() external { _registrar.setListContains(_EARNERS_LIST, _alice, true); _registrar.setListContains(_EARNERS_LIST, _bob, true); - _wrappedM.startEarningFor(_alice); + _wrappedMToken.startEarningFor(_alice); - _wrappedM.startEarningFor(_bob); + _wrappedMToken.startEarningFor(_bob); vm.prank(_alice); - _wrappedM.deposit(_alice, 100_000000); + _wrappedMToken.deposit(_alice, 100_000000); - _mToken.setBalanceOf(address(_wrappedM), 100_000000); + _mToken.setBalanceOf(address(_wrappedMToken), 100_000000); // Assert Alice (Earner) - assertEq(_wrappedM.balanceOf(_alice), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_alice), 0); + assertEq(_wrappedMToken.balanceOf(_alice), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 100_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 0); - assertEq(_wrappedM.totalSupply(), 100_000000); - assertEq(_wrappedM.totalAccruedYield(), 0); - assertEq(_wrappedM.excess(), 0); + assertEq(_wrappedMToken.totalEarningSupply(), 100_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.totalSupply(), 100_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 0); + assertEq(_wrappedMToken.excess(), 0); vm.prank(_carol); - _wrappedM.deposit(_carol, 100_000000); + _wrappedMToken.deposit(_carol, 100_000000); - _mToken.setBalanceOf(address(_wrappedM), 200_000000); + _mToken.setBalanceOf(address(_wrappedMToken), 200_000000); // Assert Carol (Non-Earner) - assertEq(_wrappedM.balanceOf(_carol), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_carol), 0); + assertEq(_wrappedMToken.balanceOf(_carol), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_carol), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 100_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 100_000000); - assertEq(_wrappedM.totalSupply(), 200_000000); - assertEq(_wrappedM.totalAccruedYield(), 0); - assertEq(_wrappedM.excess(), 0); + assertEq(_wrappedMToken.totalEarningSupply(), 100_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 100_000000); + assertEq(_wrappedMToken.totalSupply(), 200_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 0); + assertEq(_wrappedMToken.excess(), 0); _mToken.setCurrentIndex(2 * _EXP_SCALED_ONE); - _mToken.setBalanceOf(address(_wrappedM), 400_000000); + _mToken.setBalanceOf(address(_wrappedMToken), 400_000000); // Assert Alice (Earner) - assertEq(_wrappedM.balanceOf(_alice), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_alice), 100_000000); + assertEq(_wrappedMToken.balanceOf(_alice), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 100_000000); // Assert Carol (Non-Earner) - assertEq(_wrappedM.balanceOf(_carol), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_carol), 0); + assertEq(_wrappedMToken.balanceOf(_carol), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_carol), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 100_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 100_000000); - assertEq(_wrappedM.totalSupply(), 200_000000); - assertEq(_wrappedM.totalAccruedYield(), 100_000000); - assertEq(_wrappedM.excess(), 100_000000); + assertEq(_wrappedMToken.totalEarningSupply(), 100_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 100_000000); + assertEq(_wrappedMToken.totalSupply(), 200_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 100_000000); + assertEq(_wrappedMToken.excess(), 100_000000); vm.prank(_bob); - _wrappedM.deposit(_bob, 100_000000); + _wrappedMToken.deposit(_bob, 100_000000); - _mToken.setBalanceOf(address(_wrappedM), 500_000000); + _mToken.setBalanceOf(address(_wrappedMToken), 500_000000); // Assert Bob (Earner) - assertEq(_wrappedM.balanceOf(_bob), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_bob), 0); + assertEq(_wrappedMToken.balanceOf(_bob), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_bob), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 200_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 100_000000); - assertEq(_wrappedM.totalSupply(), 300_000000); + assertEq(_wrappedMToken.totalEarningSupply(), 200_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 100_000000); + assertEq(_wrappedMToken.totalSupply(), 300_000000); - assertEq(_wrappedM.totalAccruedYield(), 100_000000); - assertEq(_wrappedM.excess(), 100_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 100_000000); + assertEq(_wrappedMToken.excess(), 100_000000); vm.prank(_dave); - _wrappedM.deposit(_dave, 100_000000); + _wrappedMToken.deposit(_dave, 100_000000); - _mToken.setBalanceOf(address(_wrappedM), 600_000000); + _mToken.setBalanceOf(address(_wrappedMToken), 600_000000); // Assert Dave (Non-Earner) - assertEq(_wrappedM.balanceOf(_dave), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_dave), 0); + assertEq(_wrappedMToken.balanceOf(_dave), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_dave), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 200_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 200_000000); - assertEq(_wrappedM.totalSupply(), 400_000000); - assertEq(_wrappedM.totalAccruedYield(), 100_000000); - assertEq(_wrappedM.excess(), 100_000000); + assertEq(_wrappedMToken.totalEarningSupply(), 200_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 200_000000); + assertEq(_wrappedMToken.totalSupply(), 400_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 100_000000); + assertEq(_wrappedMToken.excess(), 100_000000); - assertEq(_wrappedM.balanceOf(_alice), 100_000000); + assertEq(_wrappedMToken.balanceOf(_alice), 100_000000); - uint256 yield_ = _wrappedM.claimFor(_alice); + uint256 yield_ = _wrappedMToken.claimFor(_alice); assertEq(yield_, 100_000000); // Assert Alice (Earner) - assertEq(_wrappedM.balanceOf(_alice), 200_000000); - assertEq(_wrappedM.accruedYieldOf(_alice), 0); + assertEq(_wrappedMToken.balanceOf(_alice), 200_000000); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 300_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 200_000000); - assertEq(_wrappedM.totalSupply(), 500_000000); - assertEq(_wrappedM.totalAccruedYield(), 0); - assertEq(_wrappedM.excess(), 100_000000); + assertEq(_wrappedMToken.totalEarningSupply(), 300_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 200_000000); + assertEq(_wrappedMToken.totalSupply(), 500_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 0); + assertEq(_wrappedMToken.excess(), 100_000000); _mToken.setCurrentIndex(3 * _EXP_SCALED_ONE); - _mToken.setBalanceOf(address(_wrappedM), 900_000000); // was 600 @ 2.0, so 900 @ 3.0 + _mToken.setBalanceOf(address(_wrappedMToken), 900_000000); // was 600 @ 2.0, so 900 @ 3.0 // Assert Alice (Earner) - assertEq(_wrappedM.balanceOf(_alice), 200_000000); - assertEq(_wrappedM.accruedYieldOf(_alice), 100_000000); + assertEq(_wrappedMToken.balanceOf(_alice), 200_000000); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 100_000000); // Assert Bob (Earner) - assertEq(_wrappedM.balanceOf(_bob), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_bob), 50_000000); + assertEq(_wrappedMToken.balanceOf(_bob), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_bob), 50_000000); // Assert Carol (Non-Earner) - assertEq(_wrappedM.balanceOf(_carol), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_carol), 0); + assertEq(_wrappedMToken.balanceOf(_carol), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_carol), 0); // Assert Dave (Non-Earner) - assertEq(_wrappedM.balanceOf(_dave), 100_000000); - assertEq(_wrappedM.accruedYieldOf(_dave), 0); + assertEq(_wrappedMToken.balanceOf(_dave), 100_000000); + assertEq(_wrappedMToken.accruedYieldOf(_dave), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 300_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 200_000000); - assertEq(_wrappedM.totalSupply(), 500_000000); - assertEq(_wrappedM.totalAccruedYield(), 150_000000); - assertEq(_wrappedM.excess(), 250_000000); + assertEq(_wrappedMToken.totalEarningSupply(), 300_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 200_000000); + assertEq(_wrappedMToken.totalSupply(), 500_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 150_000000); + assertEq(_wrappedMToken.excess(), 250_000000); vm.prank(_alice); - _wrappedM.transfer(_carol, 100_000000); + _wrappedMToken.transfer(_carol, 100_000000); // Assert Alice (Earner) - assertEq(_wrappedM.balanceOf(_alice), 199_999998); - assertEq(_wrappedM.accruedYieldOf(_alice), 0); + assertEq(_wrappedMToken.balanceOf(_alice), 199_999998); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); // Assert Carol (Non-Earner) - assertEq(_wrappedM.balanceOf(_carol), 200_000000); - assertEq(_wrappedM.accruedYieldOf(_carol), 0); + assertEq(_wrappedMToken.balanceOf(_carol), 200_000000); + assertEq(_wrappedMToken.accruedYieldOf(_carol), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 300_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 300_000000); - assertEq(_wrappedM.totalSupply(), 600_000000); - assertEq(_wrappedM.totalAccruedYield(), 50_000001); - assertEq(_wrappedM.excess(), 249_999999); + 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); vm.prank(_dave); - _wrappedM.transfer(_bob, 50_000000); + _wrappedMToken.transfer(_bob, 50_000000); // Assert Bob (Earner) - assertEq(_wrappedM.balanceOf(_bob), 199_999998); - assertEq(_wrappedM.accruedYieldOf(_bob), 0); + assertEq(_wrappedMToken.balanceOf(_bob), 199_999998); + assertEq(_wrappedMToken.accruedYieldOf(_bob), 0); // Assert Dave (Non-Earner) - assertEq(_wrappedM.balanceOf(_dave), 50_000000); - assertEq(_wrappedM.accruedYieldOf(_dave), 0); + assertEq(_wrappedMToken.balanceOf(_dave), 50_000000); + assertEq(_wrappedMToken.accruedYieldOf(_dave), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 400_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 250_000000); - assertEq(_wrappedM.totalSupply(), 650_000000); - assertEq(_wrappedM.totalAccruedYield(), 0); - assertEq(_wrappedM.excess(), 250_000000); + assertEq(_wrappedMToken.totalEarningSupply(), 400_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 250_000000); + assertEq(_wrappedMToken.totalSupply(), 650_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 0); + assertEq(_wrappedMToken.excess(), 250_000000); _mToken.setCurrentIndex(4 * _EXP_SCALED_ONE); - _mToken.setBalanceOf(address(_wrappedM), 1_200_000000); // was 900 @ 3.0, so 1200 @ 4.0 + _mToken.setBalanceOf(address(_wrappedMToken), 1_200_000000); // was 900 @ 3.0, so 1200 @ 4.0 // Assert Alice (Earner) - assertEq(_wrappedM.balanceOf(_alice), 199_999998); - assertEq(_wrappedM.accruedYieldOf(_alice), 66_666666); + assertEq(_wrappedMToken.balanceOf(_alice), 199_999998); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 66_666666); // Assert Bob (Earner) - assertEq(_wrappedM.balanceOf(_bob), 199_999998); - assertEq(_wrappedM.accruedYieldOf(_bob), 66_666666); + assertEq(_wrappedMToken.balanceOf(_bob), 199_999998); + assertEq(_wrappedMToken.accruedYieldOf(_bob), 66_666666); // Assert Carol (Non-Earner) - assertEq(_wrappedM.balanceOf(_carol), 200_000000); - assertEq(_wrappedM.accruedYieldOf(_carol), 0); + assertEq(_wrappedMToken.balanceOf(_carol), 200_000000); + assertEq(_wrappedMToken.accruedYieldOf(_carol), 0); // Assert Dave (Non-Earner) - assertEq(_wrappedM.balanceOf(_dave), 50_000000); - assertEq(_wrappedM.accruedYieldOf(_dave), 0); + assertEq(_wrappedMToken.balanceOf(_dave), 50_000000); + assertEq(_wrappedMToken.accruedYieldOf(_dave), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 400_000000); - assertEq(_wrappedM.totalNonEarningSupply(), 250_000000); - assertEq(_wrappedM.totalSupply(), 650_000000); - assertEq(_wrappedM.totalAccruedYield(), 133_333332); - assertEq(_wrappedM.excess(), 416_666668); + assertEq(_wrappedMToken.totalEarningSupply(), 400_000000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 250_000000); + assertEq(_wrappedMToken.totalSupply(), 650_000000); + assertEq(_wrappedMToken.totalAccruedYield(), 133_333332); + assertEq(_wrappedMToken.excess(), 416_666668); _registrar.setListContains(_EARNERS_LIST, _alice, false); - _wrappedM.stopEarningFor(_alice); + _wrappedMToken.stopEarningFor(_alice); // Assert Alice (Non-Earner) - assertEq(_wrappedM.balanceOf(_alice), 266_666664); - assertEq(_wrappedM.accruedYieldOf(_alice), 0); + assertEq(_wrappedMToken.balanceOf(_alice), 266_666664); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 200_000002); - assertEq(_wrappedM.totalNonEarningSupply(), 516_666664); - assertEq(_wrappedM.totalSupply(), 716_666666); - assertEq(_wrappedM.totalAccruedYield(), 66_666666); - assertEq(_wrappedM.excess(), 416_666668); + assertEq(_wrappedMToken.totalEarningSupply(), 200_000002); + assertEq(_wrappedMToken.totalNonEarningSupply(), 516_666664); + assertEq(_wrappedMToken.totalSupply(), 716_666666); + assertEq(_wrappedMToken.totalAccruedYield(), 66_666666); + assertEq(_wrappedMToken.excess(), 416_666668); _registrar.setListContains(_EARNERS_LIST, _carol, true); - _wrappedM.startEarningFor(_carol); + _wrappedMToken.startEarningFor(_carol); // Assert Carol (Earner) - assertEq(_wrappedM.balanceOf(_carol), 200_000000); - assertEq(_wrappedM.accruedYieldOf(_carol), 0); + assertEq(_wrappedMToken.balanceOf(_carol), 200_000000); + assertEq(_wrappedMToken.accruedYieldOf(_carol), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 400_000002); - assertEq(_wrappedM.totalNonEarningSupply(), 316_666664); - assertEq(_wrappedM.totalSupply(), 716_666666); - assertEq(_wrappedM.totalAccruedYield(), 66_666666); - assertEq(_wrappedM.excess(), 416_666668); + assertEq(_wrappedMToken.totalEarningSupply(), 400_000002); + assertEq(_wrappedMToken.totalNonEarningSupply(), 316_666664); + assertEq(_wrappedMToken.totalSupply(), 716_666666); + assertEq(_wrappedMToken.totalAccruedYield(), 66_666666); + assertEq(_wrappedMToken.excess(), 416_666668); _mToken.setCurrentIndex(5 * _EXP_SCALED_ONE); - _mToken.setBalanceOf(address(_wrappedM), 1_500_000000); // was 1200 @ 4.0, so 1500 @ 5.0 + _mToken.setBalanceOf(address(_wrappedMToken), 1_500_000000); // was 1200 @ 4.0, so 1500 @ 5.0 // Assert Alice (Non-Earner) - assertEq(_wrappedM.balanceOf(_alice), 266_666664); - assertEq(_wrappedM.accruedYieldOf(_alice), 0); + assertEq(_wrappedMToken.balanceOf(_alice), 266_666664); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); // Assert Bob (Earner) - assertEq(_wrappedM.balanceOf(_bob), 199_999998); - assertEq(_wrappedM.accruedYieldOf(_bob), 133_333332); + assertEq(_wrappedMToken.balanceOf(_bob), 199_999998); + assertEq(_wrappedMToken.accruedYieldOf(_bob), 133_333332); // Assert Carol (Earner) - assertEq(_wrappedM.balanceOf(_carol), 200_000000); - assertEq(_wrappedM.accruedYieldOf(_carol), 50_000000); + assertEq(_wrappedMToken.balanceOf(_carol), 200_000000); + assertEq(_wrappedMToken.accruedYieldOf(_carol), 50_000000); // Assert Dave (Non-Earner) - assertEq(_wrappedM.balanceOf(_dave), 50_000000); - assertEq(_wrappedM.accruedYieldOf(_dave), 0); + assertEq(_wrappedMToken.balanceOf(_dave), 50_000000); + assertEq(_wrappedMToken.accruedYieldOf(_dave), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 400_000002); - assertEq(_wrappedM.totalNonEarningSupply(), 316_666664); - assertEq(_wrappedM.totalSupply(), 716_666666); - assertEq(_wrappedM.totalAccruedYield(), 183_333333); - assertEq(_wrappedM.excess(), 600_000001); + assertEq(_wrappedMToken.totalEarningSupply(), 400_000002); + assertEq(_wrappedMToken.totalNonEarningSupply(), 316_666664); + assertEq(_wrappedMToken.totalSupply(), 716_666666); + assertEq(_wrappedMToken.totalAccruedYield(), 183_333333); + assertEq(_wrappedMToken.excess(), 600_000001); vm.prank(_alice); - _wrappedM.withdraw(_alice, 266_666664); + _wrappedMToken.withdraw(_alice, 266_666664); - _mToken.setBalanceOf(address(_wrappedM), 1_233_333336); + _mToken.setBalanceOf(address(_wrappedMToken), 1_233_333336); // Assert Alice (Non-Earner) - assertEq(_wrappedM.balanceOf(_alice), 0); - assertEq(_wrappedM.accruedYieldOf(_alice), 0); + assertEq(_wrappedMToken.balanceOf(_alice), 0); + assertEq(_wrappedMToken.accruedYieldOf(_alice), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 400_000002); - assertEq(_wrappedM.totalNonEarningSupply(), 50_000000); - assertEq(_wrappedM.totalSupply(), 450_000002); - assertEq(_wrappedM.totalAccruedYield(), 183_333333); - assertEq(_wrappedM.excess(), 600_000001); + assertEq(_wrappedMToken.totalEarningSupply(), 400_000002); + assertEq(_wrappedMToken.totalNonEarningSupply(), 50_000000); + assertEq(_wrappedMToken.totalSupply(), 450_000002); + assertEq(_wrappedMToken.totalAccruedYield(), 183_333333); + assertEq(_wrappedMToken.excess(), 600_000001); vm.prank(_bob); - _wrappedM.withdraw(_bob, 333_333330); + _wrappedMToken.withdraw(_bob, 333_333330); - _mToken.setBalanceOf(address(_wrappedM), 900_000006); + _mToken.setBalanceOf(address(_wrappedMToken), 900_000006); // Assert Bob (Earner) - assertEq(_wrappedM.balanceOf(_bob), 0); - assertEq(_wrappedM.accruedYieldOf(_bob), 0); + assertEq(_wrappedMToken.balanceOf(_bob), 0); + assertEq(_wrappedMToken.accruedYieldOf(_bob), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 200_000004); - assertEq(_wrappedM.totalNonEarningSupply(), 50_000000); - assertEq(_wrappedM.totalSupply(), 250_000004); - assertEq(_wrappedM.totalAccruedYield(), 50_000001); - assertEq(_wrappedM.excess(), 600_000001); + assertEq(_wrappedMToken.totalEarningSupply(), 200_000004); + assertEq(_wrappedMToken.totalNonEarningSupply(), 50_000000); + assertEq(_wrappedMToken.totalSupply(), 250_000004); + assertEq(_wrappedMToken.totalAccruedYield(), 50_000001); + assertEq(_wrappedMToken.excess(), 600_000001); vm.prank(_carol); - _wrappedM.withdraw(_carol, 250_000000); + _wrappedMToken.withdraw(_carol, 250_000000); - _mToken.setBalanceOf(address(_wrappedM), 650_000006); + _mToken.setBalanceOf(address(_wrappedMToken), 650_000006); // Assert Carol (Earner) - assertEq(_wrappedM.balanceOf(_carol), 0); - assertEq(_wrappedM.accruedYieldOf(_carol), 0); + assertEq(_wrappedMToken.balanceOf(_carol), 0); + assertEq(_wrappedMToken.accruedYieldOf(_carol), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 4); - assertEq(_wrappedM.totalNonEarningSupply(), 50_000000); - assertEq(_wrappedM.totalSupply(), 50_000004); - assertEq(_wrappedM.totalAccruedYield(), 1); - assertEq(_wrappedM.excess(), 600_000001); + assertEq(_wrappedMToken.totalEarningSupply(), 4); + assertEq(_wrappedMToken.totalNonEarningSupply(), 50_000000); + assertEq(_wrappedMToken.totalSupply(), 50_000004); + assertEq(_wrappedMToken.totalAccruedYield(), 1); + assertEq(_wrappedMToken.excess(), 600_000001); vm.prank(_dave); - _wrappedM.withdraw(_dave, 50_000000); + _wrappedMToken.withdraw(_dave, 50_000000); - _mToken.setBalanceOf(address(_wrappedM), 600_000006); + _mToken.setBalanceOf(address(_wrappedMToken), 600_000006); // Assert Dave (Non-Earner) - assertEq(_wrappedM.balanceOf(_dave), 0); - assertEq(_wrappedM.accruedYieldOf(_dave), 0); + assertEq(_wrappedMToken.balanceOf(_dave), 0); + assertEq(_wrappedMToken.accruedYieldOf(_dave), 0); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 4); - assertEq(_wrappedM.totalNonEarningSupply(), 0); - assertEq(_wrappedM.totalSupply(), 4); - assertEq(_wrappedM.totalAccruedYield(), 1); - assertEq(_wrappedM.excess(), 600_000001); + assertEq(_wrappedMToken.totalEarningSupply(), 4); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.totalSupply(), 4); + assertEq(_wrappedMToken.totalAccruedYield(), 1); + assertEq(_wrappedMToken.excess(), 600_000001); - _wrappedM.claimExcess(); + _wrappedMToken.claimExcess(); - _mToken.setBalanceOf(address(_wrappedM), 11); + _mToken.setBalanceOf(address(_wrappedMToken), 11); // Assert Globals - assertEq(_wrappedM.totalEarningSupply(), 4); - assertEq(_wrappedM.totalNonEarningSupply(), 0); - assertEq(_wrappedM.totalSupply(), 4); - assertEq(_wrappedM.totalAccruedYield(), 1); - assertEq(_wrappedM.excess(), 6); + assertEq(_wrappedMToken.totalEarningSupply(), 4); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.totalSupply(), 4); + assertEq(_wrappedMToken.totalAccruedYield(), 1); + assertEq(_wrappedMToken.excess(), 6); } function test_migration() external { - WrappedMV2 implementationV2_ = new WrappedMV2(); - address migrator_ = address(new WrappedMMigratorV1(address(implementationV2_))); + WrappedMTokenV2 implementationV2_ = new WrappedMTokenV2(); + address migrator_ = address(new WrappedMTokenMigratorV1(address(implementationV2_))); _registrar.set( - keccak256(abi.encode(_MIGRATOR_V1_PREFIX, address(_wrappedM))), + keccak256(abi.encode(_MIGRATOR_V1_PREFIX, address(_wrappedMToken))), bytes32(uint256(uint160(migrator_))) ); vm.expectRevert(); - WrappedMV2(address(_wrappedM)).foo(); + WrappedMTokenV2(address(_wrappedMToken)).foo(); - _wrappedM.migrate(); + _wrappedMToken.migrate(); - assertEq(WrappedMV2(address(_wrappedM)).foo(), 1); + assertEq(WrappedMTokenV2(address(_wrappedMToken)).foo(), 1); } } diff --git a/test/WrappedMToken.t.sol b/test/WrappedMToken.t.sol new file mode 100644 index 0000000..7d91270 --- /dev/null +++ b/test/WrappedMToken.t.sol @@ -0,0 +1,518 @@ +// SPDX-License-Identifier: UNLICENSED + +pragma solidity 0.8.23; + +import { Test, console2 } from "../lib/forge-std/src/Test.sol"; +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"; + +import { Proxy } from "../src/Proxy.sol"; + +import { MockM, MockRegistrar } from "./utils/Mocks.sol"; +import { WrappedMTokenHarness } from "./utils/WrappedMTokenHarness.sol"; + +// NOTE: Due to `_indexOfTotalEarningSupply` a helper to overestimate `totalEarningSupply()`, there is little reason +// to programmatically expect its value rather than ensuring `totalEarningSupply()` is acceptable. + +contract WrappedMTokenTests is Test { + uint56 internal constant _EXP_SCALED_ONE = 1e12; + + bytes32 internal constant _EARNERS_LIST = "earners"; + bytes32 internal constant _CLAIM_DESTINATION_PREFIX = "wm_claim_destination"; + bytes32 internal constant _MIGRATOR_V1_PREFIX = "wm_migrator_v1"; + + address internal _alice = makeAddr("alice"); + address internal _bob = makeAddr("bob"); + address internal _charlie = makeAddr("charlie"); + address internal _david = makeAddr("david"); + + address[] internal _accounts = [_alice, _bob, _charlie, _david]; + + address internal _vault = makeAddr("vault"); + + uint128 internal _expectedCurrentIndex; + + MockM internal _mToken; + MockRegistrar internal _registrar; + WrappedMTokenHarness internal _implementation; + WrappedMTokenHarness internal _wrappedMToken; + + function setUp() external { + _registrar = new MockRegistrar(); + _registrar.setVault(_vault); + + _mToken = new MockM(); + _mToken.setCurrentIndex(_EXP_SCALED_ONE); + _mToken.setTtgRegistrar(address(_registrar)); + + _implementation = new WrappedMTokenHarness(address(_mToken)); + + _wrappedMToken = WrappedMTokenHarness(address(new Proxy(address(_implementation)))); + + _mToken.setCurrentIndex(_expectedCurrentIndex = 1_100000068703); + } + + /* ============ constructor ============ */ + + function test_constructor() external view { + assertEq(_wrappedMToken.implementation(), address(_implementation)); + assertEq(_wrappedMToken.mToken(), address(_mToken)); + assertEq(_wrappedMToken.registrar(), address(_registrar)); + assertEq(_wrappedMToken.vault(), _vault); + } + + function test_constructor_zeroMToken() external { + vm.expectRevert(IWrappedMToken.ZeroMToken.selector); + new WrappedMTokenHarness(address(0)); + } + + /* ============ deposit ============ */ + + function test_deposit_insufficientAmount() external { + vm.expectRevert(abi.encodeWithSelector(IERC20Extended.InsufficientAmount.selector, 0)); + + _wrappedMToken.deposit(_alice, 0); + } + + function test_deposit_invalidRecipient() external { + vm.expectRevert(abi.encodeWithSelector(IERC20Extended.InvalidRecipient.selector, address(0))); + + _wrappedMToken.deposit(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); + } + + function test_deposit_toNonEarner() external { + vm.prank(_alice); + _wrappedMToken.deposit(_alice, 1_000); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 1_000); + assertEq(_wrappedMToken.totalNonEarningSupply(), 1_000); + assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.indexOfTotalEarningSupply(), 0); + } + + function test_deposit_toEarner() external { + _wrappedMToken.setIsEarningOf(_alice, true); + + vm.prank(_alice); + _wrappedMToken.deposit(_alice, 999); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 908); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 908); + assertEq(_wrappedMToken.totalEarningSupply(), 999); + + vm.prank(_alice); + _wrappedMToken.deposit(_alice, 1); + + // No change due to principal round down on deposit. + assertEq(_wrappedMToken.internalBalanceOf(_alice), 908); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 908); + assertEq(_wrappedMToken.totalEarningSupply(), 1000); + + vm.prank(_alice); + _wrappedMToken.deposit(_alice, 2); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 909); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 909); + assertEq(_wrappedMToken.totalEarningSupply(), 1002); + } + + /* ============ withdraw ============ */ + + function test_withdraw_insufficientAmount() external { + vm.expectRevert(abi.encodeWithSelector(IERC20Extended.InsufficientAmount.selector, 0)); + + _wrappedMToken.withdraw(_alice, 0); + } + + function test_withdraw_insufficientBalance_fromNonEarner() external { + _wrappedMToken.setRawBalanceOf(_alice, 999); + + vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); + vm.prank(_alice); + _wrappedMToken.withdraw(_alice, 1_000); + } + + function test_withdraw_insufficientBalance_fromEarner() external { + _wrappedMToken.setIsEarningOf(_alice, true); + _wrappedMToken.setRawBalanceOf(_alice, 908); + + vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 908, 910)); + vm.prank(_alice); + _wrappedMToken.withdraw(_alice, 1_000); + } + + function test_withdraw_fromNonEarner() external { + _wrappedMToken.setTotalNonEarningSupply(1_000); + + _wrappedMToken.setRawBalanceOf(_alice, 1_000); + + vm.prank(_alice); + _wrappedMToken.withdraw(_alice, 500); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 500); + assertEq(_wrappedMToken.totalNonEarningSupply(), 500); + assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.indexOfTotalEarningSupply(), 0); + + vm.prank(_alice); + _wrappedMToken.withdraw(_alice, 500); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 0); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.indexOfTotalEarningSupply(), 0); + } + + function test_withdraw_fromEarner() external { + _wrappedMToken.setPrincipalOfTotalEarningSupply(909); + _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + + _wrappedMToken.setIsEarningOf(_alice, true); + _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); + _wrappedMToken.setRawBalanceOf(_alice, 909); + + vm.prank(_alice); + _wrappedMToken.withdraw(_alice, 1); + + // Change due to principal round up on withdraw. + assertEq(_wrappedMToken.internalBalanceOf(_alice), 908); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningSupply(), 999); + + vm.prank(_alice); + _wrappedMToken.withdraw(_alice, 998); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 0); + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningSupply(), 1); + } + + /* ============ transfer ============ */ + function test_transfer_invalidRecipient() external { + _wrappedMToken.setRawBalanceOf(_alice, 1_000); + + vm.expectRevert(abi.encodeWithSelector(IERC20Extended.InvalidRecipient.selector, address(0))); + + vm.prank(_alice); + _wrappedMToken.transfer(address(0), 1_000); + } + + function test_transfer_insufficientBalance_fromNonEarner_toNonEarner() external { + _wrappedMToken.setRawBalanceOf(_alice, 999); + + vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 999, 1_000)); + vm.prank(_alice); + _wrappedMToken.transfer(_bob, 1_000); + } + + function test_transfer_insufficientBalance_fromEarner_toNonEarner() external { + _wrappedMToken.setIsEarningOf(_alice, true); + _wrappedMToken.setRawBalanceOf(_alice, 908); + + vm.expectRevert(abi.encodeWithSelector(IWrappedMToken.InsufficientBalance.selector, _alice, 908, 910)); + vm.prank(_alice); + _wrappedMToken.transfer(_bob, 1_000); + } + + function test_transfer_fromNonEarner_toNonEarner() external { + _wrappedMToken.setTotalNonEarningSupply(1_500); + + _wrappedMToken.setRawBalanceOf(_alice, 1_000); + _wrappedMToken.setRawBalanceOf(_bob, 500); + + vm.prank(_alice); + _wrappedMToken.transfer(_bob, 500); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 500); + + assertEq(_wrappedMToken.internalBalanceOf(_bob), 1_000); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 1_500); + assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.indexOfTotalEarningSupply(), 0); + } + + function testFuzz_transfer_fromNonEarner_toNonEarner( + uint256 supply_, + uint256 aliceBalance_, + uint256 transferAmount_ + ) external { + supply_ = bound(supply_, 1, type(uint112).max); + aliceBalance_ = bound(aliceBalance_, 1, supply_); + transferAmount_ = bound(transferAmount_, 1, aliceBalance_); + uint256 bobBalance = supply_ - aliceBalance_; + + _wrappedMToken.setTotalNonEarningSupply(supply_); + + _wrappedMToken.setRawBalanceOf(_alice, aliceBalance_); + _wrappedMToken.setRawBalanceOf(_bob, bobBalance); + + vm.prank(_alice); + _wrappedMToken.transfer(_bob, transferAmount_); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), aliceBalance_ - transferAmount_); + assertEq(_wrappedMToken.internalBalanceOf(_bob), bobBalance + transferAmount_); + + assertEq(_wrappedMToken.totalNonEarningSupply(), supply_); + assertEq(_wrappedMToken.principalOfTotalEarningSupply(), 0); + assertEq(_wrappedMToken.indexOfTotalEarningSupply(), 0); + } + + function test_transfer_fromEarner_toNonEarner() external { + _wrappedMToken.setPrincipalOfTotalEarningSupply(909); + _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setTotalNonEarningSupply(500); + + _wrappedMToken.setIsEarningOf(_alice, true); + _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); + _wrappedMToken.setRawBalanceOf(_alice, 909); + + _wrappedMToken.setRawBalanceOf(_bob, 500); + + vm.prank(_alice); + _wrappedMToken.transfer(_bob, 500); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 454); + + assertEq(_wrappedMToken.internalBalanceOf(_bob), 1_000); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 1_000); + assertEq(_wrappedMToken.totalEarningSupply(), 500); + + vm.prank(_alice); + _wrappedMToken.transfer(_bob, 1); + + // Change due to principal round up on burn. + assertEq(_wrappedMToken.internalBalanceOf(_alice), 453); + + assertEq(_wrappedMToken.internalBalanceOf(_bob), 1_001); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 1_001); + assertEq(_wrappedMToken.totalEarningSupply(), 499); + } + + function test_transfer_fromNonEarner_toEarner() external { + _wrappedMToken.setPrincipalOfTotalEarningSupply(455); + _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + _wrappedMToken.setTotalNonEarningSupply(1_000); + + _wrappedMToken.setRawBalanceOf(_alice, 1_000); + + _wrappedMToken.setIsEarningOf(_bob, true); + _wrappedMToken.setIndexOf(_bob, _expectedCurrentIndex); + _wrappedMToken.setRawBalanceOf(_bob, 455); + + vm.prank(_alice); + _wrappedMToken.transfer(_bob, 500); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 500); + + assertEq(_wrappedMToken.internalBalanceOf(_bob), 909); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 500); + assertEq(_wrappedMToken.totalEarningSupply(), 1_001); + } + + function test_transfer_fromEarner_toEarner() external { + _wrappedMToken.setPrincipalOfTotalEarningSupply(1_364); + _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + + _wrappedMToken.setIsEarningOf(_alice, true); + _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); + _wrappedMToken.setRawBalanceOf(_alice, 909); + + _wrappedMToken.setIsEarningOf(_bob, true); + _wrappedMToken.setIndexOf(_bob, _expectedCurrentIndex); + _wrappedMToken.setRawBalanceOf(_bob, 454); + + vm.prank(_alice); + _wrappedMToken.transfer(_bob, 500); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 454); + + assertEq(_wrappedMToken.internalBalanceOf(_bob), 908); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningSupply(), 1_501); + } + + /* ============ startEarningFor ============ */ + function test_startEarningFor_notApprovedEarner() external { + vm.expectRevert(IWrappedMToken.NotApprovedEarner.selector); + _wrappedMToken.startEarningFor(_alice); + } + + function test_startEarningFor() external { + _wrappedMToken.setTotalNonEarningSupply(1_000); + + _wrappedMToken.setRawBalanceOf(_alice, 1_000); + + _registrar.setListContains(_EARNERS_LIST, _alice, true); + + vm.expectEmit(); + emit IWrappedMToken.StartedEarning(_alice); + + _wrappedMToken.startEarningFor(_alice); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 909); + assertEq(_wrappedMToken.isEarning(_alice), true); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 0); + assertEq(_wrappedMToken.totalEarningSupply(), 1_000); + } + + function test_startEarning_overflow() external { + uint256 aliceBalance_ = uint256(type(uint112).max) + 20; + + _mToken.setCurrentIndex(_expectedCurrentIndex = _EXP_SCALED_ONE); + + _wrappedMToken.setTotalNonEarningSupply(aliceBalance_); + _wrappedMToken.setRawBalanceOf(_alice, aliceBalance_); + + _registrar.setListContains(_EARNERS_LIST, _alice, true); + + vm.expectRevert(UIntMath.InvalidUInt112.selector); + _wrappedMToken.startEarningFor(_alice); + } + + /* ============ stopEarningFor ============ */ + function test_stopEarningForAccount_isApprovedEarner() external { + _registrar.setListContains(_EARNERS_LIST, _alice, true); + + vm.expectRevert(IWrappedMToken.IsApprovedEarner.selector); + _wrappedMToken.stopEarningFor(_alice); + } + + function test_stopEarningFor() external { + _wrappedMToken.setPrincipalOfTotalEarningSupply(909); + _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + + _wrappedMToken.setIsEarningOf(_alice, true); + _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); + _wrappedMToken.setRawBalanceOf(_alice, 909); + + _registrar.setListContains(_EARNERS_LIST, _alice, false); + + vm.expectEmit(); + emit IWrappedMToken.StoppedEarning(_alice); + + _wrappedMToken.stopEarningFor(_alice); + + assertEq(_wrappedMToken.internalBalanceOf(_alice), 999); + assertEq(_wrappedMToken.isEarning(_alice), false); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 999); + assertEq(_wrappedMToken.totalEarningSupply(), 1); // TODO: Fix? + } + + /* ============ balanceOf ============ */ + function test_balanceOf_nonEarner() external { + _wrappedMToken.setRawBalanceOf(_alice, 500); + + assertEq(_wrappedMToken.balanceOf(_alice), 500); + + _wrappedMToken.setRawBalanceOf(_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); + + assertEq(_wrappedMToken.balanceOf(_alice), 454); + + _wrappedMToken.setRawBalanceOf(_alice, 909); + + assertEq(_wrappedMToken.balanceOf(_alice), 909); + + _wrappedMToken.setIndexOf(_alice, _expectedCurrentIndex); + + assertEq(_wrappedMToken.balanceOf(_alice), 999); + } + + /* ============ totalNonEarningSupply ============ */ + function test_totalNonEarningSupply() external { + _wrappedMToken.setTotalNonEarningSupply(500); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 500); + + _wrappedMToken.setTotalNonEarningSupply(1_000); + + assertEq(_wrappedMToken.totalNonEarningSupply(), 1_000); + } + + function test_totalEarningSupply() external { + // TODO: more variations + _wrappedMToken.setPrincipalOfTotalEarningSupply(909); + _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + + assertEq(_wrappedMToken.totalEarningSupply(), 1_000); + } + + /* ============ totalSupply ============ */ + function test_totalSupply_onlyTotalNonEarningSupply() external { + _wrappedMToken.setTotalNonEarningSupply(500); + + assertEq(_wrappedMToken.totalSupply(), 500); + + _wrappedMToken.setTotalNonEarningSupply(1_000); + + assertEq(_wrappedMToken.totalSupply(), 1_000); + } + + function test_totalSupply_onlyTotalEarningSupply() external { + // TODO: more variations + _wrappedMToken.setPrincipalOfTotalEarningSupply(909); + _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + + assertEq(_wrappedMToken.totalSupply(), 1_000); + } + + function test_totalSupply() external { + // TODO: more variations + _wrappedMToken.setPrincipalOfTotalEarningSupply(909); + _wrappedMToken.setIndexOfTotalEarningSupply(_expectedCurrentIndex); + + _wrappedMToken.setTotalNonEarningSupply(500); + + assertEq(_wrappedMToken.totalSupply(), 1_500); + + _wrappedMToken.setTotalNonEarningSupply(1_000); + + assertEq(_wrappedMToken.totalSupply(), 2_000); + } + + /* ============ utils ============ */ + function _getPrincipalAmountRoundedDown(uint240 presentAmount_, uint128 index_) internal pure returns (uint112) { + return IndexingMath.divide240By128Down(presentAmount_, index_); + } + + function _getPrincipalAmountRoundedUp(uint240 presentAmount_, uint128 index_) internal pure returns (uint112) { + return IndexingMath.divide240By128Up(presentAmount_, index_); + } + + function _getPresentAmountRoundedDown(uint112 principalAmount_, uint128 index_) internal pure returns (uint240) { + return IndexingMath.multiply112By128Down(principalAmount_, index_); + } + + function _getPresentAmountRoundedUp(uint112 principalAmount_, uint128 index_) internal pure returns (uint240) { + return IndexingMath.multiply112By128Up(principalAmount_, index_); + } +} diff --git a/test/utils/Mocks.sol b/test/utils/Mocks.sol new file mode 100644 index 0000000..87dda4d --- /dev/null +++ b/test/utils/Mocks.sol @@ -0,0 +1,51 @@ +// SPDX-License-Identifier: GPL-3.0 + +pragma solidity 0.8.23; + +contract MockM { + address public ttgRegistrar; + + uint128 public currentIndex; + + mapping(address account => uint256 balance) public balanceOf; + + function transfer(address, uint256) external returns (bool success_) { + return true; + } + + function transferFrom(address, address, uint256) external returns (bool success_) { + return true; + } + + function setBalanceOf(address account_, uint256 balance_) external { + balanceOf[account_] = balance_; + } + + function setCurrentIndex(uint128 currentIndex_) external { + currentIndex = currentIndex_; + } + + function setTtgRegistrar(address ttgRegistrar_) external { + ttgRegistrar = ttgRegistrar_; + } +} + +contract MockRegistrar { + address public vault; + + mapping(bytes32 key => bytes32 value) public get; + + mapping(bytes32 list => mapping(address account => bool contains)) public listContains; + + function set(bytes32 key_, bytes32 value_) external { + get[key_] = value_; + } + + function setListContains(bytes32 list_, address account_, bool contains_) external { + listContains[list_][account_] = contains_; + } + + function setVault(address vault_) external { + vault = vault_; + } +} diff --git a/test/utils/WrappedMTokenHarness.sol b/test/utils/WrappedMTokenHarness.sol new file mode 100644 index 0000000..a31c9b0 --- /dev/null +++ b/test/utils/WrappedMTokenHarness.sol @@ -0,0 +1,48 @@ +// SPDX-License-Identifier: UNLICENSED + +pragma solidity 0.8.23; + +import { WrappedMToken } from "../../src/WrappedMToken.sol"; + +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_); + } + + function setIndexOf(address account_, uint256 index_) external { + (bool isEarning_, , uint240 rawBalance_) = _getBalanceInfo(account_); + _setBalanceInfo(account_, isEarning_, uint128(index_), rawBalance_); + } + + function setRawBalanceOf(address account_, uint256 rawBalance_) external { + (bool isEarning_, uint128 index_, ) = _getBalanceInfo(account_); + _setBalanceInfo(account_, isEarning_, index_, uint240(rawBalance_)); + } + + function setTotalNonEarningSupply(uint256 totalNonEarningSupply_) external { + totalNonEarningSupply = uint240(totalNonEarningSupply_); + } + + function setPrincipalOfTotalEarningSupply(uint256 principalOfTotalEarningSupply_) external { + _principalOfTotalEarningSupply = uint112(principalOfTotalEarningSupply_); + } + + function setIndexOfTotalEarningSupply(uint256 indexOfTotalEarningSupply_) external { + _indexOfTotalEarningSupply = uint128(indexOfTotalEarningSupply_); + } + + function internalBalanceOf(address account_) external view returns (uint256 balance_) { + (, , balance_) = _getBalanceInfo(account_); + } + + function principalOfTotalEarningSupply() external view returns (uint256 principalOfTotalEarningSupply_) { + principalOfTotalEarningSupply_ = _principalOfTotalEarningSupply; + } + + function indexOfTotalEarningSupply() external view returns (uint256 indexOfTotalEarningSupply_) { + indexOfTotalEarningSupply_ = _indexOfTotalEarningSupply; + } +} From d2c1aa4f5143cbbdf49f1f17ca3290c7fe37689e Mon Sep 17 00:00:00 2001 From: Michael De Luca Date: Tue, 2 Jul 2024 12:19:24 -0400 Subject: [PATCH 2/3] fix: PR Review --- src/interfaces/IWrappedMToken.sol | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/interfaces/IWrappedMToken.sol b/src/interfaces/IWrappedMToken.sol index 4eeb4fd..0387af2 100644 --- a/src/interfaces/IWrappedMToken.sol +++ b/src/interfaces/IWrappedMToken.sol @@ -14,13 +14,13 @@ interface IWrappedMToken is IMigratable, IERC20Extended { event ExcessClaimed(uint256 yield); /** - * @notice Emitted when account starts being an M earner. + * @notice Emitted when account starts being an wM earner. * @param account The account that started earning. */ event StartedEarning(address indexed account); /** - * @notice Emitted when account stops being an M earner. + * @notice Emitted when account stops being an wM earner. * @param account The account that stopped earning. */ event StoppedEarning(address indexed account); From 02e3e44850252d8ea1141a7400e4e1f478d5f3d5 Mon Sep 17 00:00:00 2001 From: toninorair Date: Tue, 2 Jul 2024 12:53:57 -0400 Subject: [PATCH 3/3] refactor: Rename methods (#17) * Add renames * refactor: Auto `_setBalanceInfo` --------- Co-authored-by: Michael De Luca <35537333+deluca-mike@users.noreply.github.com> --- src/WrappedMToken.sol | 281 +++++++++++++--------------- src/interfaces/IWrappedMToken.sol | 18 +- test/Test.t.sol | 16 +- test/WrappedMToken.t.sol | 184 +++++++++--------- test/utils/WrappedMTokenHarness.sol | 34 ++-- 5 files changed, 265 insertions(+), 268 deletions(-) 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; } }