hotshot_example_types/membership/
randomized_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::collections::{BTreeMap, BTreeSet};
8
9use anyhow::Context;
10use hotshot_types::{
11    drb::{
12        election::{generate_stake_cdf, select_randomized_leader, RandomizedCommittee},
13        DrbResult,
14    },
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)]
24
25/// The randomized stake table election
26pub struct RandomizedStakeTable<
27    PubKey: SignatureKey,
28    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
29> {
30    quorum_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
31
32    da_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
33
34    first_epoch: Option<u64>,
35
36    epochs: BTreeSet<u64>,
37
38    drb_results: BTreeMap<u64, DrbResult>,
39
40    /// Stake tables randomized with the DRB, used (only) for leader election
41    randomized_committee: RandomizedCommittee<<PubKey as SignatureKey>::StakeTableEntry>,
42
43    da_committees: TestDaCommittees<PubKey, StatePubKey>,
44}
45
46impl<PubKey, StatePubKey> TestStakeTable<PubKey, StatePubKey>
47    for RandomizedStakeTable<PubKey, StatePubKey>
48where
49    PubKey: SignatureKey,
50    StatePubKey:
51        StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
52{
53    fn new(
54        quorum_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
55        da_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
56    ) -> Self {
57        // We use a constant value of `[0u8; 32]` for the drb, since this is just meant to be used in tests
58        let randomized_committee = generate_stake_cdf(
59            quorum_members
60                .clone()
61                .into_iter()
62                .map(|entry| entry.stake_table_entry)
63                .collect::<Vec<_>>(),
64            [0u8; 32],
65        );
66
67        Self {
68            quorum_members,
69            da_members,
70            first_epoch: None,
71            randomized_committee,
72            epochs: BTreeSet::new(),
73            drb_results: BTreeMap::new(),
74            da_committees: TestDaCommittees::new(),
75        }
76    }
77
78    fn stake_table(&self, _epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
79        self.quorum_members.clone()
80    }
81
82    fn da_stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
83        self.da_committees
84            .get(epoch)
85            .unwrap_or(self.da_members.clone())
86    }
87
88    fn full_stake_table(&self) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>> {
89        self.quorum_members.clone()
90    }
91
92    fn lookup_leader(&self, view_number: u64, _epoch: Option<u64>) -> anyhow::Result<PubKey> {
93        let res = select_randomized_leader(&self.randomized_committee, view_number);
94
95        Ok(PubKey::public_key(&res))
96    }
97
98    fn has_stake_table(&self, epoch: u64) -> bool {
99        self.epochs.contains(&epoch)
100    }
101
102    fn has_randomized_stake_table(&self, epoch: u64) -> anyhow::Result<bool> {
103        Ok(self.drb_results.contains_key(&epoch))
104    }
105
106    fn add_epoch_root(&mut self, epoch: u64) {
107        self.epochs.insert(epoch);
108    }
109
110    fn add_drb_result(&mut self, epoch: u64, drb_result: DrbResult) {
111        self.drb_results.insert(epoch, drb_result);
112    }
113
114    fn set_first_epoch(&mut self, epoch: u64, initial_drb_result: DrbResult) {
115        self.first_epoch = Some(epoch);
116
117        self.add_epoch_root(epoch);
118        self.add_epoch_root(epoch + 1);
119
120        self.add_drb_result(epoch, initial_drb_result);
121        self.add_drb_result(epoch + 1, initial_drb_result);
122    }
123
124    fn get_epoch_drb(&self, epoch: u64) -> anyhow::Result<DrbResult> {
125        self.drb_results
126            .get(&epoch)
127            .context("DRB result missing")
128            .copied()
129    }
130
131    fn first_epoch(&self) -> Option<u64> {
132        self.first_epoch
133    }
134
135    fn add_da_committee(
136        &mut self,
137        first_epoch: u64,
138        committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
139    ) {
140        self.da_committees.add(first_epoch, committee);
141    }
142}
hotshot_example_types/membership/randomized_committee.rs

hotshot_example_types/membership/
randomized_committee.rs
