espresso_types/v0/v0_4/
state.rs1use std::{
2 collections::HashSet,
3 sync::{Arc, OnceLock},
4};
5
6use alloy::primitives::{Address, U256};
7use anyhow::Context;
8use derive_more::{Display, From, Into};
9use hotshot_contract_adapter::reward::{RewardAuthData, RewardClaimInput};
10use jf_merkle_tree_compat::{MerkleTreeScheme, UniversalMerkleTreeScheme};
11use serde::{Deserialize, Serialize};
12use tokio::sync::{OwnedSemaphorePermit, Semaphore};
13
14use super::FeeAccount;
15pub use crate::v0::reward_mt::{
16 REWARD_MERKLE_TREE_V2_ARITY, REWARD_MERKLE_TREE_V2_HEIGHT, RewardMerkleCommitmentV2,
17 RewardMerkleTreeV2,
18};
19use crate::v0_3::{RewardAccountV1, RewardAmount};
20
21#[derive(Clone)]
22pub struct PermittedRewardMerkleTreeV2 {
23 pub tree: RewardMerkleTreeV2,
24 _permit: Arc<OwnedSemaphorePermit>,
25}
26
27static REWARD_MERKLE_TREE_V2_MEMORY_LOCK: OnceLock<Arc<Semaphore>> = OnceLock::new();
28
29impl std::ops::Deref for PermittedRewardMerkleTreeV2 {
30 type Target = RewardMerkleTreeV2;
31
32 fn deref(&self) -> &Self::Target {
33 &self.tree
34 }
35}
36
37impl PermittedRewardMerkleTreeV2 {
38 pub async fn try_from_kv_set(
39 balances: Vec<(RewardAccountV2, RewardAmount)>,
40 ) -> anyhow::Result<Self> {
41 let permit = REWARD_MERKLE_TREE_V2_MEMORY_LOCK
42 .get_or_init(|| Arc::new(Semaphore::new(1)))
43 .clone()
44 .acquire_owned()
45 .await
46 .context("Failed to acquire permit for RewardMerkleTreeV2")?;
47
48 let tree = RewardMerkleTreeV2::from_kv_set(REWARD_MERKLE_TREE_V2_HEIGHT, balances)
49 .context("Failed to rebuild reward merkle tree from balances")?;
50
51 Ok(PermittedRewardMerkleTreeV2 {
52 tree,
53 _permit: Arc::new(permit),
54 })
55 }
56}
57
58pub fn forgotten_accounts_include(tree: &RewardMerkleTreeV2, accounts: &[RewardAccountV2]) -> bool {
60 for account in accounts {
61 if tree.lookup(*account).expect_not_in_memory().is_ok() {
62 return true;
63 }
64 }
65
66 false
67}
68
69#[derive(
74 Default,
75 Hash,
76 Copy,
77 Clone,
78 Debug,
79 Display,
80 Deserialize,
81 Serialize,
82 PartialEq,
83 Eq,
84 PartialOrd,
85 Ord,
86 From,
87 Into,
88)]
89#[display("{_0}")]
90pub struct RewardAccountV2(pub Address);
91
92impl From<RewardAccountV2> for RewardAccountV1 {
93 fn from(account: RewardAccountV2) -> Self {
94 RewardAccountV1(account.0)
95 }
96}
97
98impl From<RewardAccountV1> for RewardAccountV2 {
99 fn from(account: RewardAccountV1) -> Self {
100 RewardAccountV2(account.0)
101 }
102}
103
104#[derive(Clone, Debug, Deserialize, Serialize)]
109pub struct RewardAccountProofV2 {
110 pub account: Address,
111 pub proof: RewardMerkleProofV2,
112}
113
114#[derive(Clone, Debug, Deserialize, Serialize)]
115pub enum RewardMerkleProofV2 {
116 Presence(<RewardMerkleTreeV2 as MerkleTreeScheme>::MembershipProof),
117 Absence(<RewardMerkleTreeV2 as UniversalMerkleTreeScheme>::NonMembershipProof),
118}
119
120#[derive(Clone, Debug, Serialize, Deserialize)]
121pub struct RewardAccountQueryDataV2 {
122 pub balance: U256,
123 pub proof: RewardAccountProofV2,
124}
125
126#[derive(Debug, thiserror::Error)]
127pub enum RewardClaimError {
128 #[error("Zero reward balance")]
129 ZeroRewardError,
130 #[error("Failed to convert proof: {0}")]
131 ProofConversionError(#[from] anyhow::Error),
132}
133
134impl RewardAccountQueryDataV2 {
135 pub fn to_reward_claim_input(self) -> Result<RewardClaimInput, RewardClaimError> {
140 if self.balance == U256::ZERO {
141 return Err(RewardClaimError::ZeroRewardError);
142 }
143
144 let account_proof = match self.proof.proof {
145 RewardMerkleProofV2::Presence(_) => self.proof,
146 RewardMerkleProofV2::Absence(_) => {
147 return Err(RewardClaimError::ZeroRewardError);
148 },
149 };
150
151 Ok(RewardClaimInput {
152 lifetime_rewards: self.balance,
153 auth_data: RewardAuthData::new(account_proof.try_into()?).into(),
154 })
155 }
156}
157
158#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
159pub struct Delta {
160 pub fees_delta: HashSet<FeeAccount>,
161 pub rewards_delta: HashSet<RewardAccountV2>,
162}