Skip to main content

commonware_consensus/marshal/resolver/
handler.rs

1use crate::types::{Height, Round};
2use bytes::{Buf, BufMut, Bytes};
3use commonware_actor::mailbox::{self, Overflow, Policy, Sender};
4use commonware_codec::{EncodeSize, Error as CodecError, Read, ReadExt, Write};
5use commonware_cryptography::Digest;
6use commonware_resolver::{Consumer, Delivery, Fetch as ResolverFetch, p2p::Producer};
7use commonware_runtime::Metrics;
8use commonware_utils::{Span, channel::oneshot};
9use std::{
10    collections::VecDeque,
11    fmt::{Debug, Display},
12    hash::{Hash, Hasher},
13    num::NonZeroUsize,
14    sync::mpsc::TryRecvError,
15};
16use tracing::info_span;
17
18/// The subject of a backfill request.
19const BLOCK_REQUEST: u8 = 0;
20const FINALIZED_REQUEST: u8 = 1;
21const NOTARIZED_REQUEST: u8 = 2;
22
23/// Messages sent from the resolver's [Consumer]/[Producer] implementation
24/// to the marshal actor.
25pub(crate) enum Message<D: Digest> {
26    /// A request to deliver a value for a given key.
27    Deliver {
28        /// The delivery metadata attached to the resolved value.
29        delivery: Delivery<Key<D>, Annotation>,
30        /// The value being delivered.
31        value: Bytes,
32        /// A channel to send the result of the delivery.
33        response: oneshot::Sender<bool>,
34    },
35    /// A request to produce a value for a given key.
36    Produce {
37        /// The key of the value to produce.
38        key: Key<D>,
39        /// A channel to send the produced value.
40        response: oneshot::Sender<Bytes>,
41    },
42}
43
44impl<D: Digest> Message<D> {
45    /// Returns true if the requester has stopped waiting for this response.
46    pub(crate) fn response_closed(&self) -> bool {
47        match self {
48            Self::Deliver { response, .. } => response.is_closed(),
49            Self::Produce { response, .. } => response.is_closed(),
50        }
51    }
52}
53
54/// Deliveries retained while the ready queue is full.
55pub(crate) struct Pending<D: Digest>(VecDeque<Message<D>>);
56
57impl<D: Digest> Default for Pending<D> {
58    fn default() -> Self {
59        Self(VecDeque::new())
60    }
61}
62
63impl<D: Digest> Overflow<Message<D>> for Pending<D> {
64    fn is_empty(&self) -> bool {
65        self.0.is_empty()
66    }
67
68    fn drain<F>(&mut self, mut push: F)
69    where
70        F: FnMut(Message<D>) -> Option<Message<D>>,
71    {
72        while let Some(message) = self.0.pop_front() {
73            if message.response_closed() {
74                continue;
75            }
76
77            if let Some(message) = push(message) {
78                self.0.push_front(message);
79                break;
80            }
81        }
82    }
83}
84
85impl<D: Digest> Policy for Message<D> {
86    type Overflow = Pending<D>;
87
88    fn handle(overflow: &mut Self::Overflow, message: Self) {
89        // Drop produce requests so the serve backlog stays bounded by the ready
90        // queue. We prefer handling our own responses over serving peers, who can
91        // ask a less loaded peer instead.
92        if matches!(message, Self::Produce { .. }) {
93            return;
94        }
95
96        // Retain deliveries that still have a waiting requester.
97        if message.response_closed() {
98            return;
99        }
100        overflow.0.push_back(message);
101    }
102}
103
104/// A handler that forwards requests from the resolver to the marshal actor.
105///
106/// This struct implements the [Consumer] and [Producer] traits from the
107/// resolver, and acts as a bridge to the main actor loop.
108#[derive(Clone)]
109pub struct Handler<D: Digest> {
110    sender: Sender<Message<D>>,
111}
112
113impl<D: Digest> Handler<D> {
114    /// Creates a new handler.
115    pub(crate) const fn new(sender: Sender<Message<D>>) -> Self {
116        Self { sender }
117    }
118}
119
120/// Creates a resolver receiver and handler pair.
121pub fn init<D: Digest>(metrics: impl Metrics, capacity: NonZeroUsize) -> (Receiver<D>, Handler<D>) {
122    let (sender, receiver) = mailbox::new(metrics, capacity);
123    (Receiver::new(receiver), Handler::new(sender))
124}
125
126/// Receiver for resolver handler messages.
127pub struct Receiver<D: Digest> {
128    inner: mailbox::Receiver<Message<D>>,
129}
130
131impl<D: Digest> Receiver<D> {
132    pub(crate) const fn new(inner: mailbox::Receiver<Message<D>>) -> Self {
133        Self { inner }
134    }
135
136    pub(crate) async fn recv(&mut self) -> Option<Message<D>> {
137        self.inner.recv().await
138    }
139
140    pub(crate) fn try_recv(&mut self) -> Result<Message<D>, TryRecvError> {
141        self.inner.try_recv()
142    }
143}
144
145impl<D: Digest> Consumer for Handler<D> {
146    type Key = Key<D>;
147    type Value = Bytes;
148    type Subscriber = Annotation;
149    type Outcome = bool;
150
151    fn deliver(
152        &mut self,
153        delivery: Delivery<Self::Key, Self::Subscriber>,
154        value: Self::Value,
155    ) -> oneshot::Receiver<bool> {
156        let (response, receiver) = oneshot::channel();
157        let _ = self.sender.enqueue(Message::Deliver {
158            delivery,
159            value,
160            response,
161        });
162        receiver
163    }
164}
165
166impl<D: Digest> Producer for Handler<D> {
167    type Key = Key<D>;
168
169    fn produce(&mut self, key: Self::Key) -> oneshot::Receiver<Bytes> {
170        let (response, receiver) = oneshot::channel();
171        let _ = self.sender.enqueue(Message::Produce { key, response });
172        receiver
173    }
174}
175
176/// Local processing annotation for a resolved key.
177///
178/// The resolver key is the peer-visible lookup. An annotation is local
179/// metadata attached to that lookup so marshal can decide how to process the
180/// response after validating it against the key. It is not part of peer
181/// response validity. Multiple local annotations may share one peer key when
182/// they depend on the same block.
183///
184/// [`Notarization`](Annotation::Notarization) carries round-bound local
185/// context. [`Certified`](Annotation::Certified) and
186/// [`Finalized`](Annotation::Finalized) describe how block-bearing responses
187/// should be processed locally.
188///
189/// This storage role is part of the annotation because a [`Key::Block`]
190/// only names the peer-visible commitment. The same block-shaped response may
191/// need to update different local stores depending on whether it was fetched
192/// for a certified chain or for the finalized chain.
193#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
194pub enum Annotation {
195    /// A notarization requested by round.
196    Notarization { round: Round },
197    /// A block requested by commitment for a certified chain.
198    ///
199    /// The expected height is local pruning metadata and should only be
200    /// supplied when the caller has a validated height bound. It must not make
201    /// a commitment-matching response invalid. A matching block above this
202    /// bound is delivered but not cached.
203    Certified { height: Height },
204    /// A block requested by commitment for the finalized chain.
205    Finalized(Finalized),
206}
207
208/// Metadata for a finalized block requested by commitment.
209#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
210pub enum Finalized {
211    /// The finalized height is known before the request.
212    ByHeight { height: Height },
213    /// Only the finalization round is known before the request.
214    ///
215    /// This happens when a finalization names the block commitment but not the
216    /// block height.
217    ByRound { round: Round },
218}
219
220/// A raw resolver key for backfilling data.
221#[derive(Clone, Copy)]
222pub enum Key<D: Digest> {
223    /// Fetch a block by consensus commitment.
224    Block(D),
225    Finalized {
226        height: Height,
227    },
228    Notarized {
229        round: Round,
230    },
231}
232
233impl<D: Digest> Key<D> {
234    /// The subject of the request.
235    const fn subject(&self) -> u8 {
236        match self {
237            Self::Block(_) => BLOCK_REQUEST,
238            Self::Finalized { .. } => FINALIZED_REQUEST,
239            Self::Notarized { .. } => NOTARIZED_REQUEST,
240        }
241    }
242}
243
244/// A valid marshal backfill fetch request.
245#[derive(Clone, Copy, Debug, Eq, PartialEq)]
246pub(crate) enum RequestKind<D: Digest> {
247    /// Fetch a notarized proposal for a round.
248    Notarized { round: Round },
249    /// Fetch a finalization for a height.
250    Finalized { height: Height },
251    /// Fetch a certified-chain block by commitment.
252    CertifiedBlock { commitment: D, height: Height },
253    /// Fetch a finalized-chain block by commitment when its height is known.
254    FinalizedBlockByHeight { commitment: D, height: Height },
255    /// Fetch a finalized-chain block by commitment when only its finalization round is known.
256    FinalizedBlockByRound { commitment: D, round: Round },
257}
258
259/// A marshal backfill fetch with a request and local processing annotation that match.
260#[derive(Clone, Copy, Debug, Eq, PartialEq)]
261pub struct Request<D: Digest> {
262    kind: RequestKind<D>,
263}
264
265impl<D: Digest> Request<D> {
266    /// Fetch a notarized proposal for `round`.
267    pub const fn notarized(round: Round) -> Self {
268        Self {
269            kind: RequestKind::Notarized { round },
270        }
271    }
272
273    /// Fetch a finalization for `height`.
274    pub const fn finalized(height: Height) -> Self {
275        Self {
276            kind: RequestKind::Finalized { height },
277        }
278    }
279
280    /// Fetch a certified-chain block by commitment.
281    pub const fn certified_block(commitment: D, height: Height) -> Self {
282        Self {
283            kind: RequestKind::CertifiedBlock { commitment, height },
284        }
285    }
286
287    /// Fetch a finalized-chain block by commitment when its height is known.
288    pub const fn finalized_block_by_height(commitment: D, height: Height) -> Self {
289        Self {
290            kind: RequestKind::FinalizedBlockByHeight { commitment, height },
291        }
292    }
293
294    /// Fetch a finalized-chain block by commitment when only its finalization round is known.
295    pub const fn finalized_block_by_round(commitment: D, round: Round) -> Self {
296        Self {
297            kind: RequestKind::FinalizedBlockByRound { commitment, round },
298        }
299    }
300
301    pub(crate) fn above_height_floor(&self, floor: Height) -> bool {
302        match self.kind {
303            RequestKind::Finalized { height }
304            | RequestKind::CertifiedBlock { height, .. }
305            | RequestKind::FinalizedBlockByHeight { height, .. } => height > floor,
306            RequestKind::Notarized { .. } | RequestKind::FinalizedBlockByRound { .. } => true,
307        }
308    }
309
310    pub(crate) fn above_round_floor(&self, floor: Round) -> bool {
311        match self.kind {
312            RequestKind::Notarized { round } | RequestKind::FinalizedBlockByRound { round, .. } => {
313                round > floor
314            }
315            RequestKind::Finalized { .. }
316            | RequestKind::CertifiedBlock { .. }
317            | RequestKind::FinalizedBlockByHeight { .. } => true,
318        }
319    }
320
321    pub(crate) fn into_inner(self) -> ResolverFetch<Key<D>, Annotation> {
322        let (key, subscriber) = match self.kind {
323            RequestKind::Notarized { round } => {
324                (Key::Notarized { round }, Annotation::Notarization { round })
325            }
326            RequestKind::Finalized { height } => (
327                Key::Finalized { height },
328                Annotation::Finalized(Finalized::ByHeight { height }),
329            ),
330            RequestKind::CertifiedBlock { commitment, height } => {
331                (Key::Block(commitment), Annotation::Certified { height })
332            }
333            RequestKind::FinalizedBlockByHeight { commitment, height } => (
334                Key::Block(commitment),
335                Annotation::Finalized(Finalized::ByHeight { height }),
336            ),
337            RequestKind::FinalizedBlockByRound { commitment, round } => (
338                Key::Block(commitment),
339                Annotation::Finalized(Finalized::ByRound { round }),
340            ),
341        };
342        let span = info_span!("marshal.resolver.fetch", key = %key);
343        ResolverFetch {
344            key,
345            subscriber,
346            span,
347        }
348    }
349}
350
351impl<D: Digest> From<Request<D>> for ResolverFetch<Key<D>, Annotation> {
352    fn from(fetch: Request<D>) -> Self {
353        fetch.into_inner()
354    }
355}
356
357/// Returns a predicate that keeps resolver requests above the processed height floor.
358///
359/// Unrelated requests are retained. Height-bound requests are pruned once the
360/// processed height reaches them.
361pub(crate) fn above_height_floor<D: Digest>(
362    height: Height,
363) -> impl Fn(&Key<D>, &Annotation) -> bool + Send + 'static {
364    move |request, annotation| match (request, annotation) {
365        (Key::Finalized { height: requested }, _) => *requested > height,
366        (
367            Key::Block(_),
368            Annotation::Certified { height: requested }
369            | Annotation::Finalized(Finalized::ByHeight { height: requested }),
370        ) => *requested > height,
371        _ => true,
372    }
373}
374
375/// Returns a predicate that keeps resolver requests above the processed round floor.
376///
377/// Unrelated requests are retained. Round-bound requests are pruned once the
378/// processed round reaches them.
379pub(crate) fn above_round_floor<D: Digest>(
380    round: Round,
381) -> impl Fn(&Key<D>, &Annotation) -> bool + Send + 'static {
382    move |request, annotation| match (request, annotation) {
383        (Key::Notarized { round: requested }, _) => *requested > round,
384        (Key::Block(_), Annotation::Finalized(Finalized::ByRound { round: requested })) => {
385            *requested > round
386        }
387        _ => true,
388    }
389}
390
391impl<D: Digest> Write for Key<D> {
392    fn write(&self, buf: &mut impl BufMut) {
393        self.subject().write(buf);
394        match self {
395            Self::Block(commitment) => commitment.write(buf),
396            Self::Finalized { height } => height.write(buf),
397            Self::Notarized { round } => round.write(buf),
398        }
399    }
400}
401
402impl<D: Digest> Read for Key<D> {
403    type Cfg = ();
404
405    fn read_cfg(buf: &mut impl Buf, _: &()) -> Result<Self, CodecError> {
406        let request = match u8::read(buf)? {
407            BLOCK_REQUEST => Self::Block(D::read(buf)?),
408            FINALIZED_REQUEST => Self::Finalized {
409                height: Height::read(buf)?,
410            },
411            NOTARIZED_REQUEST => Self::Notarized {
412                round: Round::read(buf)?,
413            },
414            i => return Err(CodecError::InvalidEnum(i)),
415        };
416        Ok(request)
417    }
418}
419
420impl<D: Digest> EncodeSize for Key<D> {
421    fn encode_size(&self) -> usize {
422        1 + match self {
423            Self::Block(commitment) => commitment.encode_size(),
424            Self::Finalized { height } => height.encode_size(),
425            Self::Notarized { round } => round.encode_size(),
426        }
427    }
428}
429
430impl<D: Digest> Span for Key<D> {}
431
432impl<D: Digest> PartialEq for Key<D> {
433    fn eq(&self, other: &Self) -> bool {
434        match (&self, &other) {
435            (Self::Block(a), Self::Block(b)) => a == b,
436            (Self::Finalized { height: a }, Self::Finalized { height: b }) => a == b,
437            (Self::Notarized { round: a }, Self::Notarized { round: b }) => a == b,
438            _ => false,
439        }
440    }
441}
442
443impl<D: Digest> Eq for Key<D> {}
444
445impl<D: Digest> Ord for Key<D> {
446    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
447        match (&self, &other) {
448            (Self::Block(a), Self::Block(b)) => a.cmp(b),
449            (Self::Finalized { height: a }, Self::Finalized { height: b }) => a.cmp(b),
450            (Self::Notarized { round: a }, Self::Notarized { round: b }) => a.cmp(b),
451            (a, b) => a.subject().cmp(&b.subject()),
452        }
453    }
454}
455
456impl<D: Digest> PartialOrd for Key<D> {
457    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
458        Some(self.cmp(other))
459    }
460}
461
462impl<D: Digest> Hash for Key<D> {
463    fn hash<H: Hasher>(&self, state: &mut H) {
464        self.subject().hash(state);
465        match self {
466            Self::Block(commitment) => commitment.hash(state),
467            Self::Finalized { height } => height.hash(state),
468            Self::Notarized { round } => round.hash(state),
469        }
470    }
471}
472
473impl<D: Digest> Display for Key<D> {
474    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
475        match self {
476            Self::Block(commitment) => write!(f, "Block({commitment:?})"),
477            Self::Finalized { height } => write!(f, "Finalized({height:?})"),
478            Self::Notarized { round } => write!(f, "Notarized({round:?})"),
479        }
480    }
481}
482
483impl<D: Digest> Debug for Key<D> {
484    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
485        match self {
486            Self::Block(commitment) => write!(f, "Block({commitment:?})"),
487            Self::Finalized { height } => write!(f, "Finalized({height:?})"),
488            Self::Notarized { round } => write!(f, "Notarized({round:?})"),
489        }
490    }
491}
492
493#[cfg(feature = "arbitrary")]
494impl<D: Digest> arbitrary::Arbitrary<'_> for Key<D>
495where
496    D: for<'a> arbitrary::Arbitrary<'a>,
497{
498    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
499        let choice = u.int_in_range(0..=2)?;
500        match choice {
501            0 => Ok(Self::Block(u.arbitrary()?)),
502            1 => Ok(Self::Finalized {
503                height: u.arbitrary()?,
504            }),
505            2 => Ok(Self::Notarized {
506                round: u.arbitrary()?,
507            }),
508            _ => unreachable!(),
509        }
510    }
511}
512
513#[cfg(test)]
514mod tests {
515    use super::*;
516    use crate::types::{Epoch, View};
517    use commonware_codec::{Encode, ReadExt};
518    use commonware_cryptography::{
519        Hasher as _,
520        sha256::{Digest as Sha256Digest, Sha256},
521    };
522    use commonware_utils::vec::NonEmptyVec;
523    use std::collections::BTreeSet;
524
525    type D = Sha256Digest;
526
527    #[test]
528    fn handle_retains_open_deliveries_only() {
529        let mut overflow = Pending::<D>::default();
530        let deliver = |height: u64, response| Message::Deliver {
531            delivery: Delivery {
532                key: Key::Finalized {
533                    height: Height::new(height),
534                },
535                subscribers: NonEmptyVec::new((
536                    Annotation::Finalized(Finalized::ByHeight {
537                        height: Height::new(height),
538                    }),
539                    tracing::Span::none(),
540                )),
541            },
542            value: Bytes::new(),
543            response,
544        };
545
546        // An overflowed produce request is dropped and its requester sees the
547        // closed response.
548        let (response, mut produce) = oneshot::channel();
549        Message::handle(
550            &mut overflow,
551            Message::Produce {
552                key: Key::Finalized {
553                    height: Height::new(1),
554                },
555                response,
556            },
557        );
558        assert!(matches!(
559            produce.try_recv(),
560            Err(oneshot::error::TryRecvError::Closed)
561        ));
562
563        // Deliveries are retained, and drain skips one whose requester left.
564        let (response, closed) = oneshot::channel();
565        Message::handle(&mut overflow, deliver(2, response));
566        let (response, _open) = oneshot::channel();
567        Message::handle(&mut overflow, deliver(3, response));
568        drop(closed);
569
570        let mut messages = Vec::new();
571        Overflow::drain(&mut overflow, |message| {
572            messages.push(message);
573            None
574        });
575        assert_eq!(messages.len(), 1);
576        assert!(matches!(
577            messages.pop(),
578            Some(Message::Deliver {
579                delivery: Delivery {
580                    key: Key::Finalized { height },
581                    ..
582                },
583                ..
584            }) if height == Height::new(3)
585        ));
586    }
587
588    #[test]
589    fn test_cross_variant_hash_differs() {
590        use std::{
591            collections::hash_map::DefaultHasher,
592            hash::{Hash, Hasher},
593        };
594
595        fn hash_of<T: Hash>(t: &T) -> u64 {
596            let mut h = DefaultHasher::new();
597            t.hash(&mut h);
598            h.finish()
599        }
600
601        let finalized = Key::<D>::Finalized {
602            height: Height::new(1),
603        };
604        let notarized = Key::<D>::Notarized {
605            round: Round::new(Epoch::new(0), View::new(1)),
606        };
607        assert_ne!(hash_of(&finalized), hash_of(&notarized));
608    }
609
610    #[test]
611    fn test_subject_block_encoding() {
612        let commitment = Sha256::hash(&[b"test"]);
613        let request = Key::<D>::Block(commitment);
614
615        // Test encoding
616        let encoded = request.encode();
617        assert_eq!(encoded.len(), 33); // 1 byte for enum variant + 32 bytes for commitment
618        assert_eq!(encoded[0], 0); // Block variant
619
620        // Test decoding
621        let mut buf = encoded.as_ref();
622        let decoded = Key::<D>::read(&mut buf).unwrap();
623        assert_eq!(request, decoded);
624        assert_eq!(decoded, Key::Block(commitment));
625    }
626
627    #[test]
628    fn test_subject_finalized_encoding() {
629        let height = Height::new(12345u64);
630        let request = Key::<D>::Finalized { height };
631
632        // Test encoding
633        let encoded = request.encode();
634        assert_eq!(encoded[0], 1); // Finalized variant
635
636        // Test decoding
637        let mut buf = encoded.as_ref();
638        let decoded = Key::<D>::read(&mut buf).unwrap();
639        assert_eq!(request, decoded);
640        assert_eq!(decoded, Key::Finalized { height });
641    }
642
643    #[test]
644    fn test_subject_notarized_encoding() {
645        let round = Round::new(Epoch::new(67890), View::new(12345));
646        let request = Key::<D>::Notarized { round };
647
648        // Test encoding
649        let encoded = request.encode();
650        assert_eq!(encoded[0], 2); // Notarized variant
651
652        // Test decoding
653        let mut buf = encoded.as_ref();
654        let decoded = Key::<D>::read(&mut buf).unwrap();
655        assert_eq!(request, decoded);
656        assert_eq!(decoded, Key::Notarized { round });
657    }
658
659    #[test]
660    fn test_subject_decode_rejects_invalid_enum_tag() {
661        let bad = [3u8];
662        let mut buf = bad.as_ref();
663        assert!(matches!(
664            Key::<D>::read(&mut buf),
665            Err(CodecError::InvalidEnum(3))
666        ));
667    }
668
669    #[test]
670    fn test_subject_hash() {
671        use std::collections::HashSet;
672
673        let r1 = Key::<D>::Finalized {
674            height: Height::new(100),
675        };
676        let r2 = Key::<D>::Finalized {
677            height: Height::new(100),
678        };
679        let r3 = Key::<D>::Finalized {
680            height: Height::new(200),
681        };
682
683        let mut set = HashSet::new();
684        set.insert(r1);
685        assert!(!set.insert(r2)); // Should not insert duplicate
686        assert!(set.insert(r3)); // Should insert different value
687    }
688
689    #[test]
690    fn test_height_floor_predicate() {
691        let floor = Height::new(100);
692        let higher_finalized = Key::<D>::Finalized {
693            height: Height::new(200),
694        };
695        let notarized = Key::<D>::Notarized {
696            round: Round::new(Epoch::new(333), View::new(150)),
697        };
698        let block = Key::<D>::Block(Sha256::hash(&[b"block"]));
699        let stale_finalized = Annotation::Finalized(Finalized::ByHeight {
700            height: Height::new(100),
701        });
702        let fresh_certified = Annotation::Certified {
703            height: Height::new(101),
704        };
705        let stale_certified = Annotation::Certified {
706            height: Height::new(100),
707        };
708
709        let predicate = above_height_floor(floor);
710        assert!(predicate(
711            &higher_finalized,
712            &Annotation::Finalized(Finalized::ByHeight {
713                height: Height::new(200),
714            })
715        ));
716        assert!(predicate(
717            &notarized,
718            &Annotation::Notarization {
719                round: Round::new(Epoch::new(333), View::new(150)),
720            }
721        ));
722        assert!(predicate(&block, &fresh_certified));
723
724        let same_height = Key::<D>::Finalized {
725            height: Height::new(100),
726        };
727        assert!(!predicate(
728            &same_height,
729            &Annotation::Finalized(Finalized::ByHeight {
730                height: Height::new(100),
731            })
732        ));
733        assert!(!predicate(&block, &stale_finalized));
734        assert!(!predicate(&block, &stale_certified));
735    }
736
737    #[test]
738    fn test_round_floor_predicate() {
739        let floor = Round::new(Epoch::new(1), View::new(10));
740        let block = Key::<D>::Block(Sha256::hash(&[b"block"]));
741        let higher_notarized = Key::<D>::Notarized {
742            round: Round::new(Epoch::new(1), View::new(11)),
743        };
744        let same_notarized = Key::<D>::Notarized {
745            round: Round::new(Epoch::new(1), View::new(10)),
746        };
747        let finalized = Key::<D>::Finalized {
748            height: Height::new(100),
749        };
750
751        let predicate = above_round_floor(floor);
752        assert!(predicate(
753            &higher_notarized,
754            &Annotation::Notarization {
755                round: Round::new(Epoch::new(1), View::new(11)),
756            }
757        ));
758        assert!(predicate(
759            &finalized,
760            &Annotation::Finalized(Finalized::ByHeight {
761                height: Height::new(100),
762            })
763        ));
764        assert!(predicate(
765            &block,
766            &Annotation::Finalized(Finalized::ByRound {
767                round: Round::new(Epoch::new(1), View::new(11)),
768            })
769        ));
770        assert!(!predicate(
771            &same_notarized,
772            &Annotation::Notarization {
773                round: Round::new(Epoch::new(1), View::new(10)),
774            }
775        ));
776        assert!(!predicate(
777            &block,
778            &Annotation::Finalized(Finalized::ByRound {
779                round: Round::new(Epoch::new(1), View::new(10)),
780            })
781        ));
782    }
783
784    #[test]
785    fn test_encode_size() {
786        let commitment = Sha256::hash(&[&[0u8; 32]]);
787        let r1 = Key::<D>::Block(commitment);
788        let r2 = Key::<D>::Finalized {
789            height: Height::new(u64::MAX),
790        };
791        let r3 = Key::<D>::Notarized {
792            round: Round::new(Epoch::new(333), View::new(0)),
793        };
794
795        // Verify encode_size matches actual encoded length
796        assert_eq!(r1.encode_size(), r1.encode().len());
797        assert_eq!(r2.encode_size(), r2.encode().len());
798        assert_eq!(r3.encode_size(), r3.encode().len());
799    }
800
801    #[test]
802    fn test_request_ord_same_variant() {
803        // Test ordering within the same variant
804        let commitment1 = Sha256::hash(&[b"test1"]);
805        let commitment2 = Sha256::hash(&[b"test2"]);
806        let block1 = Key::<D>::Block(commitment1);
807        let block2 = Key::<D>::Block(commitment2);
808
809        // Block ordering depends on commitment ordering
810        if commitment1 < commitment2 {
811            assert!(block1 < block2);
812            assert!(block2 > block1);
813        } else {
814            assert!(block1 > block2);
815            assert!(block2 < block1);
816        }
817
818        // Finalized ordering by height
819        let fin1 = Key::<D>::Finalized {
820            height: Height::new(100),
821        };
822        let fin2 = Key::<D>::Finalized {
823            height: Height::new(200),
824        };
825        let fin3 = Key::<D>::Finalized {
826            height: Height::new(200),
827        };
828
829        assert!(fin1 < fin2);
830        assert!(fin2 > fin1);
831        assert_eq!(fin2.cmp(&fin3), std::cmp::Ordering::Equal);
832
833        // Notarized ordering by view
834        let not1 = Key::<D>::Notarized {
835            round: Round::new(Epoch::new(333), View::new(50)),
836        };
837        let not2 = Key::<D>::Notarized {
838            round: Round::new(Epoch::new(333), View::new(150)),
839        };
840        let not3 = Key::<D>::Notarized {
841            round: Round::new(Epoch::new(333), View::new(150)),
842        };
843
844        assert!(not1 < not2);
845        assert!(not2 > not1);
846        assert_eq!(not2.cmp(&not3), std::cmp::Ordering::Equal);
847    }
848
849    #[test]
850    fn test_request_ord_cross_variant() {
851        let commitment = Sha256::hash(&[b"test"]);
852        let block = Key::<D>::Block(commitment);
853        let finalized = Key::<D>::Finalized {
854            height: Height::new(100),
855        };
856        let notarized = Key::<D>::Notarized {
857            round: Round::new(Epoch::new(333), View::new(200)),
858        };
859
860        // Block < Finalized < Notarized
861        assert!(block < finalized);
862        assert!(block < notarized);
863        assert!(finalized < notarized);
864
865        assert!(finalized > block);
866        assert!(notarized > block);
867        assert!(notarized > finalized);
868
869        // Test all combinations
870        assert_eq!(block.cmp(&finalized), std::cmp::Ordering::Less);
871        assert_eq!(block.cmp(&notarized), std::cmp::Ordering::Less);
872        assert_eq!(finalized.cmp(&notarized), std::cmp::Ordering::Less);
873        assert_eq!(finalized.cmp(&block), std::cmp::Ordering::Greater);
874        assert_eq!(notarized.cmp(&block), std::cmp::Ordering::Greater);
875        assert_eq!(notarized.cmp(&finalized), std::cmp::Ordering::Greater);
876    }
877
878    #[test]
879    fn test_request_partial_ord() {
880        let commitment1 = Sha256::hash(&[b"test1"]);
881        let commitment2 = Sha256::hash(&[b"test2"]);
882        let block1 = Key::<D>::Block(commitment1);
883        let block2 = Key::<D>::Block(commitment2);
884        let finalized = Key::<D>::Finalized {
885            height: Height::new(100),
886        };
887        let notarized = Key::<D>::Notarized {
888            round: Round::new(Epoch::new(333), View::new(200)),
889        };
890
891        // PartialOrd should always return Some
892        assert!(block1.partial_cmp(&block2).is_some());
893        assert!(block1.partial_cmp(&finalized).is_some());
894        assert!(finalized.partial_cmp(&notarized).is_some());
895
896        // Verify consistency with Ord
897        assert_eq!(
898            block1.partial_cmp(&finalized),
899            Some(std::cmp::Ordering::Less)
900        );
901        assert_eq!(
902            finalized.partial_cmp(&notarized),
903            Some(std::cmp::Ordering::Less)
904        );
905        assert_eq!(
906            notarized.partial_cmp(&block1),
907            Some(std::cmp::Ordering::Greater)
908        );
909    }
910
911    #[test]
912    fn test_request_ord_sorting() {
913        let commitment1 = Sha256::hash(&[b"a"]);
914        let commitment2 = Sha256::hash(&[b"b"]);
915        let commitment3 = Sha256::hash(&[b"c"]);
916
917        let requests = vec![
918            Key::<D>::Notarized {
919                round: Round::new(Epoch::new(333), View::new(300)),
920            },
921            Key::<D>::Block(commitment2),
922            Key::<D>::Finalized {
923                height: Height::new(200),
924            },
925            Key::<D>::Block(commitment1),
926            Key::<D>::Notarized {
927                round: Round::new(Epoch::new(333), View::new(250)),
928            },
929            Key::<D>::Finalized {
930                height: Height::new(100),
931            },
932            Key::<D>::Block(commitment3),
933        ];
934
935        // Sort using BTreeSet (uses Ord)
936        let sorted: Vec<_> = requests
937            .into_iter()
938            .collect::<BTreeSet<_>>()
939            .into_iter()
940            .collect();
941
942        // Verify order: all Blocks first (sorted by commitment), then Finalized (by height), then Notarized (by view)
943        assert_eq!(sorted.len(), 7);
944
945        // Check that all blocks come first
946        assert!(matches!(sorted[0], Key::<D>::Block(_)));
947        assert!(matches!(sorted[1], Key::<D>::Block(_)));
948        assert!(matches!(sorted[2], Key::<D>::Block(_)));
949
950        // Check that finalized come next
951        assert_eq!(
952            sorted[3],
953            Key::<D>::Finalized {
954                height: Height::new(100)
955            }
956        );
957        assert_eq!(
958            sorted[4],
959            Key::<D>::Finalized {
960                height: Height::new(200)
961            }
962        );
963
964        // Check that notarized come last
965        assert_eq!(
966            sorted[5],
967            Key::<D>::Notarized {
968                round: Round::new(Epoch::new(333), View::new(250))
969            }
970        );
971        assert_eq!(
972            sorted[6],
973            Key::<D>::Notarized {
974                round: Round::new(Epoch::new(333), View::new(300))
975            }
976        );
977    }
978
979    #[test]
980    fn test_request_ord_edge_cases() {
981        // Test with extreme values
982        let min_finalized = Key::<D>::Finalized {
983            height: Height::new(0),
984        };
985        let max_finalized = Key::<D>::Finalized {
986            height: Height::new(u64::MAX),
987        };
988        let min_notarized = Key::<D>::Notarized {
989            round: Round::new(Epoch::new(333), View::new(0)),
990        };
991        let max_notarized = Key::<D>::Notarized {
992            round: Round::new(Epoch::new(333), View::new(u64::MAX)),
993        };
994
995        assert!(min_finalized < max_finalized);
996        assert!(min_notarized < max_notarized);
997        assert!(max_finalized < min_notarized);
998
999        // Test self-comparison
1000        let commitment = Sha256::hash(&[b"self"]);
1001        let block = Key::<D>::Block(commitment);
1002        assert_eq!(block.cmp(&block), std::cmp::Ordering::Equal);
1003        assert_eq!(min_finalized.cmp(&min_finalized), std::cmp::Ordering::Equal);
1004        assert_eq!(max_notarized.cmp(&max_notarized), std::cmp::Ordering::Equal);
1005    }
1006
1007    #[cfg(feature = "arbitrary")]
1008    mod conformance {
1009        use super::*;
1010        use commonware_codec::conformance::CodecConformance;
1011
1012        commonware_conformance::conformance_tests! {
1013            CodecConformance<Key<D>>
1014        }
1015    }
1016}