staking: rework epoch transition rate calculations

crates/core/app/src/app/mod.rs

@@ -127,9 +127,9 @@ impl App {
                     },
                 );
 
-                Distributions::init_chain(&mut state_tx, Some(&())).await;
                 ShieldedPool::init_chain(&mut state_tx, Some(&app_state.shielded_pool_content))
                     .await;
+                Distributions::init_chain(&mut state_tx, Some(&())).await;
                 Staking::init_chain(&mut state_tx, Some(&app_state.stake_content)).await;
                 IBCComponent::init_chain(&mut state_tx, Some(&())).await;
                 Dex::init_chain(&mut state_tx, Some(&())).await;
@@ -138,8 +138,8 @@ impl App {
             }
             genesis::AppState::Checkpoint(_) => {
                 /* perform upgrade specific check */
-                Distributions::init_chain(&mut state_tx, None).await;
                 ShieldedPool::init_chain(&mut state_tx, None).await;
+                Distributions::init_chain(&mut state_tx, None).await;
                 Staking::init_chain(&mut state_tx, None).await;
                 IBCComponent::init_chain(&mut state_tx, None).await;
                 Dex::init_chain(&mut state_tx, None).await;
@@ -196,10 +196,10 @@ impl App {
 
         // Run each of the begin block handlers for each component, in sequence:
         let mut arc_state_tx = Arc::new(state_tx);
+        ShieldedPool::begin_block(&mut arc_state_tx, begin_block).await;
         Distributions::begin_block(&mut arc_state_tx, begin_block).await;
         IBCComponent::begin_block(&mut arc_state_tx, begin_block).await;
         Governance::begin_block(&mut arc_state_tx, begin_block).await;
-        ShieldedPool::begin_block(&mut arc_state_tx, begin_block).await;
         Staking::begin_block(&mut arc_state_tx, begin_block).await;
         Fee::begin_block(&mut arc_state_tx, begin_block).await;
 
@@ -301,11 +301,11 @@ impl App {
         let state_tx = StateDelta::new(self.state.clone());
 
         let mut arc_state_tx = Arc::new(state_tx);
+        ShieldedPool::end_block(&mut arc_state_tx, end_block).await;
         Distributions::end_block(&mut arc_state_tx, end_block).await;
         IBCComponent::end_block(&mut arc_state_tx, end_block).await;
         Dex::end_block(&mut arc_state_tx, end_block).await;
         Governance::end_block(&mut arc_state_tx, end_block).await;
-        ShieldedPool::end_block(&mut arc_state_tx, end_block).await;
         Staking::end_block(&mut arc_state_tx, end_block).await;
         Fee::end_block(&mut arc_state_tx, end_block).await;
         let mut state_tx = Arc::try_unwrap(arc_state_tx)

crates/core/component/distributions/Cargo.toml

@@ -13,11 +13,13 @@ docsrs = []
 [dependencies]
 
 # Workspace dependencies
-penumbra-proto = { path = "../../../proto", default-features = false }
-penumbra-storage = { path = "../../../storage", optional = true }
+penumbra-asset = { path = "../../asset" }
 penumbra-component = { path = "../component", optional = true }
 penumbra-chain = { path = "../chain", default-features = false }
+penumbra-num = { path = "../../../core/num", default-features = false }
+penumbra-proto = { path = "../../../proto", default-features = false }
 penumbra-shielded-pool = { path = "../shielded-pool", default-features = false }
+penumbra-storage = { path = "../../../storage", optional = true }
 
 # Crates.io deps
 async-trait = "0.1.52"

crates/core/component/distributions/src/component.rs

@@ -7,8 +7,14 @@ use std::sync::Arc;
 use anyhow::Result;
 use async_trait::async_trait;
 use penumbra_component::Component;
+// use penumbra_dex::{component::StateReadExt as _, component::StateWriteExt as _};
+// use penumbra_stake::{component::StateWriteExt as _, StateReadExt as _};
+use penumbra_asset::STAKING_TOKEN_ASSET_ID;
+use penumbra_num::Amount;
+use penumbra_shielded_pool::component::SupplyRead;
 use penumbra_storage::StateWrite;
-use tendermint::abci;
+use tendermint::v0_37::abci;
+use tracing::instrument;
 pub use view::{StateReadExt, StateWriteExt};
 
 pub struct Distributions {}
@@ -17,21 +23,69 @@ pub struct Distributions {}
 impl Component for Distributions {
     type AppState = ();
 
-    async fn init_chain<S: StateWrite>(_state: S, _app_state: Option<&()>) {}
+    #[instrument(name = "distributions", skip(state, app_state))]
+    async fn init_chain<S: StateWrite>(mut state: S, app_state: Option<&Self::AppState>) {
+        match app_state {
+            None => { /* Checkpoint -- no-op */ }
+            Some(_) => {
+                let genesis_issuance = state
+                    .token_supply(&*STAKING_TOKEN_ASSET_ID)
+                    .await
+                    .expect("supply is valid")
+                    .expect("shielded pool component has tallied genesis issuance");
+                tracing::debug!(
+                    "total genesis issuance of staking token: {}",
+                    genesis_issuance
+                );
+                // TODO(erwan): it's not yet totally clear if it is necessary, or even desirable, for the
+                // distributions component to track the total issuance. The shielded pool component
+                // already does that. We do it anyway for now so that we can write the rest of the scaffolding.
+                state.set_total_issued(genesis_issuance);
+            }
+        };
+    }
 
+    #[instrument(name = "distributions", skip(_state, _begin_block))]
     async fn begin_block<S: StateWrite + 'static>(
         _state: &mut Arc<S>,
         _begin_block: &abci::request::BeginBlock,
     ) {
     }
 
+    #[instrument(name = "distributions", skip(_state, _end_block))]
     async fn end_block<S: StateWrite + 'static>(
         _state: &mut Arc<S>,
         _end_block: &abci::request::EndBlock,
     ) {
     }
 
+    #[instrument(name = "distributions", skip(_state))]
     async fn end_epoch<S: StateWrite + 'static>(_state: &mut Arc<S>) -> Result<()> {
         Ok(())
     }
 }
+
+#[async_trait]
+trait DistributionManager: StateWriteExt {
+    /// Compute the total new issuance of staking tokens for this epoch.
+    /// TODO(erwan): this is a stub implementation.
+    async fn compute_new_issuance(&self) -> Result<Amount> {
+        let base_reward_rate: u64 = 0;
+        let total_issued = self
+            .total_issued()
+            .await?
+            .expect("total issuance has been initialized");
+        const BPS_SQUARED: u64 = 1_0000_0000; // reward rate is measured in basis points squared
+        let new_issuance = total_issued * base_reward_rate / BPS_SQUARED;
+        Ok(new_issuance.into())
+    }
+
+    /// Update the object store with the new issuance of staking tokens for this epoch.
+    /// TODO(erwan): this is a stub implementation.
+    async fn distribute(&mut self) -> Result<()> {
+        let new_issuance = self.compute_new_issuance().await?;
+        tracing::debug!(?new_issuance, "computed new issuance for epoch");
+        self.set_total_issued(new_issuance);
+        todo!()
+    }
+}

crates/core/component/distributions/src/component/state_key.rs

@@ -1 +1,4 @@
-
+// The cumulative total issuance of staking token.
+pub fn total_issued() -> &'static str {
+    "distributions/total_issued"
+}

crates/core/component/distributions/src/component/view.rs

@@ -1,13 +1,27 @@
 use async_trait::async_trait;
 
+use crate::component::state_key;
+use anyhow::Result;
+use penumbra_num::Amount;
+use penumbra_proto::{StateReadProto, StateWriteProto};
 use penumbra_storage::{StateRead, StateWrite};
 
 #[async_trait]
-pub trait StateReadExt: StateRead {}
+pub trait StateReadExt: StateRead {
+    async fn total_issued(&self) -> Result<Option<u64>> {
+        self.get_proto(&state_key::total_issued()).await
+    }
+}
 
-impl<T> StateReadExt for T where T: StateRead + ?Sized {}
+impl<T: StateRead> StateReadExt for T {}
 
 #[async_trait]
-pub trait StateWriteExt: StateWrite {}
 
-impl<T> StateWriteExt for T where T: StateWrite + ?Sized {}
+pub trait StateWriteExt: StateWrite + StateReadExt {
+    /// Set the total amount of staking tokens issued.
+    fn set_total_issued(&mut self, total_issued: Amount) {
+        let total = Amount::from(total_issued);
+        self.put(state_key::total_issued().to_string(), total)
+    }
+}
+impl<T: StateWrite> StateWriteExt for T {}

crates/core/component/shielded-pool/src/component/supply.rs

@@ -12,8 +12,8 @@ use crate::state_key;
 
 #[async_trait]
 pub trait SupplyRead: StateRead {
-    async fn token_supply(&self, asset_id: &asset::Id) -> Result<Option<u64>> {
-        self.get_proto(&state_key::token_supply(asset_id)).await
+    async fn token_supply(&self, asset_id: &asset::Id) -> Result<Option<Amount>> {
+        self.get(&state_key::token_supply(asset_id)).await
     }
 
     // TODO: refactor for new state model -- no more list of known asset IDs with fixed key
@@ -60,7 +60,7 @@ pub trait SupplyWrite: StateWrite {
     // #[instrument(skip(self, change))]
     async fn update_token_supply(&mut self, asset_id: &asset::Id, change: i128) -> Result<()> {
         let key = state_key::token_supply(asset_id);
-        let current_supply: Amount = self.get_proto(&key).await?.unwrap_or(0u64).into();
+        let current_supply: Amount = self.get(&key).await?.unwrap_or(0u64.into());
 
         // TODO: replace with a single checked_add_signed call when mixed_integer_ops lands in stable (1.66)
         let new_supply: Amount = if change < 0 {

crates/core/component/stake/src/action_handler/validator_definition.rs

@@ -113,24 +113,15 @@ impl ActionHandler for validator::Definition {
             // Set the default rates and state.
             let validator_key = v.validator.identity_key;
 
-            // Delegations require knowing the rates for the
-            // next epoch, so pre-populate with 0 reward => exchange rate 1 for
-            // the current and next epochs.
             let cur_rate_data = RateData {
                 identity_key: validator_key,
                 epoch_index: cur_epoch.index,
                 validator_reward_rate: 0,
                 validator_exchange_rate: 1_0000_0000, // 1 represented as 1e8
             };
-            let next_rate_data = RateData {
-                identity_key: validator_key,
-                epoch_index: cur_epoch.index + 1,
-                validator_reward_rate: 0,
-                validator_exchange_rate: 1_0000_0000, // 1 represented as 1e8
-            };
 
             state
-                .add_validator(v.validator.clone(), cur_rate_data, next_rate_data)
+                .add_validator(v.validator.clone(), cur_rate_data)
                 .await
                 .context("should be able to add validator during validator definition execution")?;
         }