hotshot_example_types/membership/
two_static_committees.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::collections::{BTreeMap, BTreeSet};
8
9use anyhow::Context;
10use hotshot_types::{
11    drb::DrbResult,
12    traits::signature_key::{
13        LCV1StateSignatureKey, LCV2StateSignatureKey, LCV3StateSignatureKey, SignatureKey,
14        StateSignatureKey,
15    },
16};
17
18use crate::membership::stake_table::{TestStakeTable, TestStakeTableEntry};
19
20#[derive(Clone, Debug, Eq, PartialEq, Hash)]
21pub struct TwoStakeTables<
22    PubKey: SignatureKey,
23    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
24> {
25    quorum_1_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
26
27    da_1_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
28
29    quorum_2_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
30
31    da_2_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
32
33    epochs: BTreeSet<u64>,
34
35    drb_results: BTreeMap<u64, DrbResult>,
36
37    first_epoch: Option<u64>,
38}
39
40impl<PubKey, StatePubKey> TestStakeTable<PubKey, StatePubKey>
41    for TwoStakeTables<PubKey, StatePubKey>
42where
43    PubKey: SignatureKey,
44    StatePubKey:
45        StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
46{
47    fn new(
48        quorum_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
49        da_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
50    ) -> Self {
51        Self {
52            quorum_1_members: quorum_members
53                .iter()
54                .enumerate()
55                .filter(|(idx, _)| idx % 2 == 0)
56                .map(|(_, entry)| entry.clone())
57                .collect(),
58            da_1_members: da_members
59                .iter()
60                .enumerate()
61                .filter(|(idx, _)| idx % 2 == 0)
62                .map(|(_, entry)| entry.clone())
63                .collect(),
64            quorum_2_members: quorum_members
65                .iter()
66                .enumerate()
67                .filter(|(idx, _)| idx % 2 == 1)
68                .map(|(_, entry)| entry.clone())
69                .collect(),
70            da_2_members: da_members
71                .iter()
72                .enumerate()
73                .filter(|(idx, _)| idx % 2 == 1)
74                .map(|(_, entry)| entry.clone())
75                .collect(),
76            first_epoch: None,
77            epochs: BTreeSet::new(),
78            drb_results: BTreeMap::new(),
79        }
80    }
81
82    fn stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
83        let epoch = epoch.expect("epochs cannot be disabled with TwoStakeTables");
84        if epoch != 0 && epoch.is_multiple_of(2) {
85            self.quorum_1_members.clone()
86        } else {
87            self.quorum_2_members.clone()
88        }
89    }
90
91    fn da_stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
92        let epoch = epoch.expect("epochs cannot be disabled with TwoStakeTables");
93        if epoch != 0 && epoch.is_multiple_of(2) {
94            self.da_1_members.clone()
95        } else {
96            self.da_2_members.clone()
97        }
98    }
99
100    fn full_stake_table(&self) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
101        self.quorum_1_members
102            .iter()
103            .chain(self.quorum_2_members.iter())
104            .cloned()
105            .collect()
106    }
107
108    fn lookup_leader(&self, view_number: u64, epoch: Option<u64>) -> anyhow::Result<PubKey> {
109        let stake_table = self.stake_table(epoch);
110
111        let index = view_number as usize % stake_table.len();
112        let leader = stake_table[index].clone();
113
114        Ok(leader.signature_key)
115    }
116
117    fn has_stake_table(&self, epoch: u64) -> bool {
118        self.epochs.contains(&epoch)
119    }
120
121    fn has_randomized_stake_table(&self, epoch: u64) -> anyhow::Result<bool> {
122        Ok(self.drb_results.contains_key(&epoch))
123    }
124
125    fn add_epoch_root(&mut self, epoch: u64) {
126        self.epochs.insert(epoch);
127    }
128
129    fn add_drb_result(&mut self, epoch: u64, drb_result: DrbResult) {
130        self.drb_results.insert(epoch, drb_result);
131    }
132
133    fn set_first_epoch(&mut self, epoch: u64, initial_drb_result: DrbResult) {
134        self.first_epoch = Some(epoch);
135
136        self.add_epoch_root(epoch);
137        self.add_epoch_root(epoch + 1);
138
139        self.add_drb_result(epoch, initial_drb_result);
140        self.add_drb_result(epoch + 1, initial_drb_result);
141    }
142
143    fn get_epoch_drb(&self, epoch: u64) -> anyhow::Result<DrbResult> {
144        self.drb_results
145            .get(&epoch)
146            .context("DRB result missing")
147            .copied()
148    }
149
150    fn first_epoch(&self) -> Option<u64> {
151        self.first_epoch
152    }
153
154    fn add_da_committee(
155        &mut self,
156        _first_epoch: u64,
157        _committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
158    ) {
159        panic!("Dynamic DA committees are not supported for TwoStakeTables")
160    }
161}
hotshot_example_types/membership/two_static_committees.rs

hotshot_example_types/membership/
two_static_committees.rs
