diff --git a/src/option1/wM.sol b/src/option1/wM.sol index 23bfb36..e004d78 100644 --- a/src/option1/wM.sol +++ b/src/option1/wM.sol @@ -17,29 +17,31 @@ contract wM is ERC20 { } address public immutable mToken; - address public immutable earner; - address public immutable nonEarner; + address public immutable earnerPosition; + address public immutable nonEarnerPosition; mapping(address earner => EarnerBalance balance) public earningSuppliers; - error OnlyOnePositionPerEarner(); // for simplicity purposes to start error InsufficientBalance(address account, uint256 rawBalance, uint256 amount); constructor(address mToken_) ERC20("Wrapped M Token", "wM", 18) { mToken = mToken_; - earner = address(new MPosition(address(this), mToken_)); // on M earner list - nonEarner = address(new MPosition(address(this), mToken_)); + earnerPosition = address(new MPosition(address(this), mToken_)); // on M earner list + nonEarnerPosition = address(new MPosition(address(this), mToken_)); } function wrap(uint256 amount) external { if (IMToken(mToken).isEarning(msg.sender)) { - if (earningSuppliers[msg.sender].principal != 0) revert OnlyOnePositionPerEarner(); + (uint256 initialBalance_, uint256 interest_) = _getAccruedInterestAndInitialBalance(msg.sender); + if (interest_ > 0) { + MPosition(earnerPosition).withdraw(interest_, msg.sender); // withdraw M interest + } uint256 currentIndex_ = IMToken(mToken).currentIndex(); - earningSuppliers[msg.sender] = EarnerBalance(currentIndex_, amount / currentIndex_); - IMToken(mToken).transferFrom(msg.sender, earner, amount); + earningSuppliers[msg.sender] = EarnerBalance(currentIndex_, (initialBalance_ + amount) / currentIndex_); + IMToken(mToken).transferFrom(msg.sender, earnerPosition, amount); } else { - IMToken(mToken).transferFrom(msg.sender, nonEarner, amount); + IMToken(mToken).transferFrom(msg.sender, nonEarnerPosition, amount); } _mint(msg.sender, amount); // mint wM tokens @@ -57,15 +59,24 @@ contract wM is ERC20 { uint256 mInterest_ = currentBalance_ - initialBalance_; // Withdraw interest and required amount - MPosition(earner).withdraw(amount_ + mInterest_, msg.sender); + MPosition(earnerPosition).withdraw(amount_ + mInterest_, msg.sender); earnerBalance_.principal = (initialBalance_ - amount_) / currentIndex_; earnerBalance_.lastIndex = currentIndex_; } else { // non-earner - MPosition(nonEarner).withdraw(amount_, msg.sender); + MPosition(nonEarnerPosition).withdraw(amount_, msg.sender); } _burn(msg.sender, amount_); // burn wM tokens } + + function _getAccruedInterestAndInitialBalance(address earner) internal view returns (uint256, uint256) { + uint256 currentIndex_ = IMToken(mToken).currentIndex(); + EarnerBalance storage earnerBalance_ = earningSuppliers[earner]; + uint256 initialBalance_ = earnerBalance_.principal * earnerBalance_.lastIndex; + uint256 currentBalance_ = earnerBalance_.principal * currentIndex_; + + return (initialBalance_, currentBalance_ - initialBalance_); + } }