Improved readability of API for cumulative rewards

contracts/staking/contracts/src/staking_pools/MixinCumulativeRewards.sol

@@ -41,12 +41,20 @@ contract MixinCumulativeRewards is
     using LibSafeMath for uint256;
 
     /// @dev Initializes Cumulative Rewards for a given pool.
+    /// @param poolId Unique id of pool.
     function _initializeCumulativeRewards(bytes32 poolId)
         internal
     {
         uint256 currentEpoch = getCurrentEpoch();
-        _setCumulativeReward(poolId, currentEpoch, IStructs.Fraction({numerator: 0, denominator: MIN_TOKEN_VALUE}));
+        // sets the default cumulative reward
-        _setMostRecentCumulativeReward(poolId, currentEpoch);
+        _forceSetCumulativeReward(
+            poolId,
+            currentEpoch,
+            IStructs.Fraction({
+                numerator: 0,
+                denominator: MIN_TOKEN_VALUE
+            })
+        );
     }
 
     /// @dev returns true iff Cumulative Rewards are set
@@ -61,11 +69,45 @@ contract MixinCumulativeRewards is
         return cumulativeReward.denominator != 0;
     }
 
-    /// @dev Sets a cumulative reward for `poolId` at `epoch`.
+    /// Returns true iff the cumulative reward for `poolId` at `epoch` can be unset.
+    /// @param poolId Unique id of pool.
+    /// @param epoch of the cumulative reward.
+    function _canUnsetCumulativeReward(bytes32 poolId, uint256 epoch)
+        internal
+        view
+        returns (bool)
+    {
+        return (
+            _isCumulativeRewardSet(cumulativeRewardsByPool[poolId][epoch]) &&   // is there a value to unset
+            cumulativeRewardsByPoolReferenceCounter[poolId][epoch] == 0 &&      // no references to this CR
+            cumulativeRewardsByPoolLastStored[poolId] > epoch                   // this is *not* the most recent CR
+        );
+    }
+
+    /// @dev Tries to set a cumulative reward for `poolId` at `epoch`.
     /// @param poolId Unique Id of pool.
     /// @param epoch of cumulative reward.
     /// @param value of cumulative reward.
-    function _setCumulativeReward(
+    function _trySetCumulativeReward(
+        bytes32 poolId,
+        uint256 epoch,
+        IStructs.Fraction memory value
+    )
+        internal
+    {
+        if (_isCumulativeRewardSet(cumulativeRewardsByPool[poolId][epoch])) {
+            // do nothing; we don't want to override the current value
+            return;
+        }
+        _forceSetCumulativeReward(poolId, epoch, value);
+    }
+
+    /// @dev Sets a cumulative reward for `poolId` at `epoch`.
+    /// This can be used to overwrite an existing value.
+    /// @param poolId Unique Id of pool.
+    /// @param epoch of cumulative reward.
+    /// @param value of cumulative reward.
+    function _forceSetCumulativeReward(
         bytes32 poolId,
         uint256 epoch,
         IStructs.Fraction memory value
@@ -73,15 +115,27 @@ contract MixinCumulativeRewards is
         internal
     {
         cumulativeRewardsByPool[poolId][epoch] = value;
+        _trySetMostRecentCumulativeRewardEpoch(poolId, epoch);
+    }
+
+    /// @dev Tries to unset the cumulative reward for `poolId` at `epoch`.
+    /// @param poolId Unique id of pool.
+    /// @param epoch of cumulative reward to unset.
+    function _tryUnsetCumulativeReward(bytes32 poolId, uint256 epoch)
+        internal
+    {
+        if (!_canUnsetCumulativeReward(poolId, epoch)) {
+            return;
+        }
+        _forceUnsetCumulativeReward(poolId, epoch);
     }
 
     /// @dev Unsets the cumulative reward for `poolId` at `epoch`.
     /// @param poolId Unique id of pool.
     /// @param epoch of cumulative reward to unset.
-    function _unsetCumulativeReward(bytes32 poolId, uint256 epoch)
+    function _forceUnsetCumulativeReward(bytes32 poolId, uint256 epoch)
         internal
     {
-        assert(cumulativeRewardsByPoolReferenceCounter[poolId][epoch] == 0);
         cumulativeRewardsByPool[poolId][epoch] = IStructs.Fraction({numerator: 0, denominator: 0});
     }
 
@@ -100,11 +154,44 @@ contract MixinCumulativeRewards is
         });
     }
 
+    /// @dev Tries to set the epoch of the most recent cumulative reward.
+    /// The value will only be set if the input epoch is greater than the current
+    /// most recent value.
+    /// @param poolId Unique Id of pool.
+    /// @param epoch of the most recent cumulative reward.
+    function _trySetMostRecentCumulativeRewardEpoch(bytes32 poolId, uint256 epoch)
+        internal
+    {
+        // load the current value, sanity check that we're not trying to go back in time
+        uint256 currentMostRecentEpoch = cumulativeRewardsByPoolLastStored[poolId];
+        assert(epoch >= currentMostRecentEpoch);
+
+        // check if we should do any work
+        if (epoch == currentMostRecentEpoch) {
+            return;
+        }
+
+        // update state to reflect the most recent cumulative reward
+        _forceSetMostRecentCumulativeRewardEpoch(poolId, epoch);
+
+        // unset the previous most recent reward, if it is no longer needed
+        _tryUnsetCumulativeReward(poolId, currentMostRecentEpoch);
+    }
+
+    /// @dev Forcefully sets the epoch of the most recent cumulative reward.
+    /// @param poolId Unique Id of pool.
+    /// @param epoch of the most recent cumulative reward.
+    function _forceSetMostRecentCumulativeRewardEpoch(bytes32 poolId, uint256 epoch)
+        internal
+    {
+        cumulativeRewardsByPoolLastStored[poolId] = epoch;
+    }
+
     /// @dev Adds a dependency on a cumulative reward for a given epoch.
     /// @param poolId Unique Id of pool.
     /// @param epoch to remove dependency from.
-    /// @param mostRecentCumulativeRewardInfo Epoch of the most recent cumulative reward.
+    /// @param mostRecentCumulativeRewardInfo Info for the most recent cumulative reward (value and epoch)
-    /// @param isDependent still FGREG EEGGEGREG
+    /// @param isDependent True iff there is a dependency on the cumulative reward for `poolId` at `epoch`
     function _addOrRemoveDependencyOnCumulativeReward(
         bytes32 poolId,
         uint256 epoch,
@@ -142,11 +229,12 @@ contract MixinCumulativeRewards is
         // add dependency by increasing the reference counter
         cumulativeRewardsByPoolReferenceCounter[poolId][epoch] = cumulativeRewardsByPoolReferenceCounter[poolId][epoch].safeAdd(1);
 
-        // ensure dependency has a value, otherwise copy it from the most recent reward
+        // set CR to most recent reward (if it is not already set)
-        if (!_isCumulativeRewardSet(cumulativeRewardsByPool[poolId][epoch])) {
+        _trySetCumulativeReward(
-            _setCumulativeReward(poolId, epoch, mostRecentCumulativeRewardInfo.cumulativeReward);
+            poolId,
-            _setMostRecentCumulativeReward(poolId, epoch);
+            epoch,
-        }
+            mostRecentCumulativeRewardInfo.cumulativeReward
+        );
     }
 
     /// @dev Removes a dependency on a cumulative reward for a given epoch.
@@ -165,9 +253,7 @@ contract MixinCumulativeRewards is
         cumulativeRewardsByPoolReferenceCounter[poolId][epoch] = newReferenceCounter;
 
         // clear cumulative reward from state, if it is no longer needed
-        if (newReferenceCounter == 0 && epoch < mostRecentCumulativeRewardEpoch) {
+        _tryUnsetCumulativeReward(poolId, epoch);
-            _unsetCumulativeReward(poolId, epoch);
-        }
     }
 
     /// @dev Computes a member's reward over a given epoch interval.
@@ -203,28 +289,4 @@ contract MixinCumulativeRewards is
         );
         return reward;
     }
-
-    /// @dev Sets the most recent cumulative reward for the pool.
-    /// @param poolId Unique Id of pool.
-    /// @param epoch The epoch of the most recent cumulative reward.
-    function _setMostRecentCumulativeReward(bytes32 poolId, uint256 epoch)
-        internal
-    {
-        // load the current value, sanity check that we're not trying to go back in time
-        uint256 currentMostRecentEpoch = cumulativeRewardsByPoolLastStored[poolId];
-        assert(epoch >= currentMostRecentEpoch);
-
-        // check if we should do any work
-        if (epoch == currentMostRecentEpoch) {
-            return;
-        }
-
-        // unset the current most recent reward if it has no more references
-        if (cumulativeRewardsByPoolReferenceCounter[poolId][currentMostRecentEpoch] == 0) {
-            _unsetCumulativeReward(poolId, currentMostRecentEpoch);
-        }
-
-        // update state to reflect the most recent cumulative reward
-        cumulativeRewardsByPoolLastStored[poolId] = epoch;
-    }
 }

contracts/staking/contracts/src/staking_pools/MixinStakingPoolRewards.sol

@@ -135,7 +135,7 @@ contract MixinStakingPoolRewards is
         );
 
         // store cumulative rewards and set most recent
-        _setCumulativeReward(
+        _trySetCumulativeReward(
             poolId,
             epoch,
             IStructs.Fraction({
@@ -143,7 +143,6 @@ contract MixinStakingPoolRewards is
                 denominator: denominatorNormalized
             })
         );
-        _setMostRecentCumulativeReward(poolId, epoch);
     }
 
     /// @dev Transfers a delegators accumulated rewards from the transient pool Reward Pool vault
@@ -187,8 +186,8 @@ contract MixinStakingPoolRewards is
         view
         returns (uint256 totalReward)
     {
-        // value is always zero in these two scenarios:
+        // reward balance is always zero in these two scenarios:
-        //   1. The owner's delegated is current as of this epoch: their rewards have been moved to the ETH vault.
+        //   1. The owner's delegated stake is current as of this epoch: their rewards have been moved to the ETH vault.
         //   2. The current epoch is zero: delegation begins at epoch 1
         if (unsyncedDelegatedStakeToPoolByOwner.currentEpoch == currentEpoch || currentEpoch == 0) return 0;
 

contracts/staking/contracts/src/sys/MixinParams.sol

@@ -102,8 +102,11 @@ contract MixinParams is
         _cobbDouglasAlphaDenomintor = cobbDouglasAlphaDenomintor;
     }
 
-    /// @dev Initialzize storage belonging to this mixin.
+    /// @dev Assert param values before initializing them.
-    function _initMixinParams() internal {
+    /// This must be updated for each migration.
+    function _assertMixinParamsBeforeInit()
+        internal
+    {
         // Ensure state is uninitialized.
         if (epochDurationInSeconds != 0 &&
             rewardDelegatedStakeWeight != 0 &&
@@ -118,6 +121,15 @@ contract MixinParams is
                 )
             );
         }
+    }
+
+    /// @dev Initialize storage belonging to this mixin.
+    function _initMixinParams()
+        internal
+    {
+        // assert the current values before overwriting them.
+        _assertMixinParamsBeforeInit();
+
         // Set up defaults.
         epochDurationInSeconds = 2 weeks;
         rewardDelegatedStakeWeight = (90 * PPM_DENOMINATOR) / 100; // 90%

contracts/staking/contracts/src/sys/MixinScheduler.sol

@@ -75,9 +75,9 @@ contract MixinScheduler is
         return getCurrentEpochStartTimeInSeconds().safeAdd(epochDurationInSeconds);
     }
 
-    /// @dev Initializes state owned by this mixin.
+    /// @dev Assert scheduler state before initializing it.
-    ///      Fails if state was already initialized.
+    /// This must be updated for each migration.
-    function _initMixinScheduler()
+    function _assertMixinSchedulerBeforeInit()
         internal
     {
         if (currentEpochStartTimeInSeconds != 0) {
@@ -87,6 +87,16 @@ contract MixinScheduler is
                 )
             );
         }
+    }
+
+    /// @dev Initializes state owned by this mixin.
+    ///      Fails if state was already initialized.
+    function _initMixinScheduler()
+        internal
+    {
+        // assert the current values before overwriting them.
+        _assertMixinSchedulerBeforeInit();
+
         // solhint-disable-next-line
         currentEpochStartTimeInSeconds = block.timestamp;
     }

contracts/staking/contracts/test/TestCumulativeRewardTracking.sol

@@ -25,7 +25,7 @@ contract TestCumulativeRewardTracking is
     TestStaking
 {
 
-    event SetMostRecentCumulativeReward(
+    event SetCumulativeReward(
         bytes32 poolId,
         uint256 epoch
     );
@@ -35,29 +35,41 @@ contract TestCumulativeRewardTracking is
         uint256 epoch
     );
 
-    event SetCumulativeReward(
+    event SetMostRecentCumulativeReward(
         bytes32 poolId,
         uint256 epoch
     );
 
-    function _setMostRecentCumulativeReward(bytes32 poolId, uint256 epoch)
+    function _forceSetCumulativeReward(
-        internal
+        bytes32 poolId,
-    {
+        uint256 epoch,
-        emit SetMostRecentCumulativeReward(poolId, epoch);
+        IStructs.Fraction memory value
-        MixinCumulativeRewards._setMostRecentCumulativeReward(poolId, epoch);
+    )
-    }
-
-    function _unsetCumulativeReward(bytes32 poolId, uint256 epoch)
-        internal
-    {
-        emit UnsetCumulativeReward(poolId, epoch);
-        MixinCumulativeRewards._unsetCumulativeReward(poolId, epoch);
-    }
-
-    function _setCumulativeReward(bytes32 poolId, uint256 epoch, IStructs.Fraction memory value)
         internal
     {
         emit SetCumulativeReward(poolId, epoch);
-        MixinCumulativeRewards._setCumulativeReward(poolId, epoch, value);
+        MixinCumulativeRewards._forceSetCumulativeReward(poolId, epoch, value);
     }
+
+    function _forceUnsetCumulativeReward(bytes32 poolId, uint256 epoch)
+        internal
+    {
+        emit UnsetCumulativeReward(poolId, epoch);
+        MixinCumulativeRewards._forceUnsetCumulativeReward(poolId, epoch);
+    }
+
+    function _forceSetMostRecentCumulativeRewardEpoch(bytes32 poolId, uint256 epoch)
+        internal
+    {
+        emit SetMostRecentCumulativeReward(poolId, epoch);
+        MixinCumulativeRewards._forceSetMostRecentCumulativeRewardEpoch(poolId, epoch);
+    }
+
+    function _assertMixinParamsBeforeInit()
+        internal
+    {}
+
+    function _assertMixinSchedulerBeforeInit()
+        internal
+    {}
 }

contracts/staking/test/cumulative_reward_tracking_test.ts

@@ -34,11 +34,18 @@ blockchainTests.resets('Cumulative Reward Tracking', env => {
     });
 
     describe('Tracking Cumulative Rewards (CR)', () => {
-        it('should set CR and Most Recent CR when a pool is created', async () => {
+        it('should set CR hen a pool is created is epoch 0', async () => {
             await simulation.runTestAsync(
                 [],
                 [TestAction.CreatePool],
-                [{ event: 'SetCumulativeReward', epoch: 0 }, { event: 'SetMostRecentCumulativeReward', epoch: 0 }],
+                [{ event: 'SetCumulativeReward', epoch: 0 }],
+            );
+        });
+        it('should set CR and Most Recent CR when a pool is created in epoch >0', async () => {
+            await simulation.runTestAsync(
+                [TestAction.Finalize],
+                [TestAction.CreatePool],
+                [{ event: 'SetCumulativeReward', epoch: 1 }, { event: 'SetMostRecentCumulativeReward', epoch: 1 }],
             );
         });
         it('should not set CR or Most Recent CR when values already exist for the current epoch', async () => {

contracts/staking/test/utils/cumulative_reward_tracking_simulation.ts

@@ -37,20 +37,20 @@ export class CumulativeRewardTrackingSimulation {
     private static _extractTestLogs(txReceiptLogs: DecodedLogArgs[]): TestLog[] {
         const logs = [];
         for (const log of txReceiptLogs) {
-            if ((log as any).event === 'SetMostRecentCumulativeReward') {
+            if ((log as DecodedLogArgs).event === 'SetMostRecentCumulativeReward') {
                 logs.push({
                     event: 'SetMostRecentCumulativeReward',
-                    epoch: (log as any).args.epoch,
+                    epoch: (log as DecodedLogArgs).args.epoch.toNumber(),
                 });
-            } else if ((log as any).event === 'SetCumulativeReward') {
+            } else if ((log as DecodedLogArgs).event === 'SetCumulativeReward') {
                 logs.push({
                     event: 'SetCumulativeReward',
-                    epoch: (log as any).args.epoch,
+                    epoch: (log as DecodedLogArgs).args.epoch.toNumber(),
                 });
-            } else if ((log as any).event === 'UnsetCumulativeReward') {
+            } else if ((log as DecodedLogArgs).event === 'UnsetCumulativeReward') {
                 logs.push({
                     event: 'UnsetCumulativeReward',
-                    epoch: (log as any).args.epoch,
+                    epoch: (log as DecodedLogArgs).args.epoch.toNumber(),
                 });
             }
         }
@@ -159,7 +159,7 @@ export class CumulativeRewardTrackingSimulation {
                         { from: this._poolOperator },
                     );
                     const createStakingPoolLog = txReceipt.logs[0];
-                    this._poolId = (createStakingPoolLog as any).args.poolId;
+                    this._poolId = (createStakingPoolLog as DecodedLogArgs).args.poolId;
                     break;
 
                 default: