Logic for the bare min start/stop earning

m0-foundation · Jul 2, 2024 · cf69acc · cf69acc
1 parent a5b81bb
commit cf69acc
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 2 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/src/WrappedMToken.sol b/src/WrappedMToken.sol
@@ -35,8 +35,16 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
 
     uint240 public totalNonEarningSupply;
 
+    uint128 public mIndexWhenEarningStopped;
+
     mapping(address account => BalanceInfo balance) internal _balances;
 
+    modifier onlyWhenEarning() {
+        if (IMTokenLike(mToken).isEarning(address(this))) revert NotInEarningState();
+
+        _;
+    }
+
     /* ============ Constructor ============ */
 
     constructor(address mToken_) ERC20Extended("WrappedM by M^0", "wM", 6) {
@@ -48,7 +56,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
 
     /* ============ Interactive Functions ============ */
 
-    function wrap(address recipient_, uint256 amount_) external {
+    function wrap(address recipient_, uint256 amount_) external onlyWhenEarning {
         _mint(recipient_, UIntMath.safe240(amount_));
 
         IMTokenLike(mToken).transferFrom(msg.sender, address(this), amount_);
@@ -110,6 +118,18 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
         }
     }
 
+    function startEarningM() external {
+        if (mIndexWhenEarningStopped != 0) revert OnlyEarningOnce();
+
+        IMTokenLike(mToken).startEarning();
+    }
+
+    function stopEarningM() external {
+        mIndexWhenEarningStopped = currentIndex();
+
+        IMTokenLike(mToken).stopEarning();
+    }
+
     /* ============ View/Pure Functions ============ */
 
     function accruedYieldOf(address account_) external view returns (uint240 yield_) {
@@ -123,7 +143,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
     }
 
     function currentIndex() public view returns (uint128 index_) {
-        return IMTokenLike(mToken).currentIndex();
+        return mIndexWhenEarningStopped == 0 ? IMTokenLike(mToken).currentIndex() : mIndexWhenEarningStopped;
     }
 
     function isEarning(address account_) external view returns (bool isEarning_) {

diff --git a/src/interfaces/IMTokenLike.sol b/src/interfaces/IMTokenLike.sol
@@ -9,8 +9,14 @@ interface IMTokenLike {
 
     function transferFrom(address sender, address recipient, uint256 amount) external returns (bool success);
 
+    function startEarning() external;
+
+    function stopEarning() external;
+
     /* ============ View/Pure Functions ============ */
 
+    function isEarning(address account) external view returns (bool earning);
+
     function balanceOf(address account) external view returns (uint256 balance);
 
     function currentIndex() external view returns (uint128 currentIndex);

diff --git a/src/interfaces/IWrappedMToken.sol b/src/interfaces/IWrappedMToken.sol
@@ -41,6 +41,12 @@ interface IWrappedMToken is IMigratable, IERC20Extended {
     /// @notice Emitted when calling `startEarning` for an account not approved as earner by TTG.
     error NotApprovedEarner();
 
+    /// @notice Emitted when calling `startEarningM` after wrapper was already in earning state once.
+    error OnlyEarningOnce();
+
+    /// @notice Emitted when method is called when Wrapped M is not in earning state.
+    error NotInEarningState();
+
     /// @notice Emitted in constructor if M Token is 0x0.
     error ZeroMToken();
 
@@ -54,6 +60,10 @@ interface IWrappedMToken is IMigratable, IERC20Extended {
 
     function claimExcess() external returns (uint240 yield);
 
+    function startEarningM() external;
+
+    function stopEarningM() external;
+
     function startEarningFor(address account) external;
 
     function stopEarningFor(address account) external;