1use std::{collections::BTreeMap, sync::Arc};
2
3use alloy::primitives::Address;
4use anyhow::bail;
5#[cfg(any(test, feature = "testing"))]
6use async_lock::RwLock;
7use async_trait::async_trait;
8use hotshot_types::{
9 data::EpochNumber, epoch_membership::EpochMembershipCoordinator, traits::states::InstanceState,
10 HotShotConfig,
11};
12#[cfg(any(test, feature = "testing"))]
13use vbs::version::StaticVersionType;
14use vbs::version::Version;
15
16use super::{
17 state::ValidatedState,
18 traits::{EventsPersistenceRead, MembershipPersistence},
19 v0_1::NoStorage,
20 v0_3::{EventKey, IndexedStake, StakeTableEvent},
21 SeqTypes, UpgradeType, ViewBasedUpgrade,
22};
23#[cfg(any(test, feature = "testing"))]
24use crate::EpochCommittees;
25use crate::{
26 v0::{
27 traits::StateCatchup, v0_3::ChainConfig, GenesisHeader, L1BlockInfo, L1Client, Timestamp,
28 Upgrade, UpgradeMode,
29 },
30 v0_1::RewardAmount,
31 ValidatorMap,
32};
33
34#[derive(derive_more::Debug, Clone)]
38pub struct NodeState {
39 pub node_id: u64,
40 pub chain_config: ChainConfig,
41 pub l1_client: L1Client,
42 #[debug("{}", state_catchup.name())]
43 pub state_catchup: Arc<dyn StateCatchup>,
44 pub genesis_header: GenesisHeader,
45 pub genesis_state: ValidatedState,
46 pub l1_genesis: Option<L1BlockInfo>,
47 #[debug(skip)]
48 pub coordinator: EpochMembershipCoordinator<SeqTypes>,
49 pub epoch_height: Option<u64>,
50 pub genesis_version: Version,
51
52 pub upgrades: BTreeMap<Version, Upgrade>,
61 pub current_version: Version,
68}
69
70impl NodeState {
71 pub async fn block_reward(&self) -> RewardAmount {
72 let coordinator = self.coordinator.clone();
73 let membership = coordinator.membership().read().await;
74 membership.block_reward()
75 }
76}
77
78#[async_trait]
79impl MembershipPersistence for NoStorage {
80 async fn load_stake(&self, _epoch: EpochNumber) -> anyhow::Result<Option<ValidatorMap>> {
81 Ok(None)
82 }
83
84 async fn load_latest_stake(&self, _limit: u64) -> anyhow::Result<Option<Vec<IndexedStake>>> {
85 Ok(None)
86 }
87
88 async fn store_stake(&self, _epoch: EpochNumber, _stake: ValidatorMap) -> anyhow::Result<()> {
89 Ok(())
90 }
91
92 async fn store_events(
93 &self,
94 _l1_finalized: u64,
95 _events: Vec<(EventKey, StakeTableEvent)>,
96 ) -> anyhow::Result<()> {
97 Ok(())
98 }
99
100 async fn load_events(
101 &self,
102 _l1_block: u64,
103 ) -> anyhow::Result<(
104 Option<EventsPersistenceRead>,
105 Vec<(EventKey, StakeTableEvent)>,
106 )> {
107 bail!("unimplemented")
108 }
109}
110
111impl NodeState {
112 pub fn new(
113 node_id: u64,
114 chain_config: ChainConfig,
115 l1_client: L1Client,
116 catchup: impl StateCatchup + 'static,
117 current_version: Version,
118 coordinator: EpochMembershipCoordinator<SeqTypes>,
119 genesis_version: Version,
120 ) -> Self {
121 Self {
122 node_id,
123 chain_config,
124 l1_client,
125 state_catchup: Arc::new(catchup),
126 genesis_header: Default::default(),
127 genesis_state: ValidatedState {
128 chain_config: chain_config.into(),
129 ..Default::default()
130 },
131 l1_genesis: None,
132 upgrades: Default::default(),
133 current_version,
134 epoch_height: None,
135 coordinator,
136 genesis_version,
137 }
138 }
139
140 #[cfg(any(test, feature = "testing"))]
141 pub fn mock() -> Self {
142 use alloy::primitives::U256;
143 use hotshot_example_types::storage_types::TestStorage;
144 use vbs::version::StaticVersion;
145
146 use crate::{v0_1::RewardAmount, v0_3::Fetcher};
147
148 let chain_config = ChainConfig::default();
149 let l1 = L1Client::new(vec!["http://localhost:3331".parse().unwrap()])
150 .expect("Failed to create L1 client");
151
152 let membership = Arc::new(RwLock::new(EpochCommittees::new_stake(
153 vec![],
154 vec![],
155 RewardAmount(U256::ZERO),
156 Fetcher::mock(),
157 )));
158
159 let storage = TestStorage::default();
160 let coordinator = EpochMembershipCoordinator::new(membership, 100, &storage);
161 Self::new(
162 0,
163 chain_config,
164 l1,
165 Arc::new(mock::MockStateCatchup::default()),
166 StaticVersion::<0, 1>::version(),
167 coordinator,
168 Version { major: 0, minor: 1 },
169 )
170 }
171
172 #[cfg(any(test, feature = "testing"))]
173 pub fn mock_v2() -> Self {
174 use alloy::primitives::U256;
175 use hotshot_example_types::storage_types::TestStorage;
176 use vbs::version::StaticVersion;
177
178 use crate::{v0_1::RewardAmount, v0_3::Fetcher};
179
180 let chain_config = ChainConfig::default();
181 let l1 = L1Client::new(vec!["http://localhost:3331".parse().unwrap()])
182 .expect("Failed to create L1 client");
183
184 let membership = Arc::new(RwLock::new(EpochCommittees::new_stake(
185 vec![],
186 vec![],
187 RewardAmount(U256::ZERO),
188 Fetcher::mock(),
189 )));
190 let storage = TestStorage::default();
191 let coordinator = EpochMembershipCoordinator::new(membership, 100, &storage);
192
193 Self::new(
194 0,
195 chain_config,
196 l1,
197 Arc::new(mock::MockStateCatchup::default()),
198 StaticVersion::<0, 2>::version(),
199 coordinator,
200 Version { major: 0, minor: 1 },
201 )
202 }
203
204 #[cfg(any(test, feature = "testing"))]
205 pub fn mock_v3() -> Self {
206 use alloy::primitives::U256;
207 use hotshot_example_types::storage_types::TestStorage;
208 use vbs::version::StaticVersion;
209
210 use crate::{v0_1::RewardAmount, v0_3::Fetcher};
211 let l1 = L1Client::new(vec!["http://localhost:3331".parse().unwrap()])
212 .expect("Failed to create L1 client");
213
214 let membership = Arc::new(RwLock::new(EpochCommittees::new_stake(
215 vec![],
216 vec![],
217 RewardAmount(U256::ZERO),
218 Fetcher::mock(),
219 )));
220
221 let storage = TestStorage::default();
222 let coordinator = EpochMembershipCoordinator::new(membership, 100, &storage);
223 Self::new(
224 0,
225 ChainConfig::default(),
226 l1,
227 mock::MockStateCatchup::default(),
228 StaticVersion::<0, 3>::version(),
229 coordinator,
230 Version { major: 0, minor: 1 },
231 )
232 }
233
234 pub fn with_l1(mut self, l1_client: L1Client) -> Self {
235 self.l1_client = l1_client;
236 self
237 }
238
239 pub fn with_genesis(mut self, state: ValidatedState) -> Self {
240 self.genesis_state = state;
241 self
242 }
243
244 pub fn with_chain_config(mut self, cfg: ChainConfig) -> Self {
245 self.chain_config = cfg;
246 self
247 }
248
249 pub fn with_upgrades(mut self, upgrades: BTreeMap<Version, Upgrade>) -> Self {
250 self.upgrades = upgrades;
251 self
252 }
253
254 pub fn with_current_version(mut self, version: Version) -> Self {
255 self.current_version = version;
256 self
257 }
258
259 pub fn with_epoch_height(mut self, epoch_height: u64) -> Self {
260 self.epoch_height = Some(epoch_height);
261 self
262 }
263}
264
265pub struct UpgradeMap(pub BTreeMap<Version, Upgrade>);
267impl UpgradeMap {
268 pub fn chain_config(&self, version: Version) -> ChainConfig {
269 self.0
270 .get(&version)
271 .unwrap()
272 .upgrade_type
273 .chain_config()
274 .unwrap()
275 }
276}
277
278impl From<BTreeMap<Version, Upgrade>> for UpgradeMap {
279 fn from(inner: BTreeMap<Version, Upgrade>) -> Self {
280 Self(inner)
281 }
282}
283
284#[cfg(any(test, feature = "testing"))]
287impl Default for NodeState {
288 fn default() -> Self {
289 use alloy::primitives::U256;
290 use hotshot_example_types::storage_types::TestStorage;
291 use vbs::version::StaticVersion;
292
293 use crate::{v0_1::RewardAmount, v0_3::Fetcher};
294
295 let chain_config = ChainConfig::default();
296 let l1 = L1Client::new(vec!["http://localhost:3331".parse().unwrap()])
297 .expect("Failed to create L1 client");
298
299 let membership = Arc::new(RwLock::new(EpochCommittees::new_stake(
300 vec![],
301 vec![],
302 RewardAmount(U256::ZERO),
303 Fetcher::mock(),
304 )));
305 let storage = TestStorage::default();
306 let coordinator = EpochMembershipCoordinator::new(membership, 100, &storage);
307
308 Self::new(
309 1u64,
310 chain_config,
311 l1,
312 Arc::new(mock::MockStateCatchup::default()),
313 StaticVersion::<0, 1>::version(),
314 coordinator,
315 Version { major: 0, minor: 1 },
316 )
317 }
318}
319
320impl InstanceState for NodeState {}
321
322impl Upgrade {
323 pub fn set_hotshot_config_parameters(&self, config: &mut HotShotConfig<SeqTypes>) {
324 match &self.mode {
325 UpgradeMode::View(v) => {
326 config.start_proposing_view = v.start_proposing_view;
327 config.stop_proposing_view = v.stop_proposing_view;
328 config.start_voting_view = v.start_voting_view.unwrap_or(0);
329 config.stop_voting_view = v.stop_voting_view.unwrap_or(u64::MAX);
330 config.start_proposing_time = 0;
331 config.stop_proposing_time = u64::MAX;
332 config.start_voting_time = 0;
333 config.stop_voting_time = u64::MAX;
334 },
335 UpgradeMode::Time(t) => {
336 config.start_proposing_time = t.start_proposing_time.unix_timestamp();
337 config.stop_proposing_time = t.stop_proposing_time.unix_timestamp();
338 config.start_voting_time = t.start_voting_time.unwrap_or_default().unix_timestamp();
339 config.stop_voting_time = t
340 .stop_voting_time
341 .unwrap_or(Timestamp::max())
342 .unix_timestamp();
343 config.start_proposing_view = 0;
344 config.stop_proposing_view = u64::MAX;
345 config.start_voting_view = 0;
346 config.stop_voting_view = u64::MAX;
347 },
348 }
349 }
350 pub fn pos_view_based(address: Address) -> Upgrade {
351 let chain_config = ChainConfig {
352 base_fee: 0.into(),
353 stake_table_contract: Some(address),
354 ..Default::default()
355 };
356
357 let mode = UpgradeMode::View(ViewBasedUpgrade {
358 start_voting_view: None,
359 stop_voting_view: None,
360 start_proposing_view: 200,
361 stop_proposing_view: 1000,
362 });
363
364 let upgrade_type = UpgradeType::Epoch { chain_config };
365 Upgrade { mode, upgrade_type }
366 }
367}
368
369#[cfg(any(test, feature = "testing"))]
370pub mod mock {
371 use std::collections::HashMap;
372
373 use alloy::primitives::U256;
374 use anyhow::Context;
375 use async_trait::async_trait;
376 use committable::Commitment;
377 use hotshot_types::{data::ViewNumber, stake_table::HSStakeTable};
378 use jf_merkle_tree::{ForgetableMerkleTreeScheme, MerkleTreeScheme};
379
380 use super::*;
381 use crate::{
382 retain_accounts,
383 v0_1::{RewardAccount, RewardAccountProof, RewardMerkleCommitment},
384 BackoffParams, BlockMerkleTree, FeeAccount, FeeAccountProof, FeeMerkleCommitment, Leaf2,
385 };
386
387 #[derive(Debug, Clone, Default)]
388 pub struct MockStateCatchup {
389 backoff: BackoffParams,
390 state: HashMap<ViewNumber, Arc<ValidatedState>>,
391 }
392
393 impl FromIterator<(ViewNumber, Arc<ValidatedState>)> for MockStateCatchup {
394 fn from_iter<I: IntoIterator<Item = (ViewNumber, Arc<ValidatedState>)>>(iter: I) -> Self {
395 Self {
396 backoff: Default::default(),
397 state: iter.into_iter().collect(),
398 }
399 }
400 }
401
402 #[async_trait]
403 impl StateCatchup for MockStateCatchup {
404 async fn try_fetch_leaf(
405 &self,
406 _retry: usize,
407 _height: u64,
408 _stake_table: HSStakeTable<SeqTypes>,
409 _success_threshold: U256,
410 ) -> anyhow::Result<Leaf2> {
411 Err(anyhow::anyhow!("todo"))
412 }
413
414 async fn try_fetch_accounts(
415 &self,
416 _retry: usize,
417 _instance: &NodeState,
418 _height: u64,
419 view: ViewNumber,
420 fee_merkle_tree_root: FeeMerkleCommitment,
421 accounts: &[FeeAccount],
422 ) -> anyhow::Result<Vec<FeeAccountProof>> {
423 let src = &self.state[&view].fee_merkle_tree;
424 assert_eq!(src.commitment(), fee_merkle_tree_root);
425
426 tracing::info!("catchup: fetching accounts {accounts:?} for view {view}");
427 let tree = retain_accounts(src, accounts.iter().copied())
428 .with_context(|| "failed to retain accounts")?;
429
430 let mut proofs = Vec::new();
432 for account in accounts {
433 let (proof, _) = FeeAccountProof::prove(&tree, (*account).into())
434 .context(format!("response missing fee account {account}"))?;
435 proof
436 .verify(&fee_merkle_tree_root)
437 .context(format!("invalid proof for fee account {account}"))?;
438 proofs.push(proof);
439 }
440
441 Ok(proofs)
442 }
443
444 async fn try_remember_blocks_merkle_tree(
445 &self,
446 _retry: usize,
447 _instance: &NodeState,
448 _height: u64,
449 view: ViewNumber,
450 mt: &mut BlockMerkleTree,
451 ) -> anyhow::Result<()> {
452 tracing::info!("catchup: fetching frontier for view {view}");
453 let src = &self.state[&view].block_merkle_tree;
454
455 assert_eq!(src.commitment(), mt.commitment());
456 assert!(
457 src.num_leaves() > 0,
458 "catchup should not be triggered when blocks tree is empty"
459 );
460
461 let index = src.num_leaves() - 1;
462 let (elem, proof) = src.lookup(index).expect_ok().unwrap();
463 mt.remember(index, elem, proof.clone())
464 .expect("Proof verifies");
465
466 Ok(())
467 }
468
469 async fn try_fetch_chain_config(
470 &self,
471 _retry: usize,
472 _commitment: Commitment<ChainConfig>,
473 ) -> anyhow::Result<ChainConfig> {
474 Ok(ChainConfig::default())
475 }
476
477 async fn try_fetch_reward_accounts(
478 &self,
479 _retry: usize,
480 _instance: &NodeState,
481 _height: u64,
482 _view: ViewNumber,
483 _reward_merkle_tree_root: RewardMerkleCommitment,
484 _accounts: &[RewardAccount],
485 ) -> anyhow::Result<Vec<RewardAccountProof>> {
486 anyhow::bail!("unimplemented")
487 }
488
489 fn backoff(&self) -> &BackoffParams {
490 &self.backoff
491 }
492
493 fn name(&self) -> String {
494 "MockStateCatchup".into()
495 }
496
497 fn is_local(&self) -> bool {
498 true
499 }
500 }
501}