Skip to content

Commit

Permalink
Add renames
Browse files Browse the repository at this point in the history
  • Loading branch information
toninorair authored and deluca-mike committed Jun 27, 2024
1 parent bda198a commit f5f3f4c
Show file tree
Hide file tree
Showing 4 changed files with 96 additions and 92 deletions.
126 changes: 65 additions & 61 deletions src/WrappedMToken.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
Expand All @@ -58,18 +70,6 @@ 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();

Expand Down Expand Up @@ -121,17 +121,6 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
_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_) {
Expand Down Expand Up @@ -177,17 +166,41 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {

/* ============ Internal Interactive Functions ============ */

function _addAmount(address recipient_, uint240 amount_) internal {
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 _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 {
(, , uint240 rawBalance_) = _getBalanceInfo(recipient_);

Expand All @@ -197,6 +210,17 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
}
}

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

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

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

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

Expand All @@ -212,6 +236,19 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
_addTotalEarningSupply(amount_, currentIndex_);
}

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

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

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

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

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

Expand All @@ -222,20 +259,22 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {

if (yield_ == 0) return 0;

emit Claimed(account_, yield_);

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

address claimOverrideRecipient_ = _getClaimOverrideRecipient(account_);

if (claimOverrideRecipient_ == address(0)) {
emit Claimed(account_, account_, yield_);

// NOTE: The `Transfer` event for a non-zero `claimOverrideRecipient_`
emit Transfer(address(0), account_, yield_);
return yield_;
}

emit Claimed(account_, claimOverrideRecipient_, yield_);

// NOTE: Watch out for a long chain of claim override recipients.
// TODO: Maybe can be optimized since we know `account_` is an earner and already claimed.
_transfer(account_, claimOverrideRecipient_, yield_, currentIndex_);
Expand All @@ -247,41 +286,6 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
: 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_);

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

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

uint112 principalAmount_ = IndexingMath.getPrincipalAmountRoundedUp(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_);

Expand Down
10 changes: 5 additions & 5 deletions src/interfaces/IWrappedMToken.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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);
Expand Down
16 changes: 8 additions & 8 deletions test/Test.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand All @@ -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);

Expand Down Expand Up @@ -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);

Expand All @@ -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);

Expand Down Expand Up @@ -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);

Expand All @@ -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);

Expand All @@ -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);

Expand All @@ -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);

Expand Down
Loading

0 comments on commit f5f3f4c

Please sign in to comment.