espresso_types/v0/v0_4/
state.rs

1use 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
58/// Return `true` if any of the given accounts have been forgotten in the `ValidatedState` reward_merkle_tree_v2
59pub 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// New Type for `Address` in order to implement `CanonicalSerialize` and
70// `CanonicalDeserialize`
71// This is the same as `RewardAccountV1` but the `ToTraversal` trait implementation
72// for this type is different
73#[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/// A proof of the balance of an account in the fee ledger.
105///
106/// If the account of interest does not exist in the fee state, this is a Merkle non-membership
107/// proof, and the balance is implicitly zero. Otherwise, this is a normal Merkle membership proof.
108#[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    /// Convert query data to reward claim input for contract submission.
136    ///
137    /// Auth root inputs (other than the reward merkle tree root) are currently
138    /// all zero placeholder values.
139    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}