penumbra_sdk_community_pool/component/
view.rs

1use std::{collections::BTreeMap, str::FromStr};
2
3use anyhow::Result;
4use async_trait::async_trait;
5
6use cnidarium::{StateRead, StateWrite};
7use futures::{StreamExt, TryStreamExt};
8use penumbra_sdk_asset::{asset, Value};
9use penumbra_sdk_num::Amount;
10use penumbra_sdk_proto::{StateReadProto, StateWriteProto};
11
12use crate::params::CommunityPoolParameters;
13
14use super::state_key;
15
16#[async_trait]
17pub trait StateReadExt: StateRead {
18    /// Gets the Community Pool parameters from the JMT.
19    async fn get_community_pool_params(&self) -> Result<CommunityPoolParameters> {
20        self.get(state_key::community_pool_params())
21            .await?
22            .ok_or_else(|| anyhow::anyhow!("Missing CommunityPoolParameters"))
23    }
24
25    async fn community_pool_asset_balance(&self, asset_id: asset::Id) -> Result<Amount> {
26        Ok(self
27            .get(&state_key::balance_for_asset(asset_id))
28            .await?
29            .unwrap_or_else(|| Amount::from(0u64)))
30    }
31
32    async fn community_pool_balance(&self) -> Result<BTreeMap<asset::Id, Amount>> {
33        let prefix = state_key::all_assets_balance();
34        self.prefix(prefix)
35            .map(|result| {
36                let (key, amount) = result?;
37                let asset_id = key.rsplit('/').next().expect("key is well-formed");
38                let asset_id = asset::Id::from_str(asset_id)?;
39                Ok((asset_id, amount))
40            })
41            .try_collect()
42            .await
43    }
44}
45
46impl<T> StateReadExt for T where T: StateRead + ?Sized {}
47
48#[async_trait]
49pub trait StateWriteExt: StateWrite {
50    /// Writes the provided Community Pool parameters to the JMT.
51    fn put_community_pool_params(&mut self, params: CommunityPoolParameters) {
52        self.put(state_key::community_pool_params().into(), params)
53    }
54
55    async fn community_pool_deposit(&mut self, value: Value) {
56        let key = state_key::balance_for_asset(value.asset_id);
57        let current = self
58            .get(&key)
59            .await
60            .expect("no deserialization errors")
61            .unwrap_or_else(|| Amount::from(0u64));
62        self.put(key, current + value.amount);
63    }
64
65    async fn community_pool_withdraw(&mut self, value: Value) -> Result<()> {
66        let key = state_key::balance_for_asset(value.asset_id);
67        let current = self.get(&key).await?.unwrap_or_else(|| Amount::from(0u64));
68        if let Some(remaining) = u128::from(current).checked_sub(u128::from(value.amount)) {
69            if remaining > 0 {
70                self.put(key, Amount::from(remaining));
71            } else {
72                self.delete(key);
73            }
74        } else {
75            anyhow::bail!(
76                "insufficient balance to withdraw {} of asset ID {} from the Community Pool",
77                value.amount,
78                value.asset_id
79            );
80        }
81        Ok(())
82    }
83}
84
85impl<T> StateWriteExt for T where T: StateWrite + ?Sized {}