hotshot_example_types/membership/
stake_table.rs

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

hotshot_example_types/membership/
stake_table.rs
