feat: Continuous index across disabling/enabling of earning [DO NOT MERGE]

Makefile

@@ -38,10 +38,10 @@ invariant:
 	MAINNET_RPC_URL=$(MAINNET_RPC_URL) ./test.sh -d test/invariant -p $(profile)
 
 coverage:
-	MAINNET_RPC_URL=$(MAINNET_RPC_URL) forge coverage --no-match-path 'test/in*/**/*.sol' --report lcov && lcov --extract lcov.info --rc lcov_branch_coverage=1 --rc derive_function_end_line=0 -o lcov.info 'src/*' && genhtml lcov.info --rc branch_coverage=1 --rc derive_function_end_line=0 -o coverage
+	FOUNDRY_PROFILE=production forge coverage --fork-url $(MAINNET_RPC_URL) --report lcov && lcov --extract lcov.info --rc lcov_branch_coverage=1 --rc derive_function_end_line=0 -o lcov.info 'src/*' && genhtml lcov.info --rc branch_coverage=1 --rc derive_function_end_line=0 -o coverage
 
 gas-report:
-	FOUNDRY_PROFILE=production forge test --no-match-path 'test/integration/**/*.sol' --gas-report > gasreport.ansi
+	FOUNDRY_PROFILE=production forge test --fork-url $(MAINNET_RPC_URL) --gas-report > gasreport.ansi
 
 sizes:
 	./build.sh -p production -s

src/MigratorV1.sol

@@ -1,12 +1,14 @@
-// SPDX-License-Identifier: BUSL-1.1
-
 pragma solidity 0.8.26;
 
+import { IndexingMath } from "../lib/common/src/libs/IndexingMath.sol";
+
 /**
  * @title  Migrator contract for migrating a WrappedMToken contract from V1 to V2.
  * @author M^0 Labs
  */
 contract MigratorV1 {
+    error InvalidEnableDisableEarningIndicesArrayLength();
+
     /// @dev Storage slot with the address of the current factory. `keccak256('eip1967.proxy.implementation') - 1`.
     uint256 private constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
 
@@ -17,10 +19,71 @@ contract MigratorV1 {
     }
 
     fallback() external virtual {
+        (bool earningEnabled_, uint128 disableIndex_) = _clearEnableDisableEarningIndices();
+
+        if (earningEnabled_) {
+            _setEnableMIndex(IndexingMath.EXP_SCALED_ONE);
+        } else {
+            _setDisableIndex(disableIndex_);
+        }
+
         address newImplementation_ = newImplementation;
 
         assembly {
             sstore(_IMPLEMENTATION_SLOT, newImplementation_)
         }
     }
+
+    /**
+     * @dev    Clears the entire `_enableDisableEarningIndices` array in storage, returning useful information.
+     * @return earningEnabled_ Whether earning is enabled.
+     * @return disableIndex_   The index when earning was disabled, if any.
+     */
+    function _clearEnableDisableEarningIndices() internal returns (bool earningEnabled_, uint128 disableIndex_) {
+        uint128[] storage array_;
+
+        assembly {
+            array_.slot := 7 // `_enableDisableEarningIndices` was slot 7 in v1.
+        }
+
+        // If the array is empty, earning is disabled and thus the disable index was non-existent.
+        if (array_.length == 0) return (false, 0);
+
+        // If the array has one element, earning is enabled and the disable index is non-existent.
+        if (array_.length == 1) {
+            array_.pop();
+            return (true, 0);
+        }
+
+        // If the array has two elements, earning is disabled and the disable index is the second element.
+        if (array_.length == 2) {
+            disableIndex_ = array_[1];
+            array_.pop();
+            array_.pop();
+            return (false, disableIndex_);
+        }
+
+        // In v1, it is not possible for the `_enableDisableEarningIndices` array to have more than two elements.
+        revert InvalidEnableDisableEarningIndicesArrayLength();
+    }
+
+    /**
+     * @dev   Sets the `enableMIndex` slot to `index_`.
+     * @param index_ The index to set the `enableMIndex .
+     */
+    function _setEnableMIndex(uint128 index_) internal {
+        assembly {
+            sstore(7, index_) // `enableMIndex` is the lower half of slot 7 in v2.
+        }
+    }
+
+    /**
+     * @dev   Sets the `disableIndex` slot to `index_`.
+     * @param index_ The index to set the `disableIndex .
+     */
+    function _setDisableIndex(uint128 index_) internal {
+        assembly {
+            sstore(7, shl(128, index_)) // `disableIndex` is the upper half of slot 7 in v2.
+        }
+    }
 }

src/WrappedMToken.sol

@@ -84,8 +84,11 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
     /// @dev Mapping of accounts to their respective `AccountInfo` structs.
     mapping(address account => Account balance) internal _accounts;
 
-    /// @dev Array of indices at which earning was enabled or disabled.
-    uint128[] internal _enableDisableEarningIndices;
+    /// @inheritdoc IWrappedMToken
+    uint128 public enableMIndex;
+
+    /// @inheritdoc IWrappedMToken
+    uint128 public disableIndex;
 
     /* ============ Constructor ============ */
 
@@ -175,17 +178,9 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
 
         if (isEarningEnabled()) revert EarningIsEnabled();
 
-        // NOTE: This is a temporary measure to prevent re-enabling earning after it has been disabled.
-        //       This line will be removed in the future.
-        if (wasEarningEnabled()) revert EarningCannotBeReenabled();
-
-        uint128 currentMIndex_ = _currentMIndex();
-
-        _enableDisableEarningIndices.push(currentMIndex_);
+        emit EarningEnabled(enableMIndex = _currentMIndex());
 
         IMTokenLike(mToken).startEarning();
-
-        emit EarningEnabled(currentMIndex_);
     }
 
     /// @inheritdoc IWrappedMToken
@@ -194,21 +189,16 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
 
         if (!isEarningEnabled()) revert EarningIsDisabled();
 
-        uint128 currentMIndex_ = _currentMIndex();
+        emit EarningDisabled(disableIndex = currentIndex());
 
-        _enableDisableEarningIndices.push(currentMIndex_);
+        delete enableMIndex;
 
         IMTokenLike(mToken).stopEarning();
-
-        emit EarningDisabled(currentMIndex_);
     }
 
     /// @inheritdoc IWrappedMToken
     function startEarningFor(address account_) external {
-        if (!isEarningEnabled()) revert EarningIsDisabled();
-
-        // NOTE: Use `currentIndex()` if/when upgrading to support `startEarningFor` while earning is disabled.
-        _startEarningFor(account_, _currentMIndex());
+        _startEarningFor(account_, currentIndex());
     }
 
     /// @inheritdoc IWrappedMToken
@@ -264,7 +254,9 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
 
     /// @inheritdoc IWrappedMToken
     function currentIndex() public view returns (uint128 index_) {
-        return isEarningEnabled() ? _currentMIndex() : _lastDisableEarningIndex();
+        uint128 disableIndex_ = disableIndex == 0 ? IndexingMath.EXP_SCALED_ONE : disableIndex;
+
+        return enableMIndex == 0 ? disableIndex_ : (disableIndex_ * _currentMIndex()) / enableMIndex;
     }
 
     /// @inheritdoc IWrappedMToken
@@ -274,12 +266,7 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
 
     /// @inheritdoc IWrappedMToken
     function isEarningEnabled() public view returns (bool isEnabled_) {
-        return _enableDisableEarningIndices.length % 2 == 1;
-    }
-
-    /// @inheritdoc IWrappedMToken
-    function wasEarningEnabled() public view returns (bool wasEarning_) {
-        return _enableDisableEarningIndices.length != 0;
+        return enableMIndex != 0;
     }
 
     /// @inheritdoc IWrappedMToken
@@ -663,11 +650,6 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
         return IMTokenLike(mToken).currentIndex();
     }
 
-    /// @dev Returns the earning index from the last `disableEarning` call.
-    function _lastDisableEarningIndex() internal view returns (uint128 index_) {
-        return wasEarningEnabled() ? _unsafeAccess(_enableDisableEarningIndices, 1) : 0;
-    }
-
     /**
      * @dev    Compute the yield given an account's balance, last index, and the current index.
      * @param  balance_      The token balance of an earning account.
@@ -788,23 +770,4 @@ contract WrappedMToken is IWrappedMToken, Migratable, ERC20Extended {
     function _revertIfNotApprovedEarner(address account_) internal view {
         if (!_isApprovedEarner(account_)) revert NotApprovedEarner(account_);
     }
-
-    /**
-     * @dev   Reads the uint128 value at some index of an array of uint128 values whose storage pointer is given,
-     *        assuming the index is valid, without wasting gas checking for out-of-bounds errors.
-     * @param array_ The storage pointer of an array of uint128 values.
-     * @param i_     The index of the array to read.
-     */
-    function _unsafeAccess(uint128[] storage array_, uint256 i_) internal view returns (uint128 value_) {
-        assembly {
-            mstore(0, array_.slot)
-
-            value_ := sload(add(keccak256(0, 0x20), div(i_, 2)))
-
-            // Since uint128 values take up either the top half or bottom half of a slot, shift the result accordingly.
-            if eq(mod(i_, 2), 1) {
-                value_ := shr(128, value_)
-            }
-        }
-    }
 }

src/interfaces/IWrappedMToken.sol

@@ -22,13 +22,13 @@ interface IWrappedMToken is IMigratable, IERC20Extended {
 
     /**
      * @notice Emitted when Wrapped M earning is enabled.
-     * @param  index The index at the moment earning is enabled.
+     * @param  index The M index at the moment earning is enabled.
      */
     event EarningEnabled(uint128 index);
 
     /**
      * @notice Emitted when Wrapped M earning is disabled.
-     * @param  index The index at the moment earning is disabled.
+     * @param  index The WrappedM index at the moment earning is disabled.
      */
     event EarningDisabled(uint128 index);
 
@@ -58,11 +58,8 @@ interface IWrappedMToken is IMigratable, IERC20Extended {
     /// @notice Emitted when performing an operation that is not allowed when earning is enabled.
     error EarningIsEnabled();
 
-    /// @notice Emitted when trying to enable earning after it has been explicitly disabled.
-    error EarningCannotBeReenabled();
-
     /**
-     * @notice Emitted when calling `stopEarning` for an account approved as earner by the Registrar.
+     * @notice Emitted when calling `stopEarning` for an account approved as an earner by the Registrar.
      * @param  account The account that is an approved earner.
      */
     error IsApprovedEarner(address account);
@@ -76,7 +73,7 @@ interface IWrappedMToken is IMigratable, IERC20Extended {
     error InsufficientBalance(address account, uint240 balance, uint240 amount);
 
     /**
-     * @notice Emitted when calling `startEarning` for an account not approved as earner by the Registrar.
+     * @notice Emitted when calling `startEarning` for an account not approved as an earner by the Registrar.
      * @param  account The account that is not an approved earner.
      */
     error NotApprovedEarner(address account);
@@ -246,9 +243,15 @@ interface IWrappedMToken is IMigratable, IERC20Extended {
     /// @notice The current index of Wrapped M's earning mechanism.
     function currentIndex() external view returns (uint128 index);
 
+    /// @notice The M token's index when earning was most recently enabled.
+    function enableMIndex() external view returns (uint128 enableMIndex);
+
     /// @notice This contract's current excess M that is not earmarked for account balances or accrued yield.
     function excess() external view returns (uint240 excess);
 
+    /// @notice The wrapper's index when earning was most recently disabled.
+    function disableIndex() external view returns (uint128 disableIndex);
+
     /**
      * @notice Returns whether `account` is a wM earner.
      * @param  account   The account being queried.
@@ -259,9 +262,6 @@ interface IWrappedMToken is IMigratable, IERC20Extended {
     /// @notice Whether Wrapped M earning is enabled.
     function isEarningEnabled() external view returns (bool isEnabled);
 
-    /// @notice Whether Wrapped M earning has been enabled at least once.
-    function wasEarningEnabled() external view returns (bool wasEnabled);
-
     /// @notice The account that can bypass the Registrar and call the `migrate(address migrator)` function.
     function migrationAdmin() external view returns (address migrationAdmin);
 

test/integration/Migration.sol

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: GPL-3.0
+
+pragma solidity 0.8.26;
+
+import { IndexingMath } from "../../lib/common/src/libs/IndexingMath.sol";
+
+import { TestBase } from "./TestBase.sol";
+
+contract MigrationIntegrationTests is TestBase {
+    function test_index_noMigration() external {
+        assertEq(_wrappedMToken.currentIndex(), 1_023463403719);
+
+        vm.warp(vm.getBlockTimestamp() + 365 days);
+
+        assertEq(_wrappedMToken.currentIndex(), 1_073787769981);
+    }
+
+    function test_index_migrate_beforeEarningDisabled() external {
+        assertEq(_wrappedMToken.currentIndex(), 1_023463403719);
+
+        _deployV2Components();
+        _migrate();
+
+        assertEq(_wrappedMToken.disableIndex(), 0);
+        assertEq(_wrappedMToken.enableMIndex(), IndexingMath.EXP_SCALED_ONE);
+
+        assertEq(_wrappedMToken.currentIndex(), 1_023463403719);
+
+        vm.warp(vm.getBlockTimestamp() + 365 days);
+
+        assertEq(_wrappedMToken.currentIndex(), 1_073787769981);
+    }
+
+    function test_index_migrate_afterEarningDisabled() external {
+        assertEq(_wrappedMToken.currentIndex(), 1_023463403719);
+
+        _removeFromList(_EARNERS_LIST_NAME, address(_wrappedMToken));
+
+        _wrappedMToken.disableEarning();
+
+        _deployV2Components();
+        _migrate();
+
+        assertEq(_wrappedMToken.disableIndex(), 1_023463403719);
+        assertEq(_wrappedMToken.enableMIndex(), 0);
+
+        assertEq(_wrappedMToken.currentIndex(), 1_023463403719);
+
+        vm.warp(vm.getBlockTimestamp() + 365 days);
+
+        assertEq(_wrappedMToken.currentIndex(), 1_023463403719);
+    }
+}