fix bitvector test random impl

eserilev · May 28, 2024 · e87cf1d · e87cf1d
1 parent df983a8
commit e87cf1d
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/consensus/types/src/test_utils/test_random/bitfield.rs b/consensus/types/src/test_utils/test_random/bitfield.rs
@@ -26,6 +26,15 @@ impl<N: Unsigned + Clone> TestRandom for BitVector<N> {
     fn random_for_test(rng: &mut impl RngCore) -> Self {
         let mut raw_bytes = smallvec![0; std::cmp::max(1, (N::to_usize() + 7) / 8)];
         rng.fill_bytes(&mut raw_bytes);
+        // If N isn't divisible by 8
+        // zero out bits greater than N
+        if let Some(last_byte) = raw_bytes.last_mut() {
+            let mut mask = 0;
+            for i in 0..N::to_usize() % 8 {
+                mask |= 1 << i;
+            }
+            *last_byte &= mask;
+        }
         Self::from_bytes(raw_bytes).expect("we generate a valid BitVector")
     }
 }