Do not store global inactive state

contracts/staking/contracts/src/ZrxVault.sol

@@ -52,6 +52,24 @@ contract ZrxVault is
     // Asset data for the ERC20 Proxy
     bytes internal _zrxAssetData;
 
+    /// @dev Only stakingProxy can call this function.
+    modifier onlyStakingProxy() {
+        _assertSenderIsStakingProxy();
+        _;
+    }
+
+    /// @dev Function can only be called in catastrophic failure mode.
+    modifier onlyInCatastrophicFailure() {
+        _assertInCatastrophicFailure();
+        _;
+    }
+
+    /// @dev Function can only be called not in catastropic failure mode
+    modifier onlyNotInCatastrophicFailure() {
+        _assertNotInCatastrophicFailure();
+        _;
+    }
+
     /// @dev Constructor.
     /// @param _zrxProxyAddress Address of the 0x Zrx Proxy.
     /// @param _zrxTokenAddress Address of the Zrx Token.
@@ -169,6 +187,15 @@ contract ZrxVault is
         return _balances[staker];
     }
 
+    /// @dev Returns the entire balance of Zrx tokens in the vault.
+    function balanceOfVault()
+        external
+        view
+        returns (uint256)
+    {
+        return _zrxToken.balanceOf(address(this));
+    }
+
     /// @dev Withdraw an `amount` of Zrx Tokens to `staker` from the vault.
     /// @param staker of Zrx Tokens.
     /// @param amount of Zrx Tokens to withdraw.
@@ -190,21 +217,7 @@ contract ZrxVault is
         );
     }
 
-    modifier onlyStakingProxy() {
-        _assertSenderIsStakingProxy();
-        _;
-    }
-
-    modifier onlyInCatastrophicFailure() {
-        _assertInCatastrophicFailure();
-        _;
-    }
-
-    modifier onlyNotInCatastrophicFailure() {
-        _assertNotInCatastrophicFailure();
-        _;
-    }
-
+    /// @dev Asserts that sender is stakingProxy contract.
     function _assertSenderIsStakingProxy()
         private
         view
@@ -216,6 +229,7 @@ contract ZrxVault is
         }
     }
 
+    /// @dev Asserts that vault is in catastrophic failure mode.
     function _assertInCatastrophicFailure()
         private
         view
@@ -225,6 +239,7 @@ contract ZrxVault is
         }
     }
 
+    /// @dev Asserts that vault is not in catastrophic failure mode.
     function _assertNotInCatastrophicFailure()
         private
         view

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

@@ -41,10 +41,9 @@ contract MixinStorage is
     // address for read-only proxy to call
     address public readOnlyProxyCallee;
 
-    // mapping from StakeStatus to the total amount of stake in that status for the entire
-    // staking system.
+    // amount of stake currently being delegated in the system
     // (access using _loadSyncedBalance or _loadUnsyncedBalance)
-    mapping (uint8 => IStructs.StoredBalance) internal _globalStakeByStatus;
+    IStructs.StoredBalance internal _globalDelegatedStake;
 
     // mapping from StakeStatus to address of staker to stored balance
     // (access using _loadSyncedBalance or _loadUnsyncedBalance)

contracts/staking/contracts/src/interfaces/IZrxVault.sol

@@ -95,4 +95,10 @@ interface IZrxVault {
         external
         view
         returns (uint256);
+
+    /// @dev Returns the entire balance of Zrx tokens in the vault.
+    function balanceOfVault()
+        external
+        view
+        returns (uint256);
 }

contracts/staking/contracts/src/stake/MixinStake.sol

@@ -46,12 +46,6 @@ contract MixinStake is
             amount
         );
 
-        // update global total of active stake
-        _increaseCurrentAndNextBalance(
-            _globalStakeByStatus[uint8(IStructs.StakeStatus.INACTIVE)],
-            amount
-        );
-
         // notify
         emit Stake(
             staker,
@@ -92,12 +86,6 @@ contract MixinStake is
             amount
         );
 
-        // update global total of inactive stake
-        _decreaseCurrentAndNextBalance(
-            _globalStakeByStatus[uint8(IStructs.StakeStatus.INACTIVE)],
-            amount
-        );
-
         // withdraw equivalent amount of ZRX from vault
         getZrxVault().withdrawFrom(staker, amount);
 
@@ -149,13 +137,6 @@ contract MixinStake is
             amount
         );
 
-        // update global total of stake in the statuses being moved between
-        _moveStake(
-            _globalStakeByStatus[uint8(from.status)],
-            _globalStakeByStatus[uint8(to.status)],
-            amount
-        );
-
         // notify
         emit MoveStake(
             staker,
@@ -197,6 +178,12 @@ contract MixinStake is
             _delegatedStakeByPoolId[poolId],
             amount
         );
+
+        // Increase next balance of global delegated stake
+        _increaseNextBalance(
+            _globalDelegatedStake,
+            amount
+        );
     }
 
     /// @dev Un-Delegates a owners stake from a staking pool.
@@ -229,5 +216,11 @@ contract MixinStake is
             _delegatedStakeByPoolId[poolId],
             amount
         );
+
+        // decrease next balance of global delegated stake
+        _decreaseNextBalance(
+            _globalDelegatedStake,
+            amount
+        );
     }
 }

contracts/staking/contracts/src/stake/MixinStakeBalances.sol

@@ -30,22 +30,27 @@ contract MixinStakeBalances is
     using LibSafeMath for uint256;
 
     /// @dev Gets global stake for a given status.
-    /// @param stakeStatus ACTIVE, INACTIVE, or DELEGATED
+    /// @param stakeStatus INACTIVE or DELEGATED
     /// @return Global stake for given status.
     function getGlobalStakeByStatus(IStructs.StakeStatus stakeStatus)
         external
         view
         returns (IStructs.StoredBalance memory balance)
     {
-        balance = _loadSyncedBalance(
-            _globalStakeByStatus[uint8(stakeStatus)]
-        );
+        balance = _loadSyncedBalance(_globalDelegatedStake);
+        if (stakeStatus == IStructs.StakeStatus.INACTIVE) {
+            // Inactive stake is the difference between total stake and delegated stake
+            // Note that any Zrx erroneously sent to the vault will be counted as inactive stake
+            uint256 totalStake = getZrxVault().balanceOfVault();
+            balance.currentEpochBalance = totalStake.safeSub(balance.currentEpochBalance).downcastToUint96();
+            balance.nextEpochBalance = totalStake.safeSub(balance.nextEpochBalance).downcastToUint96();
+        }
         return balance;
     }
 
     /// @dev Gets an owner's stake balances by status.
     /// @param staker Owner of stake.
-    /// @param stakeStatus ACTIVE, INACTIVE, or DELEGATED
+    /// @param stakeStatus INACTIVE or DELEGATED
     /// @return Owner's stake balances for given status.
     function getOwnerStakeByStatus(
         address staker,
@@ -65,7 +70,7 @@ contract MixinStakeBalances is
     /// @param staker of stake.
     /// @return Total active stake for staker.
     function getTotalStake(address staker)
-        external
+        public
         view
         returns (uint256)
     {

contracts/staking/contracts/test/TestStorageLayoutAndConstants.sol

@@ -119,8 +119,8 @@ contract TestStorageLayoutAndConstants is
             slot := add(slot, 0x1)
 
             assertSlotAndOffset(
-                _globalStakeByStatus_slot,
-                _globalStakeByStatus_offset,
+                _globalDelegatedStake_slot,
+                _globalDelegatedStake_offset,
                 slot,
                 offset
             )