added totalStake and totalValidatorStake along with unit tests

contracts/artifacts/abi/LightClient.json

@@ -1256,11 +1256,6 @@
     "name": "InvalidProof",
     "inputs": []
   },
-  {
-    "type": "error",
-    "name": "InvalidScalar",
-    "inputs": []
-  },
   {
     "type": "error",
     "name": "NoChangeRequired",

contracts/artifacts/abi/LightClientMock.json

@@ -1336,11 +1336,6 @@
     "name": "InvalidProof",
     "inputs": []
   },
-  {
-    "type": "error",
-    "name": "InvalidScalar",
-    "inputs": []
-  },
   {
     "type": "error",
     "name": "NoChangeRequired",

contracts/artifacts/abi/LightClientV2.json

@@ -1777,11 +1777,6 @@
     "name": "InvalidProof",
     "inputs": []
   },
-  {
-    "type": "error",
-    "name": "InvalidScalar",
-    "inputs": []
-  },
   {
     "type": "error",
     "name": "MissingEpochRootUpdate",

contracts/artifacts/abi/LightClientV2Mock.json

@@ -1923,11 +1923,6 @@
     "name": "InvalidProof",
     "inputs": []
   },
-  {
-    "type": "error",
-    "name": "InvalidScalar",
-    "inputs": []
-  },
   {
     "type": "error",
     "name": "MissingEpochRootUpdate",

contracts/artifacts/abi/StakeTable.json

@@ -819,11 +819,6 @@
     "name": "BLSSigVerificationFailed",
     "inputs": []
   },
-  {
-    "type": "error",
-    "name": "BN254PairingProdFailed",
-    "inputs": []
-  },
   {
     "type": "error",
     "name": "BlsKeyAlreadyUsed",
@@ -887,11 +882,6 @@
     "name": "InvalidCommission",
     "inputs": []
   },
-  {
-    "type": "error",
-    "name": "InvalidG1",
-    "inputs": []
-  },
   {
     "type": "error",
     "name": "InvalidInitialization",
@@ -934,11 +924,6 @@
       }
     ]
   },
-  {
-    "type": "error",
-    "name": "PowPrecompileFailed",
-    "inputs": []
-  },
   {
     "type": "error",
     "name": "PrematureWithdrawal",

contracts/src/StakeTableV2.sol

@@ -57,6 +57,11 @@ import { SafeTransferLib } from "solmate/utils/SafeTransferLib.sol";
 ///
 /// 8. The commission rate for validators can be updated with the `updateCommission` function.
 ///
+/// 9. The `totalValidatorStake` and `totalStake` functions are added to allow governance to
+/// track the total stake in the contract. The totalValidatorStake is the total stake in active
+/// validators and the totalStake is the total stake in all states (active + exiting +
+/// pendingWithdrawal).
+///
 /// @notice The StakeTableV2 contract ABI is a superset of the original ABI. Consumers of the
 /// contract can use the V2 ABI, even if they would like to maintain backwards compatibility.
 contract StakeTableV2 is StakeTable, PausableUpgradeable, AccessControlUpgradeable {
@@ -84,6 +89,12 @@ contract StakeTableV2 is StakeTable, PausableUpgradeable, AccessControlUpgradeab
     /// @notice Maximum commission increase allowed per increase (in basis points)
     uint16 public maxCommissionIncrease;
 
+    /// @notice Total stake in active (not marked for exit) validators in the contract
+    uint256 public totalValidatorStake;
+
+    /// @notice Total stake in all states (active + exiting + pendingWithdrawal) in the contract
+    uint256 public totalStake;
+
     /// @notice Commission tracking for each validator
     mapping(address validator => CommissionTracking tracking) public commissionTracking;
 
@@ -254,6 +265,9 @@ contract StakeTableV2 is StakeTable, PausableUpgradeable, AccessControlUpgradeab
         // it's only decremented during withdrawal
         validators[validator].delegatedAmount -= amount;
 
+        // finally, update the total stake as the balance was withdrawn
+        totalStake -= amount;
+
         SafeTransferLib.safeTransfer(token, delegator, amount);
 
         emit Withdrawal(delegator, amount);
@@ -263,7 +277,26 @@ contract StakeTableV2 is StakeTable, PausableUpgradeable, AccessControlUpgradeab
     /// @param validator The validator to withdraw from
     /// @dev This function is overridden to add pausable functionality
     function claimWithdrawal(address validator) public virtual override whenNotPaused {
-        super.claimWithdrawal(validator);
+        address delegator = msg.sender;
+        // If entries are missing at any of the levels of the mapping this will return zero
+        uint256 amount = undelegations[validator][delegator].amount;
+        if (amount == 0) {
+            revert NothingToWithdraw();
+        }
+
+        if (block.timestamp < undelegations[validator][delegator].unlocksAt) {
+            revert PrematureWithdrawal();
+        }
+
+        // Mark funds as spent
+        delete undelegations[validator][delegator];
+
+        // update the total stake managed by the contract
+        totalStake -= amount;
+
+        SafeTransferLib.safeTransfer(token, delegator, amount);
+
+        emit Withdrawal(delegator, amount);
     }
 
     /// @notice Delegate funds to a validator
@@ -272,6 +305,9 @@ contract StakeTableV2 is StakeTable, PausableUpgradeable, AccessControlUpgradeab
     /// @dev This function is overridden to add pausable functionality
     function delegate(address validator, uint256 amount) public virtual override whenNotPaused {
         super.delegate(validator, amount);
+
+        totalStake += amount;
+        totalValidatorStake += amount;
     }
 
     /// @notice Undelegate funds from a validator
@@ -280,6 +316,7 @@ contract StakeTableV2 is StakeTable, PausableUpgradeable, AccessControlUpgradeab
     /// @dev This function is overridden to add pausable functionality
     function undelegate(address validator, uint256 amount) public virtual override whenNotPaused {
         super.undelegate(validator, amount);
+        totalValidatorStake -= amount;
     }
 
     /// @notice Deregister a validator
@@ -295,6 +332,7 @@ contract StakeTableV2 is StakeTable, PausableUpgradeable, AccessControlUpgradeab
         validatorExits[validator] = block.timestamp + exitEscrowPeriod;
         // in v2, the delegatedAmount is not updated until withdrawal
 
+        totalValidatorStake -= validators[validator].delegatedAmount;
         emit ValidatorExit(validator);
     }