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::{TestDaCommittees, 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    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}