Chagne StoredBalance functions to not mutate in place

contracts/integrations/test/framework/assertions/moveStake.ts

@@ -32,18 +32,42 @@ function updateNextEpochBalances(
     // Decrement next epoch balances associated with the `from` stake
     if (from.status === StakeStatus.Undelegated) {
         // Decrement owner undelegated stake
-        decreaseNextBalance(ownerStake[StakeStatus.Undelegated], amount, currentEpoch);
+        ownerStake[StakeStatus.Undelegated] = decreaseNextBalance(
+            ownerStake[StakeStatus.Undelegated],
+            amount,
+            currentEpoch,
+        );
         // Decrement global undelegated stake
-        decreaseNextBalance(globalStake[StakeStatus.Undelegated], amount, currentEpoch);
+        globalStake[StakeStatus.Undelegated] = decreaseNextBalance(
+            globalStake[StakeStatus.Undelegated],
+            amount,
+            currentEpoch,
+        );
     } else if (from.status === StakeStatus.Delegated) {
         // Decrement owner's delegated stake to this pool
-        decreaseNextBalance(ownerStake[StakeStatus.Delegated][from.poolId], amount, currentEpoch);
+        ownerStake[StakeStatus.Delegated][from.poolId] = decreaseNextBalance(
+            ownerStake[StakeStatus.Delegated][from.poolId],
+            amount,
+            currentEpoch,
+        );
         // Decrement owner's total delegated stake
-        decreaseNextBalance(ownerStake[StakeStatus.Delegated].total, amount, currentEpoch);
+        ownerStake[StakeStatus.Delegated].total = decreaseNextBalance(
+            ownerStake[StakeStatus.Delegated].total,
+            amount,
+            currentEpoch,
+        );
         // Decrement global delegated stake
-        decreaseNextBalance(globalStake[StakeStatus.Delegated], amount, currentEpoch);
+        globalStake[StakeStatus.Delegated] = decreaseNextBalance(
+            globalStake[StakeStatus.Delegated],
+            amount,
+            currentEpoch,
+        );
         // Decrement pool's delegated stake
-        decreaseNextBalance(stakingPools[from.poolId].delegatedStake, amount, currentEpoch);
+        stakingPools[from.poolId].delegatedStake = decreaseNextBalance(
+            stakingPools[from.poolId].delegatedStake,
+            amount,
+            currentEpoch,
+        );
         updatedPools.push(from.poolId);
 
         // TODO: Check that delegator rewards have been withdrawn/synced
@@ -52,22 +76,46 @@ function updateNextEpochBalances(
     // Increment next epoch balances associated with the `to` stake
     if (to.status === StakeStatus.Undelegated) {
         // Increment owner undelegated stake
-        increaseNextBalance(ownerStake[StakeStatus.Undelegated], amount, currentEpoch);
+        ownerStake[StakeStatus.Undelegated] = increaseNextBalance(
+            ownerStake[StakeStatus.Undelegated],
+            amount,
+            currentEpoch,
+        );
         // Increment global undelegated stake
-        increaseNextBalance(globalStake[StakeStatus.Undelegated], amount, currentEpoch);
+        globalStake[StakeStatus.Undelegated] = increaseNextBalance(
+            globalStake[StakeStatus.Undelegated],
+            amount,
+            currentEpoch,
+        );
     } else if (to.status === StakeStatus.Delegated) {
         // Initializes the balance for this pool if the user has not previously delegated to it
         _.defaults(ownerStake[StakeStatus.Delegated], {
             [to.poolId]: new StoredBalance(),
         });
         // Increment owner's delegated stake to this pool
-        increaseNextBalance(ownerStake[StakeStatus.Delegated][to.poolId], amount, currentEpoch);
+        ownerStake[StakeStatus.Delegated][to.poolId] = increaseNextBalance(
+            ownerStake[StakeStatus.Delegated][to.poolId],
+            amount,
+            currentEpoch,
+        );
         // Increment owner's total delegated stake
-        increaseNextBalance(ownerStake[StakeStatus.Delegated].total, amount, currentEpoch);
+        ownerStake[StakeStatus.Delegated].total = increaseNextBalance(
+            ownerStake[StakeStatus.Delegated].total,
+            amount,
+            currentEpoch,
+        );
         // Increment global delegated stake
-        increaseNextBalance(globalStake[StakeStatus.Delegated], amount, currentEpoch);
+        globalStake[StakeStatus.Delegated] = increaseNextBalance(
+            globalStake[StakeStatus.Delegated],
+            amount,
+            currentEpoch,
+        );
         // Increment pool's delegated stake
-        increaseNextBalance(stakingPools[to.poolId].delegatedStake, amount, currentEpoch);
+        stakingPools[to.poolId].delegatedStake = increaseNextBalance(
+            stakingPools[to.poolId].delegatedStake,
+            amount,
+            currentEpoch,
+        );
         updatedPools.push(to.poolId);
 
         // TODO: Check that delegator rewards have been withdrawn/synced

contracts/integrations/test/framework/assertions/stake.ts

@@ -54,8 +54,16 @@ export function validStakeAssertion(
             balanceStore.assertEquals(expectedBalances);
 
             // _increaseCurrentAndNextBalance
-            increaseCurrentAndNextBalance(ownerStake[StakeStatus.Undelegated], amount, currentEpoch);
-            increaseCurrentAndNextBalance(globalStake[StakeStatus.Undelegated], amount, currentEpoch);
+            ownerStake[StakeStatus.Undelegated] = increaseCurrentAndNextBalance(
+                ownerStake[StakeStatus.Undelegated],
+                amount,
+                currentEpoch,
+            );
+            globalStake[StakeStatus.Undelegated] = increaseCurrentAndNextBalance(
+                globalStake[StakeStatus.Undelegated],
+                amount,
+                currentEpoch,
+            );
 
             // Checks that the owner's undelegated stake has increased by the stake amount
             const ownerUndelegatedStake = await stakingWrapper

contracts/integrations/test/framework/assertions/unstake.ts

@@ -54,8 +54,16 @@ export function validUnstakeAssertion(
             balanceStore.assertEquals(expectedBalances);
 
             // _decreaseCurrentAndNextBalance
-            decreaseCurrentAndNextBalance(ownerStake[StakeStatus.Undelegated], amount, currentEpoch);
-            decreaseCurrentAndNextBalance(globalStake[StakeStatus.Undelegated], amount, currentEpoch);
+            ownerStake[StakeStatus.Undelegated] = decreaseCurrentAndNextBalance(
+                ownerStake[StakeStatus.Undelegated],
+                amount,
+                currentEpoch,
+            );
+            globalStake[StakeStatus.Undelegated] = decreaseCurrentAndNextBalance(
+                globalStake[StakeStatus.Undelegated],
+                amount,
+                currentEpoch,
+            );
 
             // Checks that the owner's undelegated stake has decreased by the stake amount
             const ownerUndelegatedStake = await stakingWrapper

contracts/staking/src/types.ts

@@ -71,56 +71,62 @@ export class StoredBalance {
  * Simulates _loadCurrentBalance. `shouldMutate` flag specifies whether or not to update the given
  * StoredBalance instance.
  */
-export function loadCurrentBalance(
-    balance: StoredBalance,
-    epoch: BigNumber,
-    shouldMutate: boolean = false,
-): StoredBalance {
-    const loadedBalance = new StoredBalance(
+export function loadCurrentBalance(balance: StoredBalance, epoch: BigNumber): StoredBalance {
+    return new StoredBalance(
         epoch,
         epoch.isGreaterThan(balance.currentEpoch) ? balance.nextEpochBalance : balance.currentEpochBalance,
         balance.nextEpochBalance,
     );
-    if (shouldMutate) {
-        balance.currentEpoch = loadedBalance.currentEpoch;
-        balance.currentEpochBalance = loadedBalance.currentEpochBalance;
-        balance.nextEpochBalance = loadedBalance.nextEpochBalance;
-    }
-    return loadedBalance;
 }
 
 /**
  * Simulates _increaseNextBalance
  */
-export function increaseNextBalance(balance: StoredBalance, amount: Numberish, epoch: BigNumber): void {
-    loadCurrentBalance(balance, epoch, true);
-    balance.nextEpochBalance = balance.nextEpochBalance.plus(amount);
+export function increaseNextBalance(balance: StoredBalance, amount: Numberish, epoch: BigNumber): StoredBalance {
+    return {
+        ...loadCurrentBalance(balance, epoch),
+        nextEpochBalance: balance.nextEpochBalance.plus(amount),
+    };
 }
 
 /**
  * Simulates _decreaseNextBalance
  */
-export function decreaseNextBalance(balance: StoredBalance, amount: Numberish, epoch: BigNumber): void {
-    loadCurrentBalance(balance, epoch, true);
-    balance.nextEpochBalance = balance.nextEpochBalance.minus(amount);
+export function decreaseNextBalance(balance: StoredBalance, amount: Numberish, epoch: BigNumber): StoredBalance {
+    return {
+        ...loadCurrentBalance(balance, epoch),
+        nextEpochBalance: balance.nextEpochBalance.minus(amount),
+    };
 }
 
 /**
  * Simulates _increaseCurrentAndNextBalance
  */
-export function increaseCurrentAndNextBalance(balance: StoredBalance, amount: Numberish, epoch: BigNumber): void {
-    loadCurrentBalance(balance, epoch, true);
-    balance.currentEpochBalance = balance.currentEpochBalance.plus(amount);
-    balance.nextEpochBalance = balance.nextEpochBalance.plus(amount);
+export function increaseCurrentAndNextBalance(
+    balance: StoredBalance,
+    amount: Numberish,
+    epoch: BigNumber,
+): StoredBalance {
+    return {
+        ...loadCurrentBalance(balance, epoch),
+        currentEpochBalance: balance.currentEpochBalance.plus(amount),
+        nextEpochBalance: balance.nextEpochBalance.plus(amount),
+    };
 }
 
 /**
  * Simulates _decreaseCurrentAndNextBalance
  */
-export function decreaseCurrentAndNextBalance(balance: StoredBalance, amount: Numberish, epoch: BigNumber): void {
-    loadCurrentBalance(balance, epoch, true);
-    balance.currentEpochBalance = balance.currentEpochBalance.minus(amount);
-    balance.nextEpochBalance = balance.nextEpochBalance.minus(amount);
+export function decreaseCurrentAndNextBalance(
+    balance: StoredBalance,
+    amount: Numberish,
+    epoch: BigNumber,
+): StoredBalance {
+    return {
+        ...loadCurrentBalance(balance, epoch),
+        currentEpochBalance: balance.currentEpochBalance.minus(amount),
+        nextEpochBalance: balance.nextEpochBalance.minus(amount),
+    };
 }
 
 export interface StakeBalanceByPool {