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(|| {
43 let num_permits: usize = std::env::var("ESPRESSO_NODE_REWARD_MERKLE_TREE_PERMITS")
44 .ok()
45 .and_then(|v| v.parse::<usize>().ok())
46 .unwrap_or(1);
47
48 tracing::warn!(
49 "Initializing RewardMerkleTreeV2 semaphore with {num_permits} permits"
50 );
51 Arc::new(Semaphore::new(num_permits))
52 })
53 .clone()
54 .acquire_owned()
55 .await
56 .context("Failed to acquire permit for RewardMerkleTreeV2")?;
57
58 let tree = RewardMerkleTreeV2::from_kv_set(REWARD_MERKLE_TREE_V2_HEIGHT, balances)
59 .context("Failed to rebuild reward merkle tree from balances")?;
60
61 Ok(PermittedRewardMerkleTreeV2 {
62 tree,
63 _permit: Arc::new(permit),
64 })
65 }
66}
67
68pub fn forgotten_accounts_include(tree: &RewardMerkleTreeV2, accounts: &[RewardAccountV2]) -> bool {
70 for account in accounts {
71 if tree.lookup(*account).expect_not_in_memory().is_ok() {
72 return true;
73 }
74 }
75
76 false
77}
78
79#[derive(
84 Default,
85 Hash,
86 Copy,
87 Clone,
88 Debug,
89 Display,
90 Deserialize,
91 Serialize,
92 PartialEq,
93 Eq,
94 PartialOrd,
95 Ord,
96 From,
97 Into,
98)]
99#[display("{_0}")]
100pub struct RewardAccountV2(pub Address);
101
102impl From<RewardAccountV2> for RewardAccountV1 {
103 fn from(account: RewardAccountV2) -> Self {
104 RewardAccountV1(account.0)
105 }
106}
107
108impl From<RewardAccountV1> for RewardAccountV2 {
109 fn from(account: RewardAccountV1) -> Self {
110 RewardAccountV2(account.0)
111 }
112}
113
114#[derive(Clone, Debug, Deserialize, Serialize)]
119pub struct RewardAccountProofV2 {
120 pub account: Address,
121 pub proof: RewardMerkleProofV2,
122}
123
124#[derive(Clone, Debug, Deserialize, Serialize)]
125pub enum RewardMerkleProofV2 {
126 Presence(<RewardMerkleTreeV2 as MerkleTreeScheme>::MembershipProof),
127 Absence(<RewardMerkleTreeV2 as UniversalMerkleTreeScheme>::NonMembershipProof),
128}
129
130#[derive(Clone, Debug, Serialize, Deserialize)]
131pub struct RewardAccountQueryDataV2 {
132 pub balance: U256,
133 pub proof: RewardAccountProofV2,
134}
135
136#[derive(Debug, thiserror::Error)]
137pub enum RewardClaimError {
138 #[error("Zero reward balance")]
139 ZeroRewardError,
140 #[error("Failed to convert proof: {0}")]
141 ProofConversionError(#[from] anyhow::Error),
142}
143
144impl RewardAccountQueryDataV2 {
145 pub fn to_reward_claim_input(self) -> Result<RewardClaimInput, RewardClaimError> {
150 if self.balance == U256::ZERO {
151 return Err(RewardClaimError::ZeroRewardError);
152 }
153
154 let account_proof = match self.proof.proof {
155 RewardMerkleProofV2::Presence(_) => self.proof,
156 RewardMerkleProofV2::Absence(_) => {
157 return Err(RewardClaimError::ZeroRewardError);
158 },
159 };
160
161 Ok(RewardClaimInput {
162 lifetime_rewards: self.balance,
163 auth_data: RewardAuthData::new(account_proof.try_into()?).into(),
164 })
165 }
166}
167
168#[derive(Clone, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
169pub struct Delta {
170 pub fees_delta: HashSet<FeeAccount>,
171 pub rewards_delta: HashSet<RewardAccountV2>,
172}