shred: Consolidate and document ShredFlags logic (#4653)

Also add a unit test to guarantee reference tick saturation behavior
anza-xyz · Jan 28, 2025 · 9ca57b1 · 9ca57b1
1 parent 4fe5092
commit 9ca57b1
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 11 deletions.
diff --git a/ledger/src/shred.rs b/ledger/src/shred.rs
@@ -123,6 +123,17 @@ bitflags! {
     }
 }
 
+impl ShredFlags {
+    /// Creates a new ShredFlags from the given reference_tick
+    ///
+    /// SHRED_TICK_REFERENCE_MASK is comprised of only six bits whereas the
+    /// reference_tick has 8 bits (u8). The reference_tick bits will saturate
+    /// in the event that reference_tick > SHRED_TICK_REFERENCE_MASK
+    pub(crate) fn from_reference_tick(reference_tick: u8) -> Self {
+        Self::from_bits_retain(Self::SHRED_TICK_REFERENCE_MASK.bits().min(reference_tick))
+    }
+}
+
 #[derive(Debug, Error)]
 pub enum Error {
     #[error(transparent)]
@@ -1480,6 +1491,15 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_shred_flags_reference_tick_saturates() {
+        const MAX_REFERENCE_TICK: u8 = ShredFlags::SHRED_TICK_REFERENCE_MASK.bits();
+        for tick in 0..=u8::MAX {
+            let flags = ShredFlags::from_reference_tick(tick);
+            assert_eq!(flags.bits(), tick.min(MAX_REFERENCE_TICK));
+        }
+    }
+
     #[test]
     fn test_version_from_hash() {
         let hash = [

diff --git a/ledger/src/shred/legacy.rs b/ledger/src/shred/legacy.rs
@@ -227,12 +227,7 @@ impl ShredData {
             fec_set_index,
         };
         let size = (data.len() + Self::SIZE_OF_HEADERS) as u16;
-        let flags = flags
-            | ShredFlags::from_bits_retain(
-                ShredFlags::SHRED_TICK_REFERENCE_MASK
-                    .bits()
-                    .min(reference_tick),
-            );
+        let flags = flags | ShredFlags::from_reference_tick(reference_tick);
         let data_header = DataShredHeader {
             parent_offset,
             flags,

diff --git a/ledger/src/shred/merkle.rs b/ledger/src/shred/merkle.rs
@@ -1113,11 +1113,7 @@ pub(super) fn make_shreds_from_data(
             .checked_sub(parent_slot)
             .and_then(|offset| u16::try_from(offset).ok())
             .ok_or(Error::InvalidParentSlot { slot, parent_slot })?;
-        let flags = ShredFlags::from_bits_retain(
-            ShredFlags::SHRED_TICK_REFERENCE_MASK
-                .bits()
-                .min(reference_tick),
-        );
+        let flags = ShredFlags::from_reference_tick(reference_tick);
         DataShredHeader {
             parent_offset,
             flags,