Skip to main content

hotshot_example_types/membership/
static_committee.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7use 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)]
24/// Static stake table that doesn't use DRB results for leader election
25pub 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    /// Epochs before the earliest scheduled first-epoch (and `None`) resolve
171    /// to the default committee; later epochs resolve to the greatest
172    /// scheduled first-epoch <= the requested epoch.
173    #[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    /// The leader rotation uses the scheduled committee for the epoch, so a
193    /// committee of a different size changes which key leads a given view.
194    #[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    /// The quorum and DA schedules are stored and resolved independently.
213    #[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}