Skip to content

Commit

Permalink
feat: No claim on wrap/unwrap/transfer
Browse files Browse the repository at this point in the history
  • Loading branch information
@deluca-mike
deluca-mike committed Jan 14, 2025
1 parent e261b7b commit ae1eea3
Show file tree
Hide file tree
Showing 6 changed files with 214 additions and 274 deletions.
22 changes: 8 additions & 14 deletions src/WrappedMToken.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -151,9 +151,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {

/// @inheritdoc IWrappedMToken
function unwrap(address recipient_) external returns (uint240 unwrapped_) {
_migrateEarner(msg.sender); // NOTE: Need to migrate before calling `balanceWithYieldOf`.

return _unwrap(msg.sender, recipient_, uint240(balanceWithYieldOf(msg.sender)));
return _unwrap(msg.sender, recipient_, uint240(balanceOf(msg.sender)));
}

/// @inheritdoc IWrappedMToken
Expand Down Expand Up @@ -262,7 +260,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
}

/// @inheritdoc IWrappedMToken
function balanceWithYieldOf(address account_) public view returns (uint256 balance_) {
function balanceWithYieldOf(address account_) external view returns (uint256 balance_) {
unchecked {
return balanceOf(account_) + accruedYieldOf(account_);
}
Expand Down Expand Up @@ -343,12 +341,10 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
if (_accounts[recipient_].earningState == EarningState.NOT_EARNING) {
_addNonEarningAmount(recipient_, amount_);
} else {
uint128 currentIndex_ = currentIndex();

_claim(recipient_, currentIndex_);
_migrateEarner(recipient_);

// NOTE: Additional principal may end up being rounded to 0 and this will not `_revertIfInsufficientAmount`.
_addEarningAmount(recipient_, amount_, currentIndex_);
_addEarningAmount(recipient_, amount_, currentIndex());
}

emit Transfer(address(0), recipient_, amount_);
Expand All @@ -365,12 +361,10 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
if (_accounts[account_].earningState == EarningState.NOT_EARNING) {
_subtractNonEarningAmount(account_, amount_);
} else {
uint128 currentIndex_ = currentIndex();

_claim(account_, currentIndex_);
_migrateEarner(account_);

// NOTE: Subtracted principal may end up being rounded to 0 and this will not `_revertIfInsufficientAmount`.
_subtractEarningAmount(account_, amount_, currentIndex_);
_subtractEarningAmount(account_, amount_, currentIndex());
}

emit Transfer(account_, address(0), amount_);
Expand Down Expand Up @@ -504,8 +498,8 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
_revertIfZeroAccount(sender_);
_revertIfZeroAccount(recipient_);

_claim(sender_, currentIndex_);
_claim(recipient_, currentIndex_);
_migrateEarner(sender_);
_migrateEarner(recipient_);

emit Transfer(sender_, recipient_, amount_);

Expand Down
24 changes: 12 additions & 12 deletions test/integration/MorphoBlue.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,8 @@ contract MorphoBlueTests is MorphoTestBase {
// borrowing 0.90 USDC.
_createMarket(_alice, _USDC);

// The market creation has triggered a wM transfer and the yield has been claimed for morpho.
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield -= _morphoAccruedYield);
// The market creation has triggered a wM transfer ubt the yield has not been claimed for morpho.
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield -= 1);

assertEq(_wrappedMToken.balanceOf(_alice), _aliceBalanceOfWM -= 1e6);
assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM += 1e6);
Expand Down Expand Up @@ -97,7 +97,7 @@ contract MorphoBlueTests is MorphoTestBase {
vm.warp(vm.getBlockTimestamp() + 365 days);

assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM);
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 49_292100);
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 49_292110);

// USDC balance is unchanged.
assertEq(IERC20(_USDC).balanceOf(_MORPHO), _morphoBalanceOfUSDC);
Expand All @@ -111,8 +111,8 @@ contract MorphoBlueTests is MorphoTestBase {

_withdrawCollateral(_bob, address(_wrappedMToken), 1_000e6, _bob, _USDC);

// The collateral withdrawal has triggered a wM transfer and the yield has been claimed for morpho.
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield -= _morphoAccruedYield);
// The collateral withdrawal has triggered a wM transfer but the yield has not been claimed for morpho.
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield -= 1);

assertEq(_wrappedMToken.balanceOf(_bob), _bobBalanceOfWM += 1_000e6);
assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM -= 1_000e6);
Expand All @@ -130,7 +130,7 @@ contract MorphoBlueTests is MorphoTestBase {
vm.warp(vm.getBlockTimestamp() + 365 days);

assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM);
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 121446);
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 2_545180);

// USDC balance is unchanged.
assertEq(IERC20(_USDC).balanceOf(_MORPHO), _morphoBalanceOfUSDC);
Expand All @@ -151,8 +151,8 @@ contract MorphoBlueTests is MorphoTestBase {
// borrowing 0.90 wM.
_createMarket(_alice, address(_wrappedMToken));

// The market creation has triggered a wM transfer and the yield has been claimed for morpho.
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield -= _morphoAccruedYield);
// The market creation has triggered a wM transfer but the yield has ont been claimed for morpho.
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield -= 2);

assertEq(_wrappedMToken.balanceOf(_alice), _aliceBalanceOfWM -= 100000);
assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM += 100000);
Expand Down Expand Up @@ -190,7 +190,7 @@ contract MorphoBlueTests is MorphoTestBase {

// `startEarningFor` has been called so wM yield has accrued in the pool.
assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM);
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 4_994256);
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 4_994266);

// USDC balance is unchanged.
assertEq(IERC20(_USDC).balanceOf(_MORPHO), _morphoBalanceOfUSDC);
Expand All @@ -199,8 +199,8 @@ contract MorphoBlueTests is MorphoTestBase {

_repay(_bob, address(_wrappedMToken), 900e6, _USDC);

// The repay has triggered a wM transfer and the yield has been claimed for morpho.
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield -= _morphoAccruedYield);
// The repay has triggered a wM transfer but the yield has not been claimed for morpho.
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield);

assertEq(_wrappedMToken.balanceOf(_bob), _bobBalanceOfWM -= 900e6);
assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM += 900e6);
Expand All @@ -224,7 +224,7 @@ contract MorphoBlueTests is MorphoTestBase {

// `startEarningFor` has been called so wM yield has accrued in the pool.
assertEq(_wrappedMToken.balanceOf(_MORPHO), _morphoBalanceOfWM);
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 77192);
assertEq(_wrappedMToken.accruedYieldOf(_MORPHO), _morphoAccruedYield += 322772);

// USDC balance is unchanged.
assertEq(IERC20(_USDC).balanceOf(_MORPHO), _morphoBalanceOfUSDC);
Expand Down
42 changes: 21 additions & 21 deletions test/integration/Protocol.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -290,20 +290,20 @@ contract ProtocolIntegrationTests is TestBase {
assertEq(_wrappedMToken.accruedYieldOf(_dave), 0);

// Alice transfers all her tokens and only keeps her accrued yield.
_transferWM(_alice, _carol, 100_000000);
_transferWM(_alice, _carol, 99_999999);

// Assert Alice (Earner)
assertEq(_wrappedMToken.balanceOf(_alice), _aliceBalance = _aliceBalance + 2_395361 - 100_000000);
assertEq(_wrappedMToken.accruedYieldOf(_alice), _aliceAccruedYield -= 2_395361);
assertEq(_wrappedMToken.balanceOf(_alice), _aliceBalance -= 99_999999);
assertEq(_wrappedMToken.accruedYieldOf(_alice), _aliceAccruedYield -= 1);

// Assert Carol (Non-Earner)
assertEq(_wrappedMToken.balanceOf(_carol), _carolBalance += 100_000000);
assertEq(_wrappedMToken.balanceOf(_carol), _carolBalance += 99_999999);
assertEq(_wrappedMToken.accruedYieldOf(_carol), 0);

// Assert Globals
assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += _bobBalance + 2_395361 - 100_000000);
assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply += _daveBalance + 100_000000);
assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 2_395361 + 1);
assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply += _bobBalance - 99_999999);
assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply += _daveBalance + 99_999999);
assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 1);
assertEq(_wrappedMToken.excess(), _excess += 2);

assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess);
Expand Down Expand Up @@ -338,7 +338,7 @@ contract ProtocolIntegrationTests is TestBase {

// Assert Alice (Earner)
assertEq(_wrappedMToken.balanceOf(_alice), _aliceBalance);
assertEq(_wrappedMToken.accruedYieldOf(_alice), _aliceAccruedYield += 57376);
assertEq(_wrappedMToken.accruedYieldOf(_alice), _aliceAccruedYield += 57378);

// Assert Bob (Earner)
assertEq(_wrappedMToken.balanceOf(_bob), _bobBalance);
Expand Down Expand Up @@ -482,35 +482,35 @@ contract ProtocolIntegrationTests is TestBase {

assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess);

// Accrued yield of Bob is claimed when unwrapping
_unwrap(_bob, _bob, _bobBalance + _bobAccruedYield);
// Accrued yield of Bob is not claimed when unwrapping
_unwrap(_bob, _bob, _bobBalance);

// Assert Bob (Earner)
assertEq(_mToken.balanceOf(_bob), 100_000000 + 3_614474 - 1);
assertEq(_mToken.balanceOf(_bob), 100_000000 - 1);
assertEq(_wrappedMToken.balanceOf(_bob), _bobBalance -= 100_000000);
assertEq(_wrappedMToken.accruedYieldOf(_bob), _bobAccruedYield -= 3_614474);
assertEq(_wrappedMToken.accruedYieldOf(_bob), _bobAccruedYield -= 1);

// Assert Globals
assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply -= 100_000000);
assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply);
assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 3_614474);
assertEq(_wrappedMToken.excess(), _excess += 1);
assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield);
assertEq(_wrappedMToken.excess(), _excess);

assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess);

// Accrued yield of Carol is claimed when unwrapping
_unwrap(_carol, _carol, _carolBalance + _carolAccruedYield);
// Accrued yield of Carol is not claimed when unwrapping
_unwrap(_carol, _carol, _carolBalance);

// Assert Carol (Earner)
assertEq(_mToken.balanceOf(_carol), 100_000000 + 2_395361 - 1);
assertEq(_mToken.balanceOf(_carol), 100_000000 - 1);
assertEq(_wrappedMToken.balanceOf(_carol), _carolBalance -= 100_000000);
assertEq(_wrappedMToken.accruedYieldOf(_carol), _carolAccruedYield -= 2_395361);
assertEq(_wrappedMToken.accruedYieldOf(_carol), _carolAccruedYield);

// Assert Globals
assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply -= 100_000000);
assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply);
assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 2_395361 + 1);
assertEq(_wrappedMToken.excess(), _excess += 1);
assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield -= 1);
assertEq(_wrappedMToken.excess(), _excess += 2);

assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess);

Expand All @@ -525,7 +525,7 @@ contract ProtocolIntegrationTests is TestBase {
assertEq(_wrappedMToken.totalEarningSupply(), _totalEarningSupply);
assertEq(_wrappedMToken.totalNonEarningSupply(), _totalNonEarningSupply -= 100_000000);
assertEq(_wrappedMToken.totalAccruedYield(), _totalAccruedYield);
assertEq(_wrappedMToken.excess(), _excess);
assertEq(_wrappedMToken.excess(), _excess += 1);

assertGe(_wrapperBalanceOfM, _totalEarningSupply + _totalNonEarningSupply + _totalAccruedYield + _excess);

Expand Down
Loading

0 comments on commit ae1eea3

Please sign in to comment.