hotshot_example_types/membership/
stake_table.rs

1use std::{collections::BTreeMap, fmt::Debug, ops::Bound};
2
3use hotshot_types::{
4    drb::DrbResult,
5    traits::{
6        node_implementation::NodeType,
7        signature_key::{
8            LCV1StateSignatureKey, LCV2StateSignatureKey, LCV3StateSignatureKey, SignatureKey,
9            StateSignatureKey,
10        },
11    },
12    PeerConfig,
13};
14
15#[derive(Clone, Debug, Eq, PartialEq, Hash)]
16pub struct TestStakeTableEntry<
17    PubKey: SignatureKey,
18    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
19> {
20    pub signature_key: PubKey,
21    pub stake_table_entry: <PubKey as SignatureKey>::StakeTableEntry,
22    pub state_ver_key: StatePubKey,
23}
24
25impl<TYPES: NodeType> From<PeerConfig<TYPES>>
26    for TestStakeTableEntry<TYPES::SignatureKey, TYPES::StateSignatureKey>
27{
28    fn from(peer_config: PeerConfig<TYPES>) -> Self {
29        Self {
30            signature_key: SignatureKey::public_key(&peer_config.stake_table_entry),
31            stake_table_entry: peer_config.stake_table_entry,
32            state_ver_key: peer_config.state_ver_key,
33        }
34    }
35}
36
37impl<TYPES: NodeType> From<TestStakeTableEntry<TYPES::SignatureKey, TYPES::StateSignatureKey>>
38    for PeerConfig<TYPES>
39{
40    fn from(
41        test_stake_table_entry: TestStakeTableEntry<TYPES::SignatureKey, TYPES::StateSignatureKey>,
42    ) -> Self {
43        PeerConfig {
44            stake_table_entry: test_stake_table_entry.stake_table_entry,
45            state_ver_key: test_stake_table_entry.state_ver_key,
46        }
47    }
48}
49
50// Map from first epoch to DA committee stake table entries
51#[derive(Debug, Clone, PartialEq, Eq, Hash)]
52pub struct TestDaCommittees<
53    PubKey: SignatureKey,
54    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
55>(BTreeMap<u64, Vec<TestStakeTableEntry<PubKey, StatePubKey>>>);
56
57impl<
58        PubKey: SignatureKey,
59        StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
60    > TestDaCommittees<PubKey, StatePubKey>
61{
62    pub fn new() -> Self {
63        Self(BTreeMap::new())
64    }
65
66    pub fn add(
67        &mut self,
68        first_epoch: u64,
69        committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
70    ) {
71        self.0.insert(first_epoch, committee);
72    }
73
74    pub fn get(&self, epoch: Option<u64>) -> Option<Vec<TestStakeTableEntry<PubKey, StatePubKey>>> {
75        if let Some(e) = epoch {
76            // returns the greatest key smaller than or equal to `e`
77            self.0
78                .range((Bound::Included(&0), Bound::Included(&e)))
79                .last()
80                .map(|(_, committee)| committee)
81                .cloned()
82        } else {
83            None
84        }
85    }
86}
87
88impl<
89        PubKey: SignatureKey,
90        StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
91    > Default for TestDaCommittees<PubKey, StatePubKey>
92{
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98pub trait TestStakeTable<
99    PubKey: SignatureKey,
100    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
101>: Debug + std::marker::Send + std::marker::Sync
102{
103    fn new(
104        quorum_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
105        da_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
106    ) -> Self;
107
108    fn stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>>;
109
110    fn full_stake_table(&self) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>>;
111
112    fn da_stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>>;
113
114    fn stake(
115        &self,
116        pub_key: PubKey,
117        epoch: Option<u64>,
118    ) -> Option<TestStakeTableEntry<PubKey, StatePubKey>> {
119        self.stake_table(epoch)
120            .iter()
121            .find(|entry| entry.signature_key == pub_key)
122            .cloned()
123    }
124
125    fn da_stake(
126        &self,
127        pub_key: PubKey,
128        epoch: Option<u64>,
129    ) -> Option<TestStakeTableEntry<PubKey, StatePubKey>> {
130        self.da_stake_table(epoch)
131            .iter()
132            .find(|entry| entry.signature_key == pub_key)
133            .cloned()
134    }
135
136    fn lookup_leader(&self, view_number: u64, epoch: Option<u64>) -> anyhow::Result<PubKey>;
137
138    fn has_stake_table(&self, epoch: u64) -> bool;
139
140    fn has_randomized_stake_table(&self, epoch: u64) -> anyhow::Result<bool>;
141
142    fn add_epoch_root(&mut self, epoch: u64);
143
144    fn add_drb_result(&mut self, epoch: u64, drb_result: DrbResult);
145
146    fn set_first_epoch(&mut self, epoch: u64, initial_drb_result: DrbResult);
147
148    fn first_epoch(&self) -> Option<u64>;
149
150    fn get_epoch_drb(&self, epoch: u64) -> anyhow::Result<DrbResult>;
151
152    fn add_da_committee(
153        &mut self,
154        first_epoch: u64,
155        committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
156    );
157}
hotshot_example_types/membership/stake_table.rs

hotshot_example_types/membership/
stake_table.rs
