polkadot_statement_distribution/v2/
statement_store.rs1use bitvec::{order::Lsb0 as BitOrderLsb0, vec::BitVec};
27use polkadot_node_network_protocol::v3::StatementFilter;
28use polkadot_primitives::{
29 CandidateHash, CompactStatement, GroupIndex, SignedStatement, ValidatorIndex,
30};
31use std::collections::hash_map::{Entry as HEntry, HashMap};
32
33use super::groups::Groups;
34
35pub enum StatementOrigin {
37 Local,
39 Remote,
41}
42
43impl StatementOrigin {
44 fn is_local(&self) -> bool {
45 match *self {
46 StatementOrigin::Local => true,
47 StatementOrigin::Remote => false,
48 }
49 }
50}
51
52struct StoredStatement {
53 statement: SignedStatement,
54 known_by_backing: bool,
55}
56
57pub struct StatementStore {
60 validator_meta: HashMap<ValidatorIndex, ValidatorMeta>,
61
62 group_statements: HashMap<(GroupIndex, CandidateHash), GroupStatements>,
66 known_statements: HashMap<Fingerprint, StoredStatement>,
67}
68
69impl StatementStore {
70 pub fn new(groups: &Groups) -> Self {
72 let mut validator_meta = HashMap::new();
73 for (g, group) in groups.all().iter().enumerate() {
74 for (i, v) in group.iter().enumerate() {
75 validator_meta.insert(
76 *v,
77 ValidatorMeta {
78 seconded_count: 0,
79 within_group_index: i,
80 group: GroupIndex(g as _),
81 },
82 );
83 }
84 }
85
86 StatementStore {
87 validator_meta,
88 group_statements: HashMap::new(),
89 known_statements: HashMap::new(),
90 }
91 }
92
93 pub fn insert(
96 &mut self,
97 groups: &Groups,
98 statement: SignedStatement,
99 origin: StatementOrigin,
100 ) -> Result<bool, Error> {
101 let validator_index = statement.validator_index();
102 let validator_meta = match self.validator_meta.get_mut(&validator_index) {
103 None => return Err(Error::ValidatorUnknown),
104 Some(m) => m,
105 };
106
107 let compact = statement.payload().clone();
108 let fingerprint = (validator_index, compact.clone());
109 match self.known_statements.entry(fingerprint) {
110 HEntry::Occupied(mut e) => {
111 if let StatementOrigin::Local = origin {
112 e.get_mut().known_by_backing = true;
113 }
114
115 return Ok(false)
116 },
117 HEntry::Vacant(e) => {
118 e.insert(StoredStatement { statement, known_by_backing: origin.is_local() });
119 },
120 }
121
122 let candidate_hash = *compact.candidate_hash();
123 let seconded = if let CompactStatement::Seconded(_) = compact { true } else { false };
124
125 {
127 let group_index = validator_meta.group;
128 let group = match groups.get(group_index) {
129 Some(g) => g,
130 None => {
131 gum::error!(
132 target: crate::LOG_TARGET,
133 ?group_index,
134 "groups passed into `insert` differ from those used at store creation"
135 );
136
137 return Err(Error::ValidatorUnknown)
138 },
139 };
140
141 let group_statements = self
142 .group_statements
143 .entry((group_index, candidate_hash))
144 .or_insert_with(|| GroupStatements::with_group_size(group.len()));
145
146 if seconded {
147 validator_meta.seconded_count += 1;
148 group_statements.note_seconded(validator_meta.within_group_index);
149 } else {
150 group_statements.note_validated(validator_meta.within_group_index);
151 }
152 }
153
154 Ok(true)
155 }
156
157 pub fn fill_statement_filter(
160 &self,
161 group_index: GroupIndex,
162 candidate_hash: CandidateHash,
163 statement_filter: &mut StatementFilter,
164 ) {
165 if let Some(statements) = self.group_statements.get(&(group_index, candidate_hash)) {
166 statement_filter.seconded_in_group |= statements.seconded.as_bitslice();
167 statement_filter.validated_in_group |= statements.valid.as_bitslice();
168 }
169 }
170
171 pub fn group_statements<'a>(
176 &'a self,
177 groups: &'a Groups,
178 group_index: GroupIndex,
179 candidate_hash: CandidateHash,
180 filter: &'a StatementFilter,
181 ) -> impl Iterator<Item = &'a SignedStatement> + 'a {
182 let group_validators = groups.get(group_index);
183
184 let seconded_statements = filter
185 .seconded_in_group
186 .iter_ones()
187 .filter_map(move |i| group_validators.as_ref().and_then(|g| g.get(i)))
188 .filter_map(move |v| {
189 self.known_statements.get(&(*v, CompactStatement::Seconded(candidate_hash)))
190 })
191 .map(|s| &s.statement);
192
193 let valid_statements = filter
194 .validated_in_group
195 .iter_ones()
196 .filter_map(move |i| group_validators.as_ref().and_then(|g| g.get(i)))
197 .filter_map(move |v| {
198 self.known_statements.get(&(*v, CompactStatement::Valid(candidate_hash)))
199 })
200 .map(|s| &s.statement);
201
202 seconded_statements.chain(valid_statements)
203 }
204
205 pub fn validator_statement(
207 &self,
208 validator_index: ValidatorIndex,
209 statement: CompactStatement,
210 ) -> Option<&SignedStatement> {
211 self.known_statements.get(&(validator_index, statement)).map(|s| &s.statement)
212 }
213
214 pub fn fresh_statements_for_backing<'a>(
217 &'a self,
218 validators: &'a [ValidatorIndex],
219 candidate_hash: CandidateHash,
220 ) -> impl Iterator<Item = &'a SignedStatement> + 'a {
221 let s_st = CompactStatement::Seconded(candidate_hash);
222 let v_st = CompactStatement::Valid(candidate_hash);
223
224 let fresh_seconded =
225 validators.iter().map(move |v| self.known_statements.get(&(*v, s_st.clone())));
226
227 let fresh_valid =
228 validators.iter().map(move |v| self.known_statements.get(&(*v, v_st.clone())));
229
230 fresh_seconded
231 .chain(fresh_valid)
232 .flatten()
233 .filter(|stored| !stored.known_by_backing)
234 .map(|stored| &stored.statement)
235 }
236
237 pub fn seconded_count(&self, validator_index: &ValidatorIndex) -> usize {
239 self.validator_meta.get(validator_index).map_or(0, |m| m.seconded_count)
240 }
241
242 pub fn note_known_by_backing(
244 &mut self,
245 validator_index: ValidatorIndex,
246 statement: CompactStatement,
247 ) {
248 if let Some(stored) = self.known_statements.get_mut(&(validator_index, statement)) {
249 stored.known_by_backing = true;
250 }
251 }
252}
253
254#[derive(Debug)]
256pub enum Error {
257 ValidatorUnknown,
259}
260
261type Fingerprint = (ValidatorIndex, CompactStatement);
262
263struct ValidatorMeta {
264 group: GroupIndex,
265 within_group_index: usize,
266 seconded_count: usize,
267}
268
269struct GroupStatements {
270 seconded: BitVec<u8, BitOrderLsb0>,
271 valid: BitVec<u8, BitOrderLsb0>,
272}
273
274impl GroupStatements {
275 fn with_group_size(group_size: usize) -> Self {
276 GroupStatements {
277 seconded: BitVec::repeat(false, group_size),
278 valid: BitVec::repeat(false, group_size),
279 }
280 }
281
282 fn note_seconded(&mut self, within_group_index: usize) {
283 self.seconded.set(within_group_index, true);
284 }
285
286 fn note_validated(&mut self, within_group_index: usize) {
287 self.valid.set(within_group_index, true);
288 }
289}
290
291#[cfg(test)]
292mod tests {
293 use super::*;
294
295 use polkadot_primitives::{Hash, SigningContext, ValidatorPair};
296 use sp_application_crypto::Pair as PairT;
297
298 #[test]
299 fn always_provides_fresh_statements_in_order() {
300 let validator_a = ValidatorIndex(1);
301 let validator_b = ValidatorIndex(2);
302 let candidate_hash = CandidateHash(Hash::repeat_byte(42));
303
304 let valid_statement = CompactStatement::Valid(candidate_hash);
305 let seconded_statement = CompactStatement::Seconded(candidate_hash);
306 let signing_context =
307 SigningContext { parent_hash: Hash::repeat_byte(0), session_index: 1 };
308
309 let groups = Groups::new(vec![vec![validator_a, validator_b]].into(), 2);
310
311 let mut store = StatementStore::new(&groups);
312
313 let signed_valid_by_a = {
315 let payload = valid_statement.signing_payload(&signing_context);
316 let pair = ValidatorPair::generate().0;
317 let signature = pair.sign(&payload[..]);
318
319 SignedStatement::new(
320 valid_statement.clone(),
321 validator_a,
322 signature,
323 &signing_context,
324 &pair.public(),
325 )
326 .unwrap()
327 };
328 store.insert(&groups, signed_valid_by_a, StatementOrigin::Remote).unwrap();
329
330 let signed_seconded_by_b = {
331 let payload = seconded_statement.signing_payload(&signing_context);
332 let pair = ValidatorPair::generate().0;
333 let signature = pair.sign(&payload[..]);
334
335 SignedStatement::new(
336 seconded_statement.clone(),
337 validator_b,
338 signature,
339 &signing_context,
340 &pair.public(),
341 )
342 .unwrap()
343 };
344 store.insert(&groups, signed_seconded_by_b, StatementOrigin::Remote).unwrap();
345
346 let vals = &[validator_a, validator_b];
349 let statements =
350 store.fresh_statements_for_backing(vals, candidate_hash).collect::<Vec<_>>();
351
352 assert_eq!(statements.len(), 2);
353 assert_eq!(statements[0].payload(), &seconded_statement);
354 assert_eq!(statements[1].payload(), &valid_statement);
355
356 let vals = &[validator_b, validator_a];
357 let statements =
358 store.fresh_statements_for_backing(vals, candidate_hash).collect::<Vec<_>>();
359
360 assert_eq!(statements.len(), 2);
361 assert_eq!(statements[0].payload(), &seconded_statement);
362 assert_eq!(statements[1].payload(), &valid_statement);
363 }
364}