Merge pull request #2151 from 0xProject/feature/contracts-staking/maintain-global-stake

Maintain global stake

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

@@ -47,6 +47,10 @@ 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.
+    mapping (uint8 => IStructs.StoredBalance) public globalStakeByStatus;
+
     // mapping from Owner to Amount of Active Stake
     // (access using _loadAndSyncBalance or _loadUnsyncedBalance)
     mapping (address => IStructs.StoredBalance) internal _activeStakeByOwner;

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

@@ -49,6 +49,9 @@ contract MixinStake is
         // mint stake
         _incrementCurrentAndNextBalance(_activeStakeByOwner[owner], amount);
 
+        // update global total of active stake
+        _incrementCurrentAndNextBalance(globalStakeByStatus[uint8(IStructs.StakeStatus.ACTIVE)], amount);
+
         // notify
         emit Stake(
             owner,
@@ -78,6 +81,9 @@ contract MixinStake is
         // burn inactive stake
         _decrementCurrentAndNextBalance(_inactiveStakeByOwner[owner], amount);
 
+        // update global total of inactive stake
+        _decrementCurrentAndNextBalance(globalStakeByStatus[uint8(IStructs.StakeStatus.INACTIVE)], amount);
+
         // update withdrawable field
         _withdrawableStakeByOwner[owner] = currentWithdrawableStake.safeSub(amount);
 
@@ -140,6 +146,13 @@ contract MixinStake is
         IStructs.StoredBalance storage toPtr = _getBalancePtrFromStatus(owner, to.status);
         _moveStake(fromPtr, toPtr, amount);
 
+        // update global total of stake in the statuses being moved between
+        _moveStake(
+            globalStakeByStatus[uint8(from.status)],
+            globalStakeByStatus[uint8(to.status)],
+            amount
+        );
+
         // update withdrawable field, if necessary
         if (from.status == IStructs.StakeStatus.INACTIVE) {
             _withdrawableStakeByOwner[owner] = _computeWithdrawableStake(owner, withdrawableStake);

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

@@ -31,6 +31,54 @@ contract MixinStakeBalances is
 {
     using LibSafeMath for uint256;
 
+    /// @dev Returns the total active stake across the entire staking system.
+    /// @return Global active stake.
+    function getGlobalActiveStake()
+        external
+        view
+        returns (IStructs.StakeBalance memory balance)
+    {
+        IStructs.StoredBalance memory storedBalance = _loadAndSyncBalance(
+            globalStakeByStatus[uint8(IStructs.StakeStatus.ACTIVE)]
+        );
+        return IStructs.StakeBalance({
+            currentEpochBalance: storedBalance.currentEpochBalance,
+            nextEpochBalance: storedBalance.nextEpochBalance
+        });
+    }
+
+    /// @dev Returns the total inactive stake across the entire staking system.
+    /// @return Global inactive stake.
+    function getGlobalInactiveStake()
+        external
+        view
+        returns (IStructs.StakeBalance memory balance)
+    {
+        IStructs.StoredBalance memory storedBalance = _loadAndSyncBalance(
+            globalStakeByStatus[uint8(IStructs.StakeStatus.INACTIVE)]
+        );
+        return IStructs.StakeBalance({
+            currentEpochBalance: storedBalance.currentEpochBalance,
+            nextEpochBalance: storedBalance.nextEpochBalance
+        });
+    }
+
+    /// @dev Returns the total stake delegated across the entire staking system.
+    /// @return Global delegated stake.
+    function getGlobalDelegatedStake()
+        external
+        view
+        returns (IStructs.StakeBalance memory balance)
+    {
+        IStructs.StoredBalance memory storedBalance = _loadAndSyncBalance(
+            globalStakeByStatus[uint8(IStructs.StakeStatus.DELEGATED)]
+        );
+        return IStructs.StakeBalance({
+            currentEpochBalance: storedBalance.currentEpochBalance,
+            nextEpochBalance: storedBalance.nextEpochBalance
+        });
+    }
+
     /// @dev Returns the total stake for a given owner.
     /// @param owner of stake.
     /// @return Total active stake for owner.

contracts/staking/test/actors/staker_actor.ts

@@ -3,13 +3,28 @@ import { BigNumber, RevertError } from '@0x/utils';
 import * as _ from 'lodash';
 
 import { StakingApiWrapper } from '../utils/api_wrapper';
-import { StakeBalances, StakeInfo, StakeStatus } from '../utils/types';
+import { StakeBalance, StakeBalances, StakeInfo, StakeStatus } from '../utils/types';
 
 import { BaseActor } from './base_actor';
 
 export class StakerActor extends BaseActor {
     private readonly _poolIds: string[];
 
+    private static _incrementNextBalance(balance: StakeBalance, amount: BigNumber): void {
+        balance.nextEpochBalance = balance.nextEpochBalance.plus(amount);
+    }
+    private static _decrementNextBalance(balance: StakeBalance, amount: BigNumber): void {
+        balance.nextEpochBalance = balance.nextEpochBalance.minus(amount);
+    }
+    private static _incrementCurrentAndNextBalance(balance: StakeBalance, amount: BigNumber): void {
+        balance.currentEpochBalance = balance.currentEpochBalance.plus(amount);
+        balance.nextEpochBalance = balance.nextEpochBalance.plus(amount);
+    }
+    private static _decrementCurrentAndNextBalance(balance: StakeBalance, amount: BigNumber): void {
+        balance.currentEpochBalance = balance.currentEpochBalance.minus(amount);
+        balance.nextEpochBalance = balance.nextEpochBalance.minus(amount);
+    }
+
     constructor(owner: string, stakingApiWrapper: StakingApiWrapper) {
         super(owner, stakingApiWrapper);
         this._poolIds = [];
@@ -22,7 +37,7 @@ export class StakerActor extends BaseActor {
         revertError?: RevertError,
     ): Promise<void> {
         const initZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
-        const initStakerBalances = await this.getBalancesAsync();
+        const initBalances = await this._getBalancesAsync();
         // move stake
         const txReceiptPromise = this._stakingApiWrapper.stakingProxyContract.batchExecute.awaitTransactionSuccessAsync(
             [
@@ -37,13 +52,13 @@ export class StakerActor extends BaseActor {
         }
         await txReceiptPromise;
         // Calculate the expected stake amount.
-        const expectedStakerBalances = await this._calculateExpectedBalancesAfterMoveAsync(
+        const expectedBalances = await this._calculateExpectedBalancesAfterMoveAsync(
             from,
             to,
             amount,
-            await this._calculateExpectedBalancesAfterStakeAsync(amount, initStakerBalances),
+            await this._calculateExpectedBalancesAfterStakeAsync(amount, initBalances),
         );
-        await this.assertBalancesAsync(expectedStakerBalances);
+        await this._assertBalancesAsync(expectedBalances);
         // check zrx balance of vault
         const finalZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
         expect(finalZrxBalanceOfVault, 'final balance of zrx vault').to.be.bignumber.equal(
@@ -53,7 +68,7 @@ export class StakerActor extends BaseActor {
 
     public async stakeAsync(amount: BigNumber, revertError?: RevertError): Promise<void> {
         const initZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
-        const initStakerBalances = await this.getBalancesAsync();
+        const initBalances = await this._getBalancesAsync();
         // deposit stake
         const txReceiptPromise = this._stakingApiWrapper.stakingContract.stake.awaitTransactionSuccessAsync(amount, {
             from: this._owner,
@@ -65,8 +80,8 @@ export class StakerActor extends BaseActor {
         await txReceiptPromise;
         // @TODO check receipt logs and return value via eth_call
         // check balances
-        const expectedStakerBalances = await this._calculateExpectedBalancesAfterStakeAsync(amount, initStakerBalances);
+        const expectedBalances = await this._calculateExpectedBalancesAfterStakeAsync(amount, initBalances);
-        await this.assertBalancesAsync(expectedStakerBalances);
+        await this._assertBalancesAsync(expectedBalances);
         // check zrx balance of vault
         const finalZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
         expect(finalZrxBalanceOfVault, 'final balance of zrx vault').to.be.bignumber.equal(
@@ -76,7 +91,7 @@ export class StakerActor extends BaseActor {
 
     public async unstakeAsync(amount: BigNumber, revertError?: RevertError): Promise<void> {
         const initZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
-        const initStakerBalances = await this.getBalancesAsync();
+        const initBalances = await this._getBalancesAsync();
         // deposit stake
         const txReceiptPromise = this._stakingApiWrapper.stakingContract.unstake.awaitTransactionSuccessAsync(amount, {
             from: this._owner,
@@ -88,17 +103,13 @@ export class StakerActor extends BaseActor {
         await txReceiptPromise;
         // @TODO check receipt logs and return value via eth_call
         // check balances
-        const expectedStakerBalances = initStakerBalances;
+        const expectedBalances = initBalances;
-        expectedStakerBalances.zrxBalance = initStakerBalances.zrxBalance.plus(amount);
+        expectedBalances.zrxBalance = initBalances.zrxBalance.plus(amount);
-        expectedStakerBalances.stakeBalanceInVault = initStakerBalances.stakeBalanceInVault.minus(amount);
+        expectedBalances.stakeBalanceInVault = initBalances.stakeBalanceInVault.minus(amount);
-        expectedStakerBalances.inactiveStakeBalance.nextEpochBalance = initStakerBalances.inactiveStakeBalance.nextEpochBalance.minus(
+        StakerActor._decrementCurrentAndNextBalance(expectedBalances.inactiveStakeBalance, amount);
-            amount,
+        StakerActor._decrementCurrentAndNextBalance(expectedBalances.globalInactiveStakeBalance, amount);
-        );
+        expectedBalances.withdrawableStakeBalance = initBalances.withdrawableStakeBalance.minus(amount);
-        expectedStakerBalances.inactiveStakeBalance.currentEpochBalance = initStakerBalances.inactiveStakeBalance.currentEpochBalance.minus(
+        await this._assertBalancesAsync(expectedBalances);
-            amount,
-        );
-        expectedStakerBalances.withdrawableStakeBalance = initStakerBalances.withdrawableStakeBalance.minus(amount);
-        await this.assertBalancesAsync(expectedStakerBalances);
         // check zrx balance of vault
         const finalZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
         expect(finalZrxBalanceOfVault, 'final balance of zrx vault').to.be.bignumber.equal(
@@ -115,7 +126,7 @@ export class StakerActor extends BaseActor {
         // Cache Initial Balances.
         const initZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
         // Calculate the expected outcome after the move.
-        const expectedStakerBalances = await this._calculateExpectedBalancesAfterMoveAsync(from, to, amount);
+        const expectedBalances = await this._calculateExpectedBalancesAfterMoveAsync(from, to, amount);
         // move stake
         const txReceiptPromise = this._stakingApiWrapper.stakingContract.moveStake.awaitTransactionSuccessAsync(
             from,
@@ -129,7 +140,7 @@ export class StakerActor extends BaseActor {
         }
         await txReceiptPromise;
         // check balances
-        await this.assertBalancesAsync(expectedStakerBalances);
+        await this._assertBalancesAsync(expectedBalances);
         // check zrx balance of vault
         const finalZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
         expect(finalZrxBalanceOfVault, 'final balance of zrx vault').to.be.bignumber.equal(initZrxBalanceOfVault);
@@ -138,37 +149,40 @@ export class StakerActor extends BaseActor {
     public async goToNextEpochAsync(): Promise<void> {
         // cache balances
         const initZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
-        const initStakerBalances = await this.getBalancesAsync();
+        const initBalances = await this._getBalancesAsync();
         // go to next epoch
         await this._stakingApiWrapper.utils.skipToNextEpochAsync();
         // check balances
-        const expectedStakerBalances = this.getNextEpochBalances(initStakerBalances);
+        const expectedBalances = this._getNextEpochBalances(initBalances);
-        await this.assertBalancesAsync(expectedStakerBalances);
+        await this._assertBalancesAsync(expectedBalances);
         // check zrx balance of vault
         const finalZrxBalanceOfVault = await this._stakingApiWrapper.utils.getZrxTokenBalanceOfZrxVaultAsync();
         expect(finalZrxBalanceOfVault, 'final balance of zrx vault').to.be.bignumber.equal(initZrxBalanceOfVault);
     }
-
+    private _getNextEpochBalances(balances: StakeBalances): StakeBalances {
-    public getNextEpochBalances(balances: StakeBalances): StakeBalances {
         const nextBalances = _.cloneDeep(balances);
         nextBalances.withdrawableStakeBalance = nextBalances.inactiveStakeBalance.nextEpochBalance.isLessThan(
             nextBalances.inactiveStakeBalance.currentEpochBalance,
         )
             ? nextBalances.inactiveStakeBalance.nextEpochBalance
             : nextBalances.inactiveStakeBalance.currentEpochBalance;
-        nextBalances.activeStakeBalance.currentEpochBalance = nextBalances.activeStakeBalance.nextEpochBalance;
+
-        nextBalances.inactiveStakeBalance.currentEpochBalance = nextBalances.inactiveStakeBalance.nextEpochBalance;
+        for (const balance of [
-        nextBalances.delegatedStakeBalance.currentEpochBalance = nextBalances.delegatedStakeBalance.nextEpochBalance;
+            nextBalances.activeStakeBalance,
-        for (const poolId of this._poolIds) {
+            nextBalances.inactiveStakeBalance,
-            nextBalances.delegatedStakeByPool[poolId].currentEpochBalance =
+            nextBalances.delegatedStakeBalance,
-                nextBalances.delegatedStakeByPool[poolId].nextEpochBalance;
+            nextBalances.globalActiveStakeBalance,
-            nextBalances.totalDelegatedStakeByPool[poolId].currentEpochBalance =
+            nextBalances.globalInactiveStakeBalance,
-                nextBalances.totalDelegatedStakeByPool[poolId].nextEpochBalance;
+            nextBalances.globalDelegatedStakeBalance,
+            ...this._poolIds.map(poolId => nextBalances.delegatedStakeByPool[poolId]),
+            ...this._poolIds.map(poolId => nextBalances.totalDelegatedStakeByPool[poolId]),
+        ]) {
+            balance.currentEpochBalance = balance.nextEpochBalance;
         }
         return nextBalances;
     }
-    public async getBalancesAsync(): Promise<StakeBalances> {
+    private async _getBalancesAsync(): Promise<StakeBalances> {
-        const stakerBalances: StakeBalances = {
+        const balances: StakeBalances = {
             zrxBalance: await this._stakingApiWrapper.zrxTokenContract.balanceOf.callAsync(this._owner),
             stakeBalance: await this._stakingApiWrapper.stakingContract.getTotalStake.callAsync(this._owner),
             stakeBalanceInVault: await this._stakingApiWrapper.zrxVaultContract.balanceOf.callAsync(this._owner),
@@ -180,6 +194,9 @@ export class StakerActor extends BaseActor {
             delegatedStakeBalance: await this._stakingApiWrapper.stakingContract.getStakeDelegatedByOwner.callAsync(
                 this._owner,
             ),
+            globalActiveStakeBalance: await this._stakingApiWrapper.stakingContract.getGlobalActiveStake.callAsync(),
+            globalInactiveStakeBalance: await this._stakingApiWrapper.stakingContract.getGlobalInactiveStake.callAsync(),
+            globalDelegatedStakeBalance: await this._stakingApiWrapper.stakingContract.getGlobalDelegatedStake.callAsync(),
             delegatedStakeByPool: {},
             totalDelegatedStakeByPool: {},
         };
@@ -192,13 +209,13 @@ export class StakerActor extends BaseActor {
             const totalDelegatedStakeBalanceByPool = await this._stakingApiWrapper.stakingContract.getTotalStakeDelegatedToPool.callAsync(
                 poolId,
             );
-            stakerBalances.delegatedStakeByPool[poolId] = delegatedStakeBalanceByPool;
+            balances.delegatedStakeByPool[poolId] = delegatedStakeBalanceByPool;
-            stakerBalances.totalDelegatedStakeByPool[poolId] = totalDelegatedStakeBalanceByPool;
+            balances.totalDelegatedStakeByPool[poolId] = totalDelegatedStakeBalanceByPool;
         }
-        return stakerBalances;
+        return balances;
     }
-    public async assertBalancesAsync(expectedBalances: StakeBalances): Promise<void> {
+    private async _assertBalancesAsync(expectedBalances: StakeBalances): Promise<void> {
-        const balances = await this.getBalancesAsync();
+        const balances = await this._getBalancesAsync();
         expect(balances.zrxBalance, 'zrx balance').to.be.bignumber.equal(expectedBalances.zrxBalance);
         expect(balances.stakeBalanceInVault, 'stake balance, recorded in vault').to.be.bignumber.equal(
             expectedBalances.stakeBalanceInVault,
@@ -226,6 +243,28 @@ export class StakerActor extends BaseActor {
         expect(balances.delegatedStakeBalance.nextEpochBalance, 'delegated stake balance (next)').to.be.bignumber.equal(
             expectedBalances.delegatedStakeBalance.nextEpochBalance,
         );
+        expect(
+            balances.globalActiveStakeBalance.currentEpochBalance,
+            'global active stake (current)',
+        ).to.bignumber.equal(expectedBalances.globalActiveStakeBalance.currentEpochBalance);
+        expect(
+            balances.globalInactiveStakeBalance.currentEpochBalance,
+            'global inactive stake (current)',
+        ).to.bignumber.equal(expectedBalances.globalInactiveStakeBalance.currentEpochBalance);
+        expect(
+            balances.globalDelegatedStakeBalance.currentEpochBalance,
+            'global delegated stake (current)',
+        ).to.bignumber.equal(expectedBalances.globalDelegatedStakeBalance.currentEpochBalance);
+        expect(balances.globalActiveStakeBalance.nextEpochBalance, 'global active stake (next)').to.bignumber.equal(
+            expectedBalances.globalActiveStakeBalance.nextEpochBalance,
+        );
+        expect(balances.globalInactiveStakeBalance.nextEpochBalance, 'global inactive stake (next)').to.bignumber.equal(
+            expectedBalances.globalInactiveStakeBalance.nextEpochBalance,
+        );
+        expect(
+            balances.globalDelegatedStakeBalance.nextEpochBalance,
+            'global delegated stake (next)',
+        ).to.bignumber.equal(expectedBalances.globalDelegatedStakeBalance.nextEpochBalance);
         expect(balances.delegatedStakeByPool, 'delegated stake by pool').to.be.deep.equal(
             expectedBalances.delegatedStakeByPool,
         );
@@ -233,99 +272,67 @@ export class StakerActor extends BaseActor {
             expectedBalances.totalDelegatedStakeByPool,
         );
     }
-    public async forceBalanceSyncAsync(): Promise<void> {
-        const initBalances = await this.getBalancesAsync();
-        await this._stakingApiWrapper.stakingContract.stake.awaitTransactionSuccessAsync(new BigNumber(0), {
-            from: this._owner,
-        });
-        await this.assertBalancesAsync(initBalances);
-    }
 
     private async _calculateExpectedBalancesAfterMoveAsync(
         from: StakeInfo,
         to: StakeInfo,
         amount: BigNumber,
-        initStakerBalances?: StakeBalances,
+        initBalances?: StakeBalances,
     ): Promise<StakeBalances> {
         // check if we're moving stake into a new pool
         if (to.status === StakeStatus.Delegated && to.poolId !== undefined && !_.includes(this._poolIds, to.poolId)) {
             this._poolIds.push(to.poolId);
         }
         // cache balances
-        const initialStakerBalances = initStakerBalances || (await this.getBalancesAsync());
+        const expectedBalances = initBalances || (await this._getBalancesAsync());
         // @TODO check receipt logs and return value via eth_call
         // check balances
-        const expectedStakerBalances = initialStakerBalances;
         // from
         if (from.status === StakeStatus.Active) {
-            expectedStakerBalances.activeStakeBalance.nextEpochBalance = initialStakerBalances.activeStakeBalance.nextEpochBalance.minus(
+            StakerActor._decrementNextBalance(expectedBalances.activeStakeBalance, amount);
-                amount,
+            StakerActor._decrementNextBalance(expectedBalances.globalActiveStakeBalance, amount);
-            );
         } else if (from.status === StakeStatus.Inactive) {
-            expectedStakerBalances.inactiveStakeBalance.nextEpochBalance = initialStakerBalances.inactiveStakeBalance.nextEpochBalance.minus(
+            StakerActor._decrementNextBalance(expectedBalances.inactiveStakeBalance, amount);
-                amount,
+            StakerActor._decrementNextBalance(expectedBalances.globalInactiveStakeBalance, amount);
-            );
             if (
-                expectedStakerBalances.inactiveStakeBalance.nextEpochBalance.isLessThan(
+                expectedBalances.inactiveStakeBalance.nextEpochBalance.isLessThan(
-                    expectedStakerBalances.withdrawableStakeBalance,
+                    expectedBalances.withdrawableStakeBalance,
                 )
             ) {
-                expectedStakerBalances.withdrawableStakeBalance =
+                expectedBalances.withdrawableStakeBalance = expectedBalances.inactiveStakeBalance.nextEpochBalance;
-                    expectedStakerBalances.inactiveStakeBalance.nextEpochBalance;
             }
         } else if (from.status === StakeStatus.Delegated && from.poolId !== undefined) {
-            expectedStakerBalances.delegatedStakeBalance.nextEpochBalance = initialStakerBalances.delegatedStakeBalance.nextEpochBalance.minus(
+            StakerActor._decrementNextBalance(expectedBalances.delegatedStakeBalance, amount);
-                amount,
+            StakerActor._decrementNextBalance(expectedBalances.globalDelegatedStakeBalance, amount);
-            );
+            StakerActor._decrementNextBalance(expectedBalances.delegatedStakeByPool[from.poolId], amount);
-            expectedStakerBalances.delegatedStakeByPool[
+            StakerActor._decrementNextBalance(expectedBalances.totalDelegatedStakeByPool[from.poolId], amount);
-                from.poolId
-            ].nextEpochBalance = initialStakerBalances.delegatedStakeByPool[from.poolId].nextEpochBalance.minus(amount);
-            expectedStakerBalances.totalDelegatedStakeByPool[
-                from.poolId
-            ].nextEpochBalance = initialStakerBalances.totalDelegatedStakeByPool[from.poolId].nextEpochBalance.minus(
-                amount,
-            );
         }
         // to
         if (to.status === StakeStatus.Active) {
-            expectedStakerBalances.activeStakeBalance.nextEpochBalance = initialStakerBalances.activeStakeBalance.nextEpochBalance.plus(
+            StakerActor._incrementNextBalance(expectedBalances.activeStakeBalance, amount);
-                amount,
+            StakerActor._incrementNextBalance(expectedBalances.globalActiveStakeBalance, amount);
-            );
         } else if (to.status === StakeStatus.Inactive) {
-            expectedStakerBalances.inactiveStakeBalance.nextEpochBalance = initialStakerBalances.inactiveStakeBalance.nextEpochBalance.plus(
+            StakerActor._incrementNextBalance(expectedBalances.inactiveStakeBalance, amount);
-                amount,
+            StakerActor._incrementNextBalance(expectedBalances.globalInactiveStakeBalance, amount);
-            );
         } else if (to.status === StakeStatus.Delegated && to.poolId !== undefined) {
-            expectedStakerBalances.delegatedStakeBalance.nextEpochBalance = initialStakerBalances.delegatedStakeBalance.nextEpochBalance.plus(
+            StakerActor._incrementNextBalance(expectedBalances.delegatedStakeBalance, amount);
-                amount,
+            StakerActor._incrementNextBalance(expectedBalances.globalDelegatedStakeBalance, amount);
-            );
+            StakerActor._incrementNextBalance(expectedBalances.delegatedStakeByPool[to.poolId], amount);
-            expectedStakerBalances.delegatedStakeByPool[
+            StakerActor._incrementNextBalance(expectedBalances.totalDelegatedStakeByPool[to.poolId], amount);
-                to.poolId
-            ].nextEpochBalance = initialStakerBalances.delegatedStakeByPool[to.poolId].nextEpochBalance.plus(amount);
-            expectedStakerBalances.totalDelegatedStakeByPool[
-                to.poolId
-            ].nextEpochBalance = initialStakerBalances.totalDelegatedStakeByPool[to.poolId].nextEpochBalance.plus(
-                amount,
-            );
         }
-        return expectedStakerBalances;
+        return expectedBalances;
     }
 
     private async _calculateExpectedBalancesAfterStakeAsync(
         amount: BigNumber,
-        initStakerBalances?: StakeBalances,
+        initBalances?: StakeBalances,
     ): Promise<StakeBalances> {
-        const initialStakerBalances = initStakerBalances || (await this.getBalancesAsync());
+        const expectedBalances = initBalances || (await this._getBalancesAsync());
         // check balances
-        const expectedStakerBalances = initialStakerBalances;
+        expectedBalances.zrxBalance = expectedBalances.zrxBalance.minus(amount);
-        expectedStakerBalances.zrxBalance = initialStakerBalances.zrxBalance.minus(amount);
+        expectedBalances.stakeBalanceInVault = expectedBalances.stakeBalanceInVault.plus(amount);
-        expectedStakerBalances.stakeBalanceInVault = initialStakerBalances.stakeBalanceInVault.plus(amount);
+        StakerActor._incrementCurrentAndNextBalance(expectedBalances.activeStakeBalance, amount);
-        expectedStakerBalances.activeStakeBalance.currentEpochBalance = initialStakerBalances.activeStakeBalance.currentEpochBalance.plus(
+        StakerActor._incrementCurrentAndNextBalance(expectedBalances.globalActiveStakeBalance, amount);
-            amount,
+        return expectedBalances;
-        );
-        expectedStakerBalances.activeStakeBalance.nextEpochBalance = initialStakerBalances.activeStakeBalance.nextEpochBalance.plus(
-            amount,
-        );
-        return expectedStakerBalances;
     }
 }

contracts/staking/test/utils/types.ts

@@ -86,6 +86,9 @@ export interface StakeBalances {
     activeStakeBalance: StakeBalance;
     inactiveStakeBalance: StakeBalance;
     delegatedStakeBalance: StakeBalance;
+    globalActiveStakeBalance: StakeBalance;
+    globalInactiveStakeBalance: StakeBalance;
+    globalDelegatedStakeBalance: StakeBalance;
     delegatedStakeByPool: StakeBalanceByPool;
     totalDelegatedStakeByPool: StakeBalanceByPool;
 }