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 d742261
Show file tree
Hide file tree
Showing 4 changed files with 88 additions and 88 deletions.
116 changes: 58 additions & 58 deletions src/WrappedMToken.sol
Original file line number Diff line number Diff line change
Expand Up @@ -48,26 +48,37 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {

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

function claimFor(address account_) external returns (uint240 yield_) {
return _claim(account_, currentIndex());
}
function wrap(address recipient_, uint256 amount_) external {
_revertIfInsufficientAmount(amount_);
_revertIfInvalidRecipient(recipient_);

function claimExcess() external returns (uint240 yield_) {
emit ExcessClaimed(yield_ = excess());
emit Transfer(address(0), recipient_, amount_);

IMTokenLike(mToken).transfer(vault, yield_);
// TODO: Technically, might want to `_revertIfInsufficientAmount` if earning principal is 0.
_mint(recipient_, UIntMath.safe240(amount_));

IMTokenLike(mToken).transferFrom(msg.sender, address(this), amount_);
}

function deposit(address recipient_, uint256 amount_) external {
function unwrap(address recipient_, uint256 amount_) external {
_revertIfInsufficientAmount(amount_);
_revertIfInvalidRecipient(recipient_);

emit Transfer(address(0), recipient_, amount_);
emit Transfer(msg.sender, address(0), amount_);

// TODO: Technically, might want to `_revertIfInsufficientAmount` if earning principal is 0.
_addAmount(recipient_, UIntMath.safe240(amount_));
_burn(msg.sender, UIntMath.safe240(amount_));

IMTokenLike(mToken).transferFrom(msg.sender, address(this), amount_);
IMTokenLike(mToken).transfer(recipient_, amount_);
}

function claimFor(address account_) external returns (uint240 yield_) {
return _claim(account_, currentIndex());
}

function claimExcess() external returns (uint240 yield_) {
emit ExcessClaimed(yield_ = excess());

IMTokenLike(mToken).transfer(vault, yield_);
}

function startEarningFor(address account_) external {
Expand Down Expand Up @@ -121,17 +132,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,7 +177,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {

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

function _addAmount(address recipient_, uint240 amount_) internal {
function _mint(address recipient_, uint240 amount_) internal {
(bool isEarning_, , ) = _getBalanceInfo(recipient_);

if (!isEarning_) return _addNonEarningAmount(recipient_, amount_);
Expand All @@ -188,6 +188,17 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
_addEarningAmount(recipient_, amount_, currentIndex_);
}

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

Expand All @@ -197,6 +208,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 +234,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 Down Expand Up @@ -247,41 +282,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
8 changes: 4 additions & 4 deletions src/interfaces/IWrappedMToken.sol
Original file line number Diff line number Diff line change
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 d742261

Please sign in to comment.