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::{TestCommitteeSchedule, 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 quorum_committees: TestCommitteeSchedule<PubKey, StatePubKey>,
34
35 first_epoch: Option<u64>,
36
37 epochs: BTreeSet<u64>,
38
39 drb_results: BTreeMap<u64, DrbResult>,
40
41 da_committees: TestCommitteeSchedule<PubKey, StatePubKey>,
42}
43
44impl<PubKey, StatePubKey> TestStakeTable<PubKey, StatePubKey>
45 for StaticStakeTable<PubKey, StatePubKey>
46where
47 PubKey: SignatureKey,
48 StatePubKey:
49 StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
50{
51 fn new(
52 quorum_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
53 da_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
54 ) -> Self {
55 Self {
56 quorum_members,
57 da_members,
58 quorum_committees: TestCommitteeSchedule::new(),
59 first_epoch: None,
60 epochs: BTreeSet::new(),
61 drb_results: BTreeMap::new(),
62 da_committees: TestCommitteeSchedule::new(),
63 }
64 }
65
66 fn stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
67 self.quorum_committees
68 .get(epoch)
69 .unwrap_or_else(|| self.quorum_members.clone())
70 }
71
72 fn da_stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
73 self.da_committees
74 .get(epoch)
75 .unwrap_or(self.da_members.clone())
76 }
77
78 fn full_stake_table(&self) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
79 self.quorum_members.clone()
80 }
81
82 fn lookup_leader(&self, view_number: u64, epoch: Option<u64>) -> anyhow::Result<PubKey> {
83 let committee = self.stake_table(epoch);
84 let index = view_number as usize % committee.len();
85 let leader = committee[index].clone();
86 Ok(leader.signature_key)
87 }
88
89 fn has_stake_table(&self, epoch: u64) -> bool {
90 self.epochs.contains(&epoch)
91 }
92
93 fn has_randomized_stake_table(&self, epoch: u64) -> anyhow::Result<bool> {
94 Ok(self.drb_results.contains_key(&epoch))
95 }
96
97 fn add_epoch_root(&mut self, epoch: u64) {
98 self.epochs.insert(epoch);
99 }
100
101 fn add_drb_result(&mut self, epoch: u64, drb_result: DrbResult) {
102 self.drb_results.insert(epoch, drb_result);
103 }
104
105 fn set_first_epoch(&mut self, epoch: u64, initial_drb_result: DrbResult) {
106 self.first_epoch = Some(epoch);
107
108 self.add_epoch_root(epoch);
109 self.add_epoch_root(epoch + 1);
110
111 self.add_drb_result(epoch, initial_drb_result);
112 self.add_drb_result(epoch + 1, initial_drb_result);
113 }
114
115 fn get_epoch_drb(&self, epoch: u64) -> anyhow::Result<DrbResult> {
116 self.drb_results
117 .get(&epoch)
118 .context("DRB result missing")
119 .copied()
120 }
121
122 fn first_epoch(&self) -> Option<u64> {
123 self.first_epoch
124 }
125
126 fn add_da_committee(
127 &mut self,
128 first_epoch: u64,
129 committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
130 ) {
131 self.da_committees.add(first_epoch, committee);
132 }
133
134 fn add_quorum_committee(
135 &mut self,
136 first_epoch: u64,
137 committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
138 ) {
139 self.quorum_committees.add(first_epoch, committee);
140 }
141}
142
143#[cfg(test)]
144mod tests {
145 use alloy::primitives::U256;
146 use hotshot_types::{
147 ValidatorConfig,
148 signature_key::{BLSPubKey, SchnorrPubKey},
149 };
150
151 use super::*;
152 use crate::node_types::TestTypes;
153
154 fn entries(indices: &[u64]) -> Vec<TestStakeTableEntry<BLSPubKey, SchnorrPubKey>> {
155 indices
156 .iter()
157 .map(|i| {
158 let config: ValidatorConfig<TestTypes> =
159 ValidatorConfig::generated_from_seed_indexed(
160 [0u8; 32],
161 *i,
162 U256::from(1),
163 false,
164 );
165 config.public_config().into()
166 })
167 .collect()
168 }
169
170 #[test]
174 fn quorum_schedule_resolution() {
175 let default = entries(&[0, 1, 2, 3, 4]);
176 let smaller = entries(&[0, 1, 2, 3]);
177 let larger = entries(&[0, 1, 2, 3, 4, 5]);
178
179 let mut table = StaticStakeTable::new(default.clone(), default.clone());
180 table.add_quorum_committee(3, smaller.clone());
181 table.add_quorum_committee(6, larger.clone());
182
183 assert_eq!(table.stake_table(None), default);
184 assert_eq!(table.stake_table(Some(1)), default);
185 assert_eq!(table.stake_table(Some(2)), default);
186 assert_eq!(table.stake_table(Some(3)), smaller);
187 assert_eq!(table.stake_table(Some(5)), smaller);
188 assert_eq!(table.stake_table(Some(6)), larger);
189 assert_eq!(table.stake_table(Some(100)), larger);
190 }
191
192 #[test]
195 fn quorum_schedule_leader_rotation() {
196 let default = entries(&[0, 1, 2, 3, 4]);
197 let scheduled = entries(&[0, 1, 2]);
198
199 let mut table = StaticStakeTable::new(default.clone(), default.clone());
200 table.add_quorum_committee(3, scheduled.clone());
201
202 assert_eq!(
203 table.lookup_leader(4, Some(2)).unwrap(),
204 default[4].signature_key
205 );
206 assert_eq!(
207 table.lookup_leader(4, Some(3)).unwrap(),
208 scheduled[4 % 3].signature_key
209 );
210 }
211
212 #[test]
214 fn quorum_and_da_schedules_independent() {
215 let default = entries(&[0, 1, 2, 3, 4]);
216 let quorum = entries(&[0, 1, 2, 3]);
217 let da = entries(&[0, 1]);
218
219 let mut table = StaticStakeTable::new(default.clone(), default.clone());
220 table.add_quorum_committee(3, quorum.clone());
221 table.add_da_committee(4, da.clone());
222
223 assert_eq!(table.stake_table(Some(3)), quorum);
224 assert_eq!(table.da_stake_table(Some(3)), default);
225 assert_eq!(table.stake_table(Some(4)), quorum);
226 assert_eq!(table.da_stake_table(Some(4)), da);
227 }
228}