Allow wrap tokens multiple times

m0-foundation · May 22, 2024 · e6e794a · e6e794a
1 parent 0b6c236
commit e6e794a
Showing 1 changed file with 22 additions and 11 deletions.
diff --git 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_);
+    }
 }