Gateway stake accounting

common/client-libs/validator-client/src/client.rs

@@ -189,6 +189,13 @@ impl<C> Client<C> {
         Ok(self.nymd.get_total_mix_stake().await?.u128())
     }
 
+    pub async fn get_total_gt_stake(&self) -> Result<u128, ValidatorClientError>
+    where
+        C: CosmWasmClient + Sync,
+    {
+        Ok(self.nymd.get_total_gt_stake().await?.u128())
+    }
+
     // basically handles paging for us
     pub async fn get_all_nymd_mixnodes(&self) -> Result<Vec<MixNodeBond>, ValidatorClientError>
     where

common/client-libs/validator-client/src/nymd/mod.rs

@@ -208,6 +208,16 @@ impl<C> NymdClient<C> {
             .await
     }
 
+    pub async fn get_total_gt_stake(&self) -> Result<Uint128, NymdError>
+    where
+        C: CosmWasmClient + Sync,
+    {
+        let request = QueryMsg::GetTotalGtStake {};
+        self.client
+            .query_contract_smart(self.contract_address()?, &request)
+            .await
+    }
+
     /// Checks whether there is a bonded mixnode associated with the provided client's address
     pub async fn owns_mixnode(&self, address: &AccountId) -> Result<bool, NymdError>
     where

common/mixnet-contract/src/gateway.rs

@@ -59,6 +59,10 @@ impl GatewayBond {
     pub fn gateway(&self) -> &Gateway {
         &self.gateway
     }
+
+    pub fn total_delegation(&self) -> Coin {
+        self.total_delegation.clone()
+    }
 }
 
 impl PartialOrd for GatewayBond {

common/mixnet-contract/src/msg.rs

@@ -108,6 +108,7 @@ pub enum QueryMsg {
     },
     LayerDistribution {},
     GetTotalMixStake {},
+    GetTotalGtStake {},
 }
 
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]

contracts/mixnet/src/contract.rs

@@ -208,13 +208,14 @@ pub fn query(deps: Deps, _env: Env, msg: QueryMsg) -> Result<QueryResponse, Cont
             address,
         )?),
         QueryMsg::GetTotalMixStake {} => to_binary(&queries::query_total_mix_stake(deps)),
+        QueryMsg::GetTotalGtStake {} => to_binary(&queries::query_total_gt_stake(deps)),
     };
 
     Ok(query_res?)
 }
 #[entry_point]
 pub fn migrate(_deps: DepsMut, _env: Env, _msg: MigrateMsg) -> Result<Response, ContractError> {
-    todo!("Calculate initial total mix stake");
+    todo!("Calculate initial total mix and gateway stake after initial deployment");
 
     #[allow(unreachable_code)]
     Ok(Default::default())

contracts/mixnet/src/queries.rs

@@ -2,7 +2,12 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::error::ContractError;
-use crate::storage::{gateway_delegations_read, gateways_owners_read, gateways_read, mix_delegations_read, mixnodes_owners_read, mixnodes_read, read_layer_distribution, read_state_params, reverse_gateway_delegations_read, reverse_mix_delegations_read, total_mix_stake, all_gateway_delegations_read, all_mix_delegations_read};
+use crate::storage::{
+    gateway_delegations_read, gateways_owners_read, gateways_read, mix_delegations_read,
+    mixnodes_owners_read, mixnodes_read, read_layer_distribution, read_state_params,
+    reverse_gateway_delegations_read, reverse_mix_delegations_read, total_gateway_stake,
+    total_mix_stake, all_mix_delegations_read, all_gateway_delegations_read
+};
 use config::defaults::DENOM;
 use cosmwasm_std::Deps;
 use cosmwasm_std::Order;
@@ -98,6 +103,11 @@ pub(crate) fn query_total_mix_stake(deps: Deps) -> Uint128 {
     total_mix_stake(deps.storage).load().unwrap()
 }
 
+pub(crate) fn query_total_gt_stake(deps: Deps) -> Uint128 {
+    // Same as above
+    total_gateway_stake(deps.storage).load().unwrap()
+}
+
 /// Adds a 0 byte to terminate the `start_after` value given. This allows CosmWasm
 /// to get the succeeding key as the start of the next page.
 // S works for both `String` and `Addr` and that's what we wanted

contracts/mixnet/src/storage.rs

@@ -24,8 +24,9 @@ use serde::Serialize;
 // singletons
 const CONFIG_KEY: &[u8] = b"config";
 const LAYER_DISTRIBUTION_KEY: &[u8] = b"layers";
-// Keeps total amount of stake towards mixnodes. Removing a mixnode bond removes all its delegations from the total, the reverse is true for adding a mixnode bond.ß
-const TOTAL_STAKE_KEY: &[u8] = b"total_mix";
+// Keeps total amount of stake towards mixnodes and gateways. Removing a bond removes all its delegations from the total, the reverse is true for adding a bond.
+const TOTAL_MN_STAKE_KEY: &[u8] = b"total_mn";
+const TOTAL_GT_STAKE_KEY: &[u8] = b"total_gt";
 
 // buckets
 const PREFIX_MIXNODES: &[u8] = b"mn";
@@ -51,11 +52,19 @@ pub fn config_read(storage: &dyn Storage) -> ReadonlySingleton<State> {
 }
 
 pub fn total_mix_stake(storage: &dyn Storage) -> ReadonlySingleton<Uint128> {
-    singleton_read(storage, TOTAL_STAKE_KEY)
+    singleton_read(storage, TOTAL_MN_STAKE_KEY)
 }
 
 pub fn mut_total_mix_stake(storage: &mut dyn Storage) -> Singleton<Uint128> {
-    singleton(storage, TOTAL_STAKE_KEY)
+    singleton(storage, TOTAL_MN_STAKE_KEY)
+}
+
+pub fn total_gateway_stake(storage: &dyn Storage) -> ReadonlySingleton<Uint128> {
+    singleton_read(storage, TOTAL_GT_STAKE_KEY)
+}
+
+pub fn mut_total_gateway_stake(storage: &mut dyn Storage) -> Singleton<Uint128> {
+    singleton(storage, TOTAL_GT_STAKE_KEY)
 }
 
 pub fn incr_total_mix_stake(
@@ -76,6 +85,24 @@ pub fn decr_total_mix_stake(
     Ok(())
 }
 
+pub fn incr_total_gt_stake(
+    amount: Uint128,
+    storage: &mut dyn Storage,
+) -> Result<(), ContractError> {
+    let stake = total_gateway_stake(storage).load()?.checked_add(amount)?;
+    mut_total_gateway_stake(storage).save(&stake)?;
+    Ok(())
+}
+
+pub fn decr_total_gt_stake(
+    amount: Uint128,
+    storage: &mut dyn Storage,
+) -> Result<(), ContractError> {
+    let stake = total_gateway_stake(storage).load()?.checked_sub(amount)?;
+    mut_total_gateway_stake(storage).save(&stake)?;
+    Ok(())
+}
+
 pub(crate) fn read_state_params(storage: &dyn Storage) -> StateParams {
     // note: In any other case, I wouldn't have attempted to unwrap this result, but in here
     // if we fail to load the stored state we would already be in the undefined behaviour land,

contracts/mixnet/src/transactions.rs

@@ -248,6 +248,9 @@ pub(crate) fn try_add_gateway(
     gateways_owners(deps.storage).save(sender_bytes, identity)?;
     increment_layer_count(deps.storage, Layer::Gateway)?;
 
+    incr_total_gt_stake(bond.bond_amount().amount, deps.storage)?;
+    incr_total_gt_stake(bond.total_delegation().amount, deps.storage)?;
+
     let attributes = vec![attr("overwritten", was_present)];
     Ok(Response {
         submessages: Vec::new(),
@@ -286,6 +289,9 @@ pub(crate) fn try_remove_gateway(
     // decrement layer count
     decrement_layer_count(deps.storage, Layer::Gateway)?;
 
+    decr_total_gt_stake(gateway_bond.bond_amount().amount, deps.storage)?;
+    decr_total_gt_stake(gateway_bond.total_delegation().amount, deps.storage)?;
+
     // log our actions
     let attributes = vec![
         attr("action", "unbond"),

validator-api/src/cache/mod.rs

@@ -39,6 +39,7 @@ struct ValidatorCacheInner {
     current_gateway_active_set_size: AtomicUsize,
 
     total_mix_stake: RwLock<Cache<u128>>,
+    total_gt_stake: RwLock<Cache<u128>>,
 }
 
 #[derive(Default, Serialize, Clone)]
@@ -78,10 +79,11 @@ impl<C> ValidatorCacheRefresher<C> {
     where
         C: CosmWasmClient + Sync,
     {
-        let (mixnodes, gateways, total_mix_stake) = tokio::try_join!(
+        let (mixnodes, gateways, total_mix_stake, total_gt_state) = tokio::try_join!(
             self.nymd_client.get_mixnodes(),
             self.nymd_client.get_gateways(),
-            self.nymd_client.get_total_mix_stake()
+            self.nymd_client.get_total_mix_stake(),
+            self.nymd_client.get_total_gt_stake()
         )?;
 
         let state_params = self.nymd_client.get_state_params().await?;
@@ -94,7 +96,13 @@ impl<C> ValidatorCacheRefresher<C> {
         );
 
         self.cache
-            .update_cache(mixnodes, gateways, total_mix_stake, state_params)
+            .update_cache(
+                mixnodes,
+                gateways,
+                total_mix_stake,
+                total_gt_state,
+                state_params,
+            )
             .await;
 
         Ok(())
@@ -182,6 +190,7 @@ impl ValidatorCache {
         mut mixnodes: Vec<MixNodeBond>,
         mut gateways: Vec<GatewayBond>,
         total_mix_stake: u128,
+        total_gt_stake: u128,
         state: StateParams,
     ) {
         // if our data is valid, it means the active sets are available,
@@ -228,6 +237,7 @@ impl ValidatorCache {
             .write()
             .await
             .set(total_mix_stake);
+        self.inner.total_gt_stake.write().await.set(total_gt_stake);
     }
 
     pub async fn mixnodes(&self) -> Cache<Vec<MixNodeBond>> {
@@ -314,6 +324,7 @@ impl ValidatorCacheInner {
             current_mixnode_active_set_size: Default::default(),
             current_gateway_active_set_size: Default::default(),
             total_mix_stake: Default::default(),
+            total_gt_stake: Default::default(),
         }
     }
 }

validator-api/src/nymd_client.rs

@@ -99,6 +99,13 @@ impl<C> Client<C> {
         Ok(self.0.read().await.get_total_mix_stake().await?)
     }
 
+    pub(crate) async fn get_total_gt_stake(&self) -> Result<u128, ValidatorClientError>
+    where
+        C: CosmWasmClient + Sync,
+    {
+        Ok(self.0.read().await.get_total_gt_stake().await?)
+    }
+
     pub(crate) async fn get_mixnodes(&self) -> Result<Vec<MixNodeBond>, ValidatorClientError>
     where
         C: CosmWasmClient + Sync,