1use core::{hash::Hash, fmt::Debug, future::Future, num::NonZero};
2
3use crate::{BlockNumber, RoundNumber};
4
5pub trait Validator: Clone + Copy + PartialEq + Eq + Hash + Debug {}
14impl<V: Clone + Copy + PartialEq + Eq + Hash + Debug> Validator for V {}
15
16pub trait Signature: Clone + AsRef<[u8]> {}
28impl<S: Clone + AsRef<[u8]>> Signature for S {}
29
30pub trait AggregateSignature: Clone + AsRef<[u8]> {}
39impl<S: Clone + AsRef<[u8]>> AggregateSignature for S {}
40
41#[derive(Debug)]
43pub struct InvalidAggregateSignature;
44
45pub trait Signer {
47 type Validator: Validator;
49 type Signature: Signature;
51
52 fn validator(&self) -> Self::Validator;
56
57 type SignFuture: Future<Output = Self::Signature>;
61
62 fn sign(&self, message: impl IntoIterator<Item = impl AsRef<[u8]>>) -> Self::SignFuture;
66}
67
68impl<S: ?Sized + Signer> Signer for &S {
69 type Validator = S::Validator;
70 type Signature = S::Signature;
71
72 fn validator(&self) -> Self::Validator {
73 S::validator(self)
74 }
75
76 type SignFuture = S::SignFuture;
77
78 fn sign(&self, message: impl IntoIterator<Item = impl AsRef<[u8]>>) -> Self::SignFuture {
79 S::sign(self, message)
80 }
81}
82
83pub trait SignatureScheme {
91 type Validator: Validator;
93 type Signature: Signature;
95 type AggregateSignature: AggregateSignature;
103
104 #[must_use]
108 fn verify(
109 &self,
110 validator: &Self::Validator,
111 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
112 signature: &Self::Signature,
113 ) -> bool;
114
115 #[must_use]
123 fn aggregate<'sig>(
124 &self,
125 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
126 signatures: impl IntoIterator<Item = (&'sig Self::Validator, &'sig Self::Signature)>,
127 ) -> Self::AggregateSignature
128 where
129 Self::Validator: 'sig,
130 Self::Signature: 'sig;
131
132 fn verify_aggregate(
143 &self,
144 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
145 aggregate_signature: &Self::AggregateSignature,
146 ) -> Result<impl IntoIterator<Item = Self::Validator>, InvalidAggregateSignature>;
147}
148
149impl<S: ?Sized + SignatureScheme> SignatureScheme for &S {
150 type Validator = S::Validator;
151 type Signature = S::Signature;
152 type AggregateSignature = S::AggregateSignature;
153
154 fn verify(
155 &self,
156 validator: &Self::Validator,
157 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
158 signature: &Self::Signature,
159 ) -> bool {
160 S::verify(self, validator, message, signature)
161 }
162
163 fn aggregate<'sig>(
164 &self,
165 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
166 signatures: impl IntoIterator<Item = (&'sig Self::Validator, &'sig Self::Signature)>,
167 ) -> Self::AggregateSignature
168 where
169 Self::Validator: 'sig,
170 Self::Signature: 'sig,
171 {
172 S::aggregate(self, message, signatures)
173 }
174
175 fn verify_aggregate(
176 &self,
177 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
178 aggregate_signature: &Self::AggregateSignature,
179 ) -> Result<impl IntoIterator<Item = Self::Validator>, InvalidAggregateSignature> {
180 S::verify_aggregate(self, message, aggregate_signature)
181 }
182}
183
184pub trait ValidatorSet {
186 type Validator: Validator;
188
189 fn total_weight(&self) -> NonZero<u16>;
193
194 fn validators(&self) -> impl IntoIterator<Item = &Self::Validator>;
202
203 fn weight(&self, validator: &Self::Validator) -> Option<NonZero<u16>>;
209
210 fn threshold(&self) -> u16 {
215 u16::try_from(((u32::from(u16::from(self.total_weight())) * 2) / 3) + 1)
216 .expect("threshold is less than or equal to the total weight, which is a `u16`")
217 }
218
219 fn fault_threshold(&self) -> u16 {
224 u16::from(self.total_weight()) - self.threshold()
225 }
226
227 fn proposer(&self, block_number: BlockNumber, round_number: RoundNumber) -> Self::Validator;
232}
233
234impl<V: ?Sized + ValidatorSet> ValidatorSet for &V {
235 type Validator = V::Validator;
236
237 fn total_weight(&self) -> NonZero<u16> {
238 V::total_weight(self)
239 }
240
241 fn validators(&self) -> impl IntoIterator<Item = &Self::Validator> {
242 V::validators(self)
243 }
244
245 fn weight(&self, validator: &Self::Validator) -> Option<NonZero<u16>> {
246 V::weight(self, validator)
247 }
248
249 fn proposer(&self, block_number: BlockNumber, round_number: RoundNumber) -> Self::Validator {
250 V::proposer(self, block_number, round_number)
251 }
252}
253
254macro_rules! map {
255 ($feature: literal, $map: path, $($generics: tt)*) => {
256 #[cfg(feature = $feature)]
257 impl<$($generics)*> ValidatorSet for $map {
258 type Validator = V;
259
260 fn total_weight(&self) -> NonZero<u16> {
263 NonZero::new(
264 self
265 .values()
266 .try_fold(0u16, |accum, value| accum.checked_add(u16::from(*value)))
267 .unwrap(),
268 )
269 .unwrap()
270 }
271
272 fn validators(&self) -> impl IntoIterator<Item = &Self::Validator> {
273 self.keys()
274 }
275
276 fn weight(&self, validator: &Self::Validator) -> Option<NonZero<u16>> {
277 self.get(validator).copied()
278 }
279
280 fn proposer(&self, block_number: BlockNumber, round_number: RoundNumber) -> Self::Validator {
305 const {
306 assert!(usize::BITS > 16);
308 }
309
310 use alloc::vec::Vec;
311
312 let total_weight = u16::from(self.total_weight());
313
314 let i = usize::from({
315 let block_i = (u64::from(block_number) - 1) % u64::from(total_weight);
317 let round_i = (u64::from(round_number) - 1) % u64::from(total_weight);
319 let i = (block_i + round_i) % u64::from(total_weight);
320 u16::try_from(i).expect(
321 "`i` indexes a validator by weight, where the weight is representable in a `u16`",
322 )
323 });
324
325 let mut all_keys = Vec::with_capacity((i + 2).min(usize::from(total_weight)));
326 for (key, value) in self {
327 let mut insert_at = 0;
328 for weight_i in 0 .. u16::from(*value) {
329 match all_keys[insert_at ..].binary_search_by(
342 |(existing_weight, existing_key): &(u16, &V)| {
343 existing_weight.cmp(&weight_i).then((*existing_key).cmp(key))
344 },
345 ) {
346 Ok(_index) => {
349 panic!("`binary_search_by` located element despite not being already present")
350 }
351 Err(insert_at_in_slice) => insert_at += insert_at_in_slice,
352 }
353
354 if insert_at > i {
360 break;
361 }
362
363 all_keys.insert(insert_at, (weight_i, key));
364 insert_at += 1;
365
366 all_keys.truncate(i + 1);
369 }
370 }
371
372 *all_keys[i].1
373 }
374 }
375 };
376}
377
378#[cfg(not(target_pointer_width = "16"))]
379map!(
380 "alloc",
381 alloc::collections::BTreeMap<V, NonZero<u16>>,
382 V: PartialOrd + Ord + Validator
383);
384
385#[cfg(not(target_pointer_width = "16"))]
386map!(
387 "std",
388 std::collections::HashMap<V, NonZero<u16>, H>,
389 V: PartialOrd + Ord + Validator, H: core::hash::BuildHasher
390);
391
392#[cfg(test)]
393mod tests {
394 use super::*;
395
396 pub(crate) struct TestSignatureScheme(u64, u64);
397 impl TestSignatureScheme {
398 pub(crate) fn new() -> Self {
399 use rand_core::{TryRngCore as _, OsRng};
400 Self(OsRng.try_next_u64().unwrap(), OsRng.try_next_u64().unwrap())
401 }
402 }
403
404 impl TestSignatureScheme {
405 fn hash(&self, message: impl IntoIterator<Item = impl AsRef<[u8]>>) -> [u8; 8] {
406 #[expect(deprecated)]
413 use core::hash::{Hasher as _, SipHasher};
414
415 #[expect(deprecated)]
416 let mut hasher = SipHasher::new_with_keys(self.0, self.1);
417 for chunk in message {
418 for byte in chunk.as_ref() {
419 hasher.write_u8(*byte);
421 }
422 }
423 hasher.finish().to_le_bytes()
424 }
425 }
426
427 impl SignatureScheme for TestSignatureScheme {
428 type Validator = u8;
429 type Signature = [u8; 1 + 8];
431 type AggregateSignature = [u8; 32 + 8];
433
434 fn verify(
435 &self,
436 validator: &<Self as SignatureScheme>::Validator,
437 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
438 signature: &<Self as SignatureScheme>::Signature,
439 ) -> bool {
440 (validator == &signature[0]) && (self.hash(message) == signature[1 ..])
441 }
442
443 fn aggregate<'sig>(
444 &self,
445 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
446 signatures: impl IntoIterator<
447 Item = (
448 &'sig <Self as SignatureScheme>::Validator,
449 &'sig <Self as SignatureScheme>::Signature,
450 ),
451 >,
452 ) -> <Self as SignatureScheme>::AggregateSignature
453 where
454 <Self as SignatureScheme>::Validator: 'sig,
455 <Self as SignatureScheme>::Signature: 'sig,
456 {
457 let hash = self.hash(message);
458 let mut aggregate_signature = [0; 32 + 8];
459 for (validator, signature) in signatures {
460 assert!((validator == &signature[0]) && (hash == signature[1 ..]));
461 aggregate_signature[usize::from(validator / 8)] |= 1 << (validator % 8);
463 }
464 aggregate_signature[32 ..].copy_from_slice(&hash);
465 aggregate_signature
466 }
467
468 fn verify_aggregate(
469 &self,
470 message: impl IntoIterator<Item = impl AsRef<[u8]>>,
471 aggregate_signature: &<Self as SignatureScheme>::AggregateSignature,
472 ) -> Result<
473 impl IntoIterator<Item = <Self as SignatureScheme>::Validator>,
474 InvalidAggregateSignature,
475 > {
476 let hash = self.hash(message);
477 if aggregate_signature[32 ..] != hash {
478 Err(InvalidAggregateSignature)?;
479 }
480
481 Ok(
483 aggregate_signature[.. 32]
484 .iter()
485 .enumerate()
486 .flat_map(|(i, b)| {
487 let i = 8 * u8::try_from(i).unwrap();
488 core::array::from_fn::<Option<u8>, 8, _>(|j| {
489 ((b & (1 << (j % 8))) != 0)
490 .then_some(i.checked_add(u8::try_from(j).unwrap()).unwrap())
491 })
492 })
493 .flatten(),
494 )
495 }
496 }
497
498 pub(crate) struct TestSigner {
499 signature_scheme: TestSignatureScheme,
500 validator: u8,
501 }
502
503 impl TestSignatureScheme {
504 pub(crate) fn signer(&self, validator: u8) -> TestSigner {
505 TestSigner { signature_scheme: TestSignatureScheme(self.0, self.1), validator }
506 }
507 }
508
509 impl Signer for TestSigner {
510 type Validator = u8;
511 type Signature = [u8; 1 + 8];
512 fn validator(&self) -> <Self as Signer>::Validator {
513 self.validator
514 }
515 type SignFuture = core::future::Ready<Self::Signature>;
516 fn sign(&self, message: impl IntoIterator<Item = impl AsRef<[u8]>>) -> Self::SignFuture {
517 let mut result = [0; 1 + 8];
518 result[0] = self.validator;
519 result[1 ..].copy_from_slice(&self.signature_scheme.hash(message));
520 core::future::ready(result)
521 }
522 }
523
524 #[test]
525 fn signature_scheme() {
526 use core::{
527 pin::pin,
528 task::{Poll, Waker, Context},
529 };
530
531 let mut context = Context::from_waker(Waker::noop());
532
533 let message = [[12].as_slice(), &[2], &[3]];
534 let other_message = [[].as_slice()];
535 let signature_scheme = loop {
536 let signature_scheme = TestSignatureScheme::new();
537 if signature_scheme.hash(message) == signature_scheme.hash(other_message) {
538 continue;
539 }
540 break signature_scheme;
541 };
542
543 let aggregated = [3, 101, 135];
544 let mut signatures = [None; 3];
545 for validator in 0 ..= u8::MAX {
546 let Poll::Ready(signature) =
547 pin!(signature_scheme.signer(validator).sign(message)).poll(&mut context)
548 else {
549 panic!("`TestSignatureScheme::sign` returned `Poll::Pending`")
550 };
551 assert!(signature_scheme.verify(&validator, message, &signature));
552 assert!(!signature_scheme.verify(&validator.wrapping_add(1), message, &signature));
553 assert!(!signature_scheme.verify(&validator, other_message, &signature));
554 if let Some(i) =
555 aggregated.iter().position(|validator_in_list| *validator_in_list == validator)
556 {
557 signatures[i] = Some((validator, signature));
558 }
559 }
560
561 let aggregate_signature = signature_scheme.aggregate(
562 message,
563 signatures.iter().map(|signature| {
564 let (validator, signature) = signature.as_ref().unwrap();
565 (validator, signature)
566 }),
567 );
568
569 let mut yielded_aggregated =
570 signature_scheme.verify_aggregate(message, &aggregate_signature).unwrap().into_iter();
571 let mut aggregated = aggregated.into_iter();
572 for (i, j) in (&mut yielded_aggregated).zip(&mut aggregated) {
573 assert_eq!(i, j);
574 }
575 assert!(yielded_aggregated.next().is_none());
576 assert!(aggregated.next().is_none());
577 }
578}
579#[cfg(test)]
580pub(crate) use tests::*;
581
582#[cfg(all(test, any(feature = "alloc", feature = "std")))]
583fn test_map<M: FromIterator<(u16, NonZero<u16>)> + ValidatorSet<Validator = u16>>() {
584 let test = |pairs: &[(u16, NonZero<u16>)]| {
585 let map = pairs.iter().copied().collect::<M>();
586
587 assert_eq!(
588 u16::from(map.total_weight()),
589 pairs.iter().map(|(_validator, weight)| u16::from(*weight)).sum::<u16>()
590 );
591
592 assert_eq!(
593 map.validators().into_iter().copied().collect::<alloc::collections::BTreeSet<_>>(),
594 pairs
595 .iter()
596 .map(|(validator, _weight)| validator)
597 .copied()
598 .collect::<alloc::collections::BTreeSet<_>>()
599 );
600
601 let proposers = {
608 let mut proposers = alloc::vec::Vec::new();
609 for (validator, weight) in pairs.iter().copied() {
610 let weight = u16::from(weight);
611 for w in 0 .. weight {
612 proposers.push((validator, w));
613 }
614 }
615 proposers.sort_by(|(a_validator, a_weight), (b_validator, b_weight)| {
616 a_weight.cmp(b_weight).then_with(|| a_validator.cmp(b_validator))
617 });
618 proposers.into_iter().map(|(validator, _weight)| validator).collect::<alloc::vec::Vec<_>>()
619 };
620
621 assert_eq!(map.proposer(BlockNumber::ONE, RoundNumber::ONE), proposers[0]);
623
624 for i in 0 .. (2 * proposers.len()) {
630 assert_eq!(
638 map.proposer(
639 BlockNumber::from(NonZero::new(1 + u64::try_from(i).unwrap()).unwrap()),
640 RoundNumber::ONE
641 ),
642 proposers[i % proposers.len()]
643 );
644 assert_eq!(
645 map.proposer(
646 BlockNumber::ONE,
647 RoundNumber(NonZero::new(1 + u64::try_from(i).unwrap()).unwrap())
648 ),
649 proposers[i % proposers.len()]
650 );
651 }
652
653 assert_eq!(
655 map.proposer(
656 BlockNumber::from(NonZero::new(u64::MAX).unwrap()),
657 RoundNumber(NonZero::new(u64::MAX).unwrap())
658 ),
659 proposers[usize::try_from(
660 (2 * u128::from(u64::MAX - 1)) % u128::try_from(proposers.len()).unwrap()
662 )
663 .unwrap()]
664 );
665 };
666
667 for weighted in [false, true] {
669 for len in 1 .. 128 {
671 test(
672 &(1 ..= len)
673 .map(|i| {
674 let weight = if weighted {
675 use rand_core::{TryRngCore as _, OsRng};
676 NonZero::new(u16::try_from(1 + (OsRng.try_next_u64().unwrap() % 16)).unwrap())
677 .unwrap()
678 } else {
679 NonZero::new(1).unwrap()
680 };
681 (i, weight)
682 })
683 .collect::<alloc::vec::Vec<_>>(),
684 );
685 }
686 }
687}
688#[cfg(feature = "alloc")]
689#[test]
690fn test_btree_map() {
691 test_map::<alloc::collections::BTreeMap<u16, NonZero<u16>>>();
692}
693#[cfg(feature = "std")]
694#[test]
695fn test_hash_map() {
696 test_map::<std::collections::HashMap<u16, NonZero<u16>>>();
697}