Skip to main content

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(|| {
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
68/// Return `true` if any of the given accounts have been forgotten in the `ValidatedState` reward_merkle_tree_v2
69pub 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// New Type for `Address` in order to implement `CanonicalSerialize` and
80// `CanonicalDeserialize`
81// This is the same as `RewardAccountV1` but the `ToTraversal` trait implementation
82// for this type is different
83#[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/// A proof of the balance of an account in the fee ledger.
115///
116/// If the account of interest does not exist in the fee state, this is a Merkle non-membership
117/// proof, and the balance is implicitly zero. Otherwise, this is a normal Merkle membership proof.
118#[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    /// Convert query data to reward claim input for contract submission.
146    ///
147    /// Auth root inputs (other than the reward merkle tree root) are currently
148    /// all zero placeholder values.
149    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}