penumbra_sdk_community_pool/component/
view.rs1use 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 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 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 {}