hotshot_example_types/membership/
static_committee.rs1use std::{
8 collections::{BTreeMap, BTreeSet},
9 fmt::Debug,
10};
11
12use anyhow::Context;
13use hotshot_types::{
14 drb::DrbResult,
15 traits::signature_key::{
16 LCV1StateSignatureKey, LCV2StateSignatureKey, LCV3StateSignatureKey, SignatureKey,
17 StateSignatureKey,
18 },
19};
20
21use crate::membership::stake_table::{TestDaCommittees, TestStakeTable, TestStakeTableEntry};
22
23#[derive(Clone, Debug, Eq, PartialEq, Hash)]
24pub struct StaticStakeTable<
26 PubKey: SignatureKey,
27 StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
28> {
29 quorum_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
30
31 da_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
32
33 first_epoch: Option<u64>,
34
35 epochs: BTreeSet<u64>,
36
37 drb_results: BTreeMap<u64, DrbResult>,
38
39 da_committees: TestDaCommittees<PubKey, StatePubKey>,
40}
41
42impl<PubKey, StatePubKey> TestStakeTable<PubKey, StatePubKey>
43 for StaticStakeTable<PubKey, StatePubKey>
44where
45 PubKey: SignatureKey,
46 StatePubKey:
47 StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
48{
49 fn new(
50 quorum_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
51 da_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
52 ) -> Self {
53 Self {
54 quorum_members,
55 da_members,
56 first_epoch: None,
57 epochs: BTreeSet::new(),
58 drb_results: BTreeMap::new(),
59 da_committees: TestDaCommittees::new(),
60 }
61 }
62
63 fn stake_table(&self, _epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
64 self.quorum_members.clone()
65 }
66
67 fn da_stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
68 self.da_committees
69 .get(epoch)
70 .unwrap_or(self.da_members.clone())
71 }
72
73 fn full_stake_table(&self) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
74 self.quorum_members.clone()
75 }
76
77 fn lookup_leader(&self, view_number: u64, _epoch: Option<u64>) -> anyhow::Result<PubKey> {
78 let index = view_number as usize % self.quorum_members.len();
79 let leader = self.quorum_members[index].clone();
80 Ok(leader.signature_key)
81 }
82
83 fn has_stake_table(&self, epoch: u64) -> bool {
84 self.epochs.contains(&epoch)
85 }
86
87 fn has_randomized_stake_table(&self, epoch: u64) -> anyhow::Result<bool> {
88 Ok(self.drb_results.contains_key(&epoch))
89 }
90
91 fn add_epoch_root(&mut self, epoch: u64) {
92 self.epochs.insert(epoch);
93 }
94
95 fn add_drb_result(&mut self, epoch: u64, drb_result: DrbResult) {
96 self.drb_results.insert(epoch, drb_result);
97 }
98
99 fn set_first_epoch(&mut self, epoch: u64, initial_drb_result: DrbResult) {
100 self.first_epoch = Some(epoch);
101
102 self.add_epoch_root(epoch);
103 self.add_epoch_root(epoch + 1);
104
105 self.add_drb_result(epoch, initial_drb_result);
106 self.add_drb_result(epoch + 1, initial_drb_result);
107 }
108
109 fn get_epoch_drb(&self, epoch: u64) -> anyhow::Result<DrbResult> {
110 self.drb_results
111 .get(&epoch)
112 .context("DRB result missing")
113 .copied()
114 }
115
116 fn first_epoch(&self) -> Option<u64> {
117 self.first_epoch
118 }
119
120 fn add_da_committee(
121 &mut self,
122 first_epoch: u64,
123 committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
124 ) {
125 self.da_committees.add(first_epoch, committee);
126 }
127}