@0x/contracts-staking: Fully implement MBF (I hope).

contracts/staking/contracts/src/fees/MixinExchangeFees.sol

@@ -62,19 +62,7 @@ contract MixinExchangeFees is
         payable
         onlyExchange
     {
-        // If the protocol fee payment is invalid, revert with a rich error.
+        _assertValidProtocolFee(protocolFeePaid);
-        if (
-            protocolFeePaid == 0 ||
-            (msg.value != protocolFeePaid && msg.value != 0)
-        ) {
-            LibRichErrors.rrevert(LibStakingRichErrors.InvalidProtocolFeePaymentError(
-                protocolFeePaid == 0 ?
-                    LibStakingRichErrors.ProtocolFeePaymentErrorCodes.ZeroProtocolFeePaid :
-                    LibStakingRichErrors.ProtocolFeePaymentErrorCodes.MismatchedFeeAndPayment,
-                protocolFeePaid,
-                msg.value
-            ));
-        }
 
         // Transfer the protocol fee to this address if it should be paid in WETH.
         if (msg.value == 0) {
@@ -88,25 +76,60 @@ contract MixinExchangeFees is
 
         // Get the pool id of the maker address.
         bytes32 poolId = getStakingPoolIdOfMaker(makerAddress);
-
+        // Only attribute the protocol fee payment to a pool if the maker is
-        // Only attribute the protocol fee payment to a pool if the maker is registered to a pool.
+        // registered to a pool.
-        if (poolId != NIL_POOL_ID) {
+        if (poolId == NIL_POOL_ID) {
-            uint256 poolStake = getTotalStakeDelegatedToPool(poolId).currentEpochBalance;
+            return;
-            // Ignore pools with dust stake.
-            if (poolStake >= minimumPoolStake) {
-                // Credit the pool.
-                uint256 _feesCollectedThisEpoch = protocolFeesThisEpochByPool[poolId];
-                protocolFeesThisEpochByPool[poolId] = _feesCollectedThisEpoch.safeAdd(protocolFeePaid);
-                if (_feesCollectedThisEpoch == 0) {
-                    activePoolsThisEpoch.push(poolId);
-                }
-            }
         }
+        uint256 poolStake = getTotalStakeDelegatedToPool(poolId).currentEpochBalance;
+        // Ignore pools with dust stake.
+        if (poolStake < minimumPoolStake) {
+            return;
+        }
+        // Look up the pool for this epoch. The epoch index is `currentEpoch % 2`
+        // because we only need to remember state in the current epoch and the
+        // epoch prior.
+        uint256 currentEpoch = getCurrentEpoch();
+        mapping (bytes32 => IStructs.ActivePool) activePoolsThisEpoch =
+            activePoolsByEpoch[currentEpoch % 2];
+        IStructs.ActivePool memory pool = activePoolsThisEpoch[poolId]
+        // If the pool was previously inactive in this epoch, initialize it.
+        if (pool.feesCollected) {
+            // Compute weighted stake.
+            uint256 operatorStake = getStakeDelegatedToPoolByOwner(
+                rewardVault.operatorOf(poolId),
+                poolId
+            ).currentEpochBalance;
+            pool.weightedStake = operatorStake.safeAdd(
+                totalStakeDelegatedToPool
+                    .safeSub(operatorStake)
+                    .safeMul(delegatedStakeWeight)
+                    .safeDiv(PPM_DENOMINATOR)
+            );
+            // Compute delegated (non-operator) stake.
+            pool.delegatedStake = poolStake.safeSub(operatorStake);
+            // Increase the total weighted stake.
+            totalWeightedStakeThisEpoch = totalWeightedStakeThisEpoch.safeAdd(
+                pool.weightedStake
+            );
+            // Increase the numberof active pools.
+            numActivePoolsThisEpoch += 1;
+            // Emit an event so keepers know what pools to pass into `finalize()`.
+            emit StakingPoolActivated(currentEpoch, poolId);
+        }
+        // Credit the fees to the pool.
+        pool.feesCollected = protocolFeePaid;
+        // Increase the total fees collected this epoch.
+        totalFeesCollectedThisEpoch = totalFeesCollectedThisEpoch.safeAdd(
+            protocolFeePaid
+        );
+        // Store the pool.
+        activePoolsThisEpoch[poolId] = pool;
     }
 
     /// @dev Pays the rebates for to market making pool that was active this epoch,
-    /// then updates the epoch and other time-based periods via the scheduler (see MixinScheduler).
+    /// then updates the epoch and other time-based periods via the scheduler
-    /// This is intentionally permissionless, and may be called by anyone.
+    /// (see MixinScheduler). This is intentionally permissionless, and may be called by anyone.
     function finalizeFees()
         external
     {
@@ -140,15 +163,17 @@ contract MixinExchangeFees is
         return address(this).balance.safeAdd(wethBalance);
     }
 
-    /// @dev Withdraws the entire WETH balance of the contract.
+    /// @dev Checks that the protocol fee passed into `payProtocolFee()` is valid.
-    function _unwrapWETH()
+    /// @param protocolFeePaid The `protocolFeePaid` parameter to `payProtocolFee.`
-        internal
+    function _assertValidProtocolFee(uint256 protocolFeePaid) private view {
-    {
+        if (protocolFeePaid == 0 || (msg.value != protocolFeePaid && msg.value != 0)) {
-        uint256 wethBalance = IEtherToken(WETH_ADDRESS).balanceOf(address(this));
+            LibRichErrors.rrevert(LibStakingRichErrors.InvalidProtocolFeePaymentError(
-
+                protocolFeePaid == 0 ?
-        // Don't withdraw WETH if the WETH balance is zero as a gas optimization.
+                    LibStakingRichErrors.ProtocolFeePaymentErrorCodes.ZeroProtocolFeePaid :
-        if (wethBalance != 0) {
+                    LibStakingRichErrors.ProtocolFeePaymentErrorCodes.MismatchedFeeAndPayment,
-            IEtherToken(WETH_ADDRESS).withdraw(wethBalance);
+                protocolFeePaid,
+                msg.value
+            ));
         }
     }
 

contracts/staking/contracts/src/finalization/MixinFinalizer.sol

@@ -35,12 +35,10 @@ import "./MixinExchangeManager.sol";
 
 
 /// @dev This mixin contains functions related to finalizing epochs.
-///      Finalization occurs over multiple calls because we can only
+///      Finalization occurs AFTER the current epoch is ended/advanced and
-///      discover the `totalRewardsPaid` to all pools by summing the
+///      over (potentially) multiple blocks/transactions. This pattern prevents
-///      the reward function across all active pools at the end of an
+///      the contract from stalling while we finalize rewards for the previous
-///      epoch. Until this value is known for epoch `e`, we cannot finalize
+///      epoch.
-///      epoch `e+1`, because the remaining balance (`balance - totalRewardsPaid`)
-///      is the reward pool for finalizing the next epoch.
 contract MixinFinalizer is
     IStakingEvents,
     MixinConstants,
@@ -68,10 +66,12 @@ contract MixinFinalizer is
         // Make sure the previous epoch has been fully finalized.
         if (unfinalizedPoolsRemaining != 0) {
             LibRichErrors.rrevert(LibStakingRichErrors.PreviousEpochNotFinalized(
-                closingEpoch.sub(1),
+                closingEpoch.safeSub(1),
                 unfinalizedPoolsRemaining
             ));
         }
+        // Unrwap any WETH protocol fees.
+        _unwrapWETH();
         // Populate finalization state.
         unfinalizedPoolsRemaining = numActivePoolsThisEpoch;
         unfinalizedRewardsAvailable = address(this).balance;
@@ -99,8 +99,8 @@ contract MixinFinalizer is
         return _unfinalizedPoolsRemaining = unfinalizedPoolsRemaining;
     }
 
-    /// @dev Finalizes a pool that was active in the previous epoch, paying out
+    /// @dev Finalizes pools that were active in the previous epoch, paying out
-    ///      its rewards to the reward vault. Keepers should call this function
+    ///      rewards to the reward vault. Keepers should call this function
     ///      repeatedly until all active pools that were emitted in in a
     ///      `StakingPoolActivated` in the prior epoch have been finalized.
     ///      Pools that have already been finalized will be silently ignored.
@@ -113,7 +113,7 @@ contract MixinFinalizer is
         external
         returns (uint256 rewardsPaid, uint256 _unfinalizedPoolsRemaining)
     {
-        uint256 epoch = getCurrentEpoch().sub(1);
+        uint256 epoch = getCurrentEpoch().safeSub(1);
         uint256 poolsRemaining = unfinalizedPoolsRemaining;
         uint256 numPoolIds = poolIds.length;
         uint256 rewardsPaid = 0;
@@ -128,12 +128,12 @@ contract MixinFinalizer is
             if (pool.feesCollected != 0) {
                 // Credit the pool with rewards.
                 // We will transfer the total rewards to the vault at the end.
-                rewardsPaid = rewardsPaid.add(_creditRewardsToPool(poolId, pool));
+                rewardsPaid = rewardsPaid.safeAdd(_creditRewardsToPool(poolId, pool));
                 // Clear the pool state so we don't finalize it again,
                 // and to recoup some gas.
                 activePools[poolId] = IStructs.ActivePool(0, 0, 0);
                 // Decrease the number of unfinalized pools left.
-                poolsRemaining = poolsRemaining.sub(1);
+                poolsRemaining = poolsRemaining.safeSub(1);
                 // Emit an event.
                 emit RewardsPaid(epoch, poolId, reward);
             }
@@ -141,7 +141,7 @@ contract MixinFinalizer is
         // Deposit all the rewards at once into the RewardVault.
         _depositIntoStakingPoolRewardVault(rewardsPaid);
         // Update finalization state.
-        totalRewardsPaidLastEpoch = totalRewardsPaidLastEpoch.add(rewardsPaid);
+        totalRewardsPaidLastEpoch = totalRewardsPaidLastEpoch.safeAdd(rewardsPaid);
         _unfinalizedPoolsRemaining = unfinalizedPoolsRemaining = poolsRemaining;
         // If there are no more unfinalized pools remaining, the epoch is
         // finalized.
@@ -149,7 +149,7 @@ contract MixinFinalizer is
             emit EpochFinalized(
                 epoch,
                 totalRewardsPaidLastEpoch,
-                unfinalizedRewardsAvailable.sub(totalRewardsPaidLastEpoch)
+                unfinalizedRewardsAvailable.safeSub(totalRewardsPaidLastEpoch)
             );
         }
     }
@@ -190,12 +190,19 @@ contract MixinFinalizer is
             _recordRewardForDelegators(
                 poolId,
                 membersPortionOfReward,
-                pool.delegatedStake,
+                pool.delegatedStake
-                epoch
             );
         }
     }
 
+    /// @dev Converts the entire WETH balance of the contract into ETH.
+    function _unwrapWETH() private {
+        uint256 wethBalance = IEtherToken(WETH_ADDRESS).balanceOf(address(this));
+        if (wethBalance != 0) {
+            IEtherToken(WETH_ADDRESS).withdraw(wethBalance);
+        }
+    }
+
     /// @dev The cobb-douglas function used to compute fee-based rewards for
     ///      staking pools in a given epoch. Note that in this function there
     ///      is no limitation on alpha; we tend to get better rounding on the

contracts/staking/contracts/src/immutable/MixinStorage.sol

@@ -74,6 +74,12 @@ contract MixinStorage is
     // mapping from Owner to Amount of Withdrawable Stake
     mapping (address => uint256) internal _withdrawableStakeByOwner;
 
+    // Mapping from Owner to Pool Id to epoch of the last rewards collected.
+    // This is the last reward epoch for a pool that a delegator collected
+    // rewards from. This is different from the epoch when the rewards were
+    // collected This will always be `<= currentEpoch`.
+    mapping (address => mapping (bytes32 => uint256)) internal lastCollectedRewardsEpochToPoolByOwner;
+
     // tracking Pool Id
     bytes32 public nextPoolId = INITIAL_POOL_ID;
 
@@ -90,12 +96,6 @@ contract MixinStorage is
     // current epoch start time
     uint256 public currentEpochStartTimeInSeconds;
 
-    // fees collected this epoch
-    mapping (bytes32 => uint256) public protocolFeesThisEpochByPool;
-
-    // pools that were active in the current epoch
-    bytes32[] public activePoolsThisEpoch;
-
     // mapping from Pool Id to Epoch to Reward Ratio
     mapping (bytes32 => mapping (uint256 => IStructs.Fraction)) internal _cumulativeRewardsByPool;
 

contracts/staking/contracts/src/interfaces/IStakingEvents.sol

@@ -43,6 +43,15 @@ interface IStakingEvents {
         address exchangeAddress
     );
 
+    /// @dev Emitted by MixinExchangeFees when a pool pays protocol fees
+    ///      for the first time in an epoch.
+    /// @param epoch The epoch in which the pool was activated.
+    /// @param poolId The ID of the pool.
+    event StakingPoolActivated(
+        uint256 epoch,
+        bytes32 poolId
+    );
+
     /// @dev Emitted by MixinFinalizer when an epoch has ended.
     /// @param epoch The closing epoch.
     /// @param numActivePools Number of active pools in the closing epoch.

contracts/staking/contracts/src/interfaces/IStructs.sol

@@ -39,12 +39,14 @@ interface IStructs {
     /// @param isInitialized
     /// @param currentEpoch the current epoch
     /// @param currentEpochBalance balance in the current epoch.
-    /// @param nextEpochBalance balance in the next epoch.
+    /// @param nextEpochBalance balance in `currentEpoch+1`.
+    /// @param prevEpochBalance balance in `currentEpoch-1`.
     struct StoredBalance {
         bool isInitialized;
         uint32 currentEpoch;
         uint96 currentEpochBalance;
         uint96 nextEpochBalance;
+        uint96 prevEpochBalance;
     }
 
     /// @dev Balance struct for stake.

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

@@ -116,8 +116,7 @@ contract MixinStakingPoolRewards is
     function _handleStakingPoolReward(
         bytes32 poolId,
         uint256 reward,
-        uint256 amountOfDelegatedStake,
+        uint256 amountOfDelegatedStake
-        uint256 epoch
     )
         internal
     {