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