penumbra_sdk_shielded_pool/component/
shielded_pool.rs

1use std::sync::Arc;
2
3use super::fmd::ClueManagerInternal as _;
4use crate::fmd::should_update_fmd_params;
5use crate::params::ShieldedPoolParameters;
6use crate::{fmd, genesis, state_key};
7use anyhow::anyhow;
8use anyhow::Result;
9use async_trait::async_trait;
10use cnidarium::{StateRead, StateWrite};
11use cnidarium_component::Component;
12use penumbra_sdk_proto::StateReadProto as _;
13use penumbra_sdk_proto::StateWriteProto as _;
14use penumbra_sdk_sct::CommitmentSource;
15use tendermint::v0_37::abci;
16use tracing::instrument;
17
18use super::{AssetRegistry, NoteManager};
19
20pub struct ShieldedPool {}
21
22#[async_trait]
23impl Component for ShieldedPool {
24    type AppState = genesis::Content;
25
26    #[instrument(name = "shielded_pool", skip(state, app_state))]
27    async fn init_chain<S: StateWrite>(mut state: S, app_state: Option<&Self::AppState>) {
28        match app_state {
29            None => { /* Checkpoint -- no-op */ }
30            Some(genesis) => {
31                // TODO(erwan): the handling of those parameters is a bit weird.
32                // rationalize it before merging
33                state.put_shielded_pool_params(genesis.shielded_pool_params.clone());
34                state.put_current_fmd_parameters(fmd::Parameters::default());
35                state.put_previous_fmd_parameters(fmd::Parameters::default());
36
37                // Register a denom for each asset in the genesis state
38                for allocation in &genesis.allocations {
39                    tracing::debug!(?allocation, "processing allocation");
40                    assert_ne!(
41                        allocation.raw_amount,
42                        0u128.into(),
43                        "Genesis allocations contain empty note",
44                    );
45
46                    state.register_denom(&allocation.denom()).await;
47                    state
48                        .mint_note(
49                            allocation.value(),
50                            &allocation.address,
51                            CommitmentSource::Genesis,
52                        )
53                        .await
54                        .expect("able to mint note for genesis allocation");
55                }
56            }
57        }
58    }
59
60    #[instrument(name = "shielded_pool", skip(_state, _begin_block))]
61    async fn begin_block<S: StateWrite + 'static>(
62        _state: &mut Arc<S>,
63        _begin_block: &abci::request::BeginBlock,
64    ) {
65    }
66
67    #[instrument(name = "shielded_pool", skip_all)]
68    async fn end_block<S: StateWrite + 'static>(
69        state: &mut Arc<S>,
70        end_block: &abci::request::EndBlock,
71    ) {
72        let height: u64 = end_block
73            .height
74            .try_into()
75            .expect("height should not be negative");
76        let state = Arc::get_mut(state).expect("the state should not be shared");
77        let meta_params = state
78            .get_shielded_pool_params()
79            .await
80            .expect("should be able to read state")
81            .fmd_meta_params;
82        if should_update_fmd_params(meta_params.fmd_grace_period_blocks, height) {
83            let old = state
84                .get_current_fmd_parameters()
85                .await
86                .expect("should be able to read state");
87            let clue_count_delta = state
88                .flush_clue_count()
89                .await
90                .expect("should be able to read state");
91            let algorithm_state = state
92                .get_fmd_algorithm_state()
93                .await
94                .expect("should be able to read state");
95            let (new, algorithm_state) =
96                meta_params.updated_fmd_params(&old, algorithm_state, height, clue_count_delta);
97            state.put_previous_fmd_parameters(old);
98            state.put_current_fmd_parameters(new);
99            state.put_fmd_algorithm_state(algorithm_state);
100        }
101    }
102
103    async fn end_epoch<S: StateWrite + 'static>(mut _state: &mut Arc<S>) -> Result<()> {
104        Ok(())
105    }
106}
107/// Extension trait providing read access to shielded pool data.
108#[async_trait]
109pub trait StateReadExt: StateRead {
110    async fn get_current_fmd_parameters(&self) -> Result<fmd::Parameters> {
111        self.get(fmd::state_key::parameters::current())
112            .await?
113            .ok_or_else(|| anyhow!("Missing FmdParameters"))
114    }
115
116    /// Gets the previous FMD parameters from the JMT.
117    async fn get_previous_fmd_parameters(&self) -> Result<fmd::Parameters> {
118        self.get(fmd::state_key::parameters::previous())
119            .await?
120            .ok_or_else(|| anyhow!("Missing FmdParameters"))
121    }
122
123    async fn get_shielded_pool_params(&self) -> Result<ShieldedPoolParameters> {
124        self.get(state_key::shielded_pool_params())
125            .await?
126            .ok_or_else(|| anyhow!("Missing ShieldedPoolParameters"))
127    }
128
129    async fn get_fmd_algorithm_state(&self) -> Result<fmd::MetaParametersAlgorithmState> {
130        Ok(self
131            .get(fmd::state_key::meta_parameters::algorithm_state())
132            .await?
133            .unwrap_or_default())
134    }
135}
136
137impl<T: StateRead + ?Sized> StateReadExt for T {}
138
139/// Extension trait providing write access to shielded pool data.
140#[async_trait]
141pub trait StateWriteExt: StateWrite + StateReadExt {
142    fn put_shielded_pool_params(&mut self, params: ShieldedPoolParameters) {
143        self.put(crate::state_key::shielded_pool_params().into(), params)
144    }
145
146    /// Writes the current FMD parameters to the JMT.
147    fn put_current_fmd_parameters(&mut self, params: fmd::Parameters) {
148        self.put(fmd::state_key::parameters::current().into(), params)
149    }
150
151    /// Writes the previous FMD parameters to the JMT.
152    fn put_previous_fmd_parameters(&mut self, params: fmd::Parameters) {
153        self.put(fmd::state_key::parameters::previous().into(), params)
154    }
155
156    fn put_fmd_algorithm_state(&mut self, state: fmd::MetaParametersAlgorithmState) {
157        self.put(
158            fmd::state_key::meta_parameters::algorithm_state().into(),
159            state,
160        )
161    }
162}
163
164impl<T: StateWrite + ?Sized> StateWriteExt for T {}