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/// Per-epoch committee schedule: each entry is the committee effective from
54/// its first epoch (inclusive) until the next scheduled change.
55#[derive(Debug, Clone, PartialEq, Eq, Hash)]
56pub struct TestCommitteeSchedule<
57    PubKey: SignatureKey,
58    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
59>(BTreeMap<u64, Vec<TestStakeTableEntry<PubKey, StatePubKey>>>);
60
61impl<
62    PubKey: SignatureKey,
63    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
64> TestCommitteeSchedule<PubKey, StatePubKey>
65{
66    pub fn new() -> Self {
67        Self(BTreeMap::new())
68    }
69
70    pub fn add(
71        &mut self,
72        first_epoch: u64,
73        committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
74    ) {
75        self.0.insert(first_epoch, committee);
76    }
77
78    pub fn get(&self, epoch: Option<u64>) -> Option<Vec<TestStakeTableEntry<PubKey, StatePubKey>>> {
79        if let Some(e) = epoch {
80            // returns the greatest key smaller than or equal to `e`
81            self.0
82                .range((Bound::Included(&0), Bound::Included(&e)))
83                .last()
84                .map(|(_, committee)| committee)
85                .cloned()
86        } else {
87            None
88        }
89    }
90}
91
92impl<
93    PubKey: SignatureKey,
94    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
95> Default for TestCommitteeSchedule<PubKey, StatePubKey>
96{
97    fn default() -> Self {
98        Self::new()
99    }
100}
101
102pub trait TestStakeTable<
103    PubKey: SignatureKey,
104    StatePubKey: StateSignatureKey + LCV1StateSignatureKey + LCV2StateSignatureKey + LCV3StateSignatureKey,
105>: Clone + Debug + std::marker::Send + std::marker::Sync
106{
107    fn new(
108        quorum_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
109        da_members: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
110    ) -> Self;
111
112    fn stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>>;
113
114    fn full_stake_table(&self) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>>;
115
116    fn da_stake_table(&self, epoch: Option<u64>) -> Vec<TestStakeTableEntry<PubKey, StatePubKey>>;
117
118    fn stake(
119        &self,
120        pub_key: PubKey,
121        epoch: Option<u64>,
122    ) -> Option<TestStakeTableEntry<PubKey, StatePubKey>> {
123        self.stake_table(epoch)
124            .iter()
125            .find(|entry| entry.signature_key == pub_key)
126            .cloned()
127    }
128
129    fn da_stake(
130        &self,
131        pub_key: PubKey,
132        epoch: Option<u64>,
133    ) -> Option<TestStakeTableEntry<PubKey, StatePubKey>> {
134        self.da_stake_table(epoch)
135            .iter()
136            .find(|entry| entry.signature_key == pub_key)
137            .cloned()
138    }
139
140    fn lookup_leader(&self, view_number: u64, epoch: Option<u64>) -> anyhow::Result<PubKey>;
141
142    fn has_stake_table(&self, epoch: u64) -> bool;
143
144    fn has_randomized_stake_table(&self, epoch: u64) -> anyhow::Result<bool>;
145
146    fn add_epoch_root(&mut self, epoch: u64);
147
148    fn add_drb_result(&mut self, epoch: u64, drb_result: DrbResult);
149
150    fn set_first_epoch(&mut self, epoch: u64, initial_drb_result: DrbResult);
151
152    fn first_epoch(&self) -> Option<u64>;
153
154    fn get_epoch_drb(&self, epoch: u64) -> anyhow::Result<DrbResult>;
155
156    fn add_da_committee(
157        &mut self,
158        first_epoch: u64,
159        committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
160    );
161
162    /// Register a quorum committee effective from `first_epoch` (inclusive).
163    /// Mirrors `add_da_committee`; implementations without per-epoch quorum
164    /// committees keep the panicking default.
165    fn add_quorum_committee(
166        &mut self,
167        _first_epoch: u64,
168        _committee: Vec<TestStakeTableEntry<PubKey, StatePubKey>>,
169    ) {
170        panic!("add_quorum_committee is not supported by this TestStakeTable implementation");
171    }
172}
hotshot_example_types/membership/stake_table.rs

hotshot_example_types/membership/
stake_table.rs
