feat(protocol): Span Batch Bits

crates/protocol/src/batch/bits.rs

@@ -0,0 +1,224 @@
+//! Module for working with span batch bits.
+
+use crate::SpanBatchError;
+use alloc::{vec, vec::Vec};
+use alloy_rlp::Buf;
+use core::cmp::Ordering;
+
+/// Type for span batch bits.
+#[derive(Debug, Default, Clone, PartialEq, Eq)]
+pub struct SpanBatchBits(Vec<u8>);
+
+impl AsRef<[u8]> for SpanBatchBits {
+    fn as_ref(&self) -> &[u8] {
+        &self.0
+    }
+}
+
+impl SpanBatchBits {
+    /// Decodes a standard span-batch bitlist from a reader.
+    /// The bitlist is encoded as big-endian integer, left-padded with zeroes to a multiple of 8
+    /// bits. The encoded bitlist cannot be longer than `bit_length`.
+    pub fn decode(b: &mut &[u8], bit_length: usize) -> Result<Self, SpanBatchError> {
+        let buffer_len = bit_length / 8 + if bit_length % 8 != 0 { 1 } else { 0 };
+        let bits = if b.len() < buffer_len {
+            let mut bits = vec![0; buffer_len];
+            bits[..b.len()].copy_from_slice(b);
+            b.advance(b.len());
+            bits
+        } else {
+            let v = b[..buffer_len].to_vec();
+            b.advance(buffer_len);
+            v
+        };
+        let sb_bits = Self(bits);
+
+        if sb_bits.bit_len() > bit_length {
+            return Err(SpanBatchError::BitfieldTooLong);
+        }
+
+        Ok(sb_bits)
+    }
+
+    /// Encodes a standard span-batch bitlist.
+    /// The bitlist is encoded as big-endian integer, left-padded with zeroes to a multiple of 8
+    /// bits. The encoded bitlist cannot be longer than `bit_length`
+    pub fn encode(w: &mut Vec<u8>, bit_length: usize, bits: &Self) -> Result<(), SpanBatchError> {
+        if bits.bit_len() > bit_length {
+            return Err(SpanBatchError::BitfieldTooLong);
+        }
+
+        // Round up, ensure enough bytes when number of bits is not a multiple of 8.
+        // Alternative of (L+7)/8 is not overflow-safe.
+        let buf_len = bit_length / 8 + if bit_length % 8 != 0 { 1 } else { 0 };
+        let mut buf = vec![0; buf_len];
+        buf[buf_len - bits.0.len()..].copy_from_slice(bits.as_ref());
+        w.extend_from_slice(&buf);
+        Ok(())
+    }
+
+    /// Get a bit from the [SpanBatchBits] bitlist.
+    pub fn get_bit(&self, index: usize) -> Option<u8> {
+        let byte_index = index / 8;
+        let bit_index = index % 8;
+
+        // Check if the byte index is within the bounds of the bitlist
+        if byte_index < self.0.len() {
+            // Retrieve the specific byte that contains the bit we're interested in
+            let byte = self.0[self.0.len() - byte_index - 1];
+
+            // Shift the bits of the byte to the right, based on the bit index, and
+            // mask it with 1 to isolate the bit we're interested in.
+            // If the result is not zero, the bit is set to 1, otherwise it's 0.
+            Some(if byte & (1 << bit_index) != 0 { 1 } else { 0 })
+        } else {
+            // Return None if the index is out of bounds
+            None
+        }
+    }
+
+    /// Sets a bit in the [SpanBatchBits] bitlist.
+    pub fn set_bit(&mut self, index: usize, value: bool) {
+        let byte_index = index / 8;
+        let bit_index = index % 8;
+
+        // Ensure the vector is large enough to contain the bit at 'index'.
+        // If not, resize the vector, filling with 0s.
+        if byte_index >= self.0.len() {
+            Self::resize_from_right(&mut self.0, byte_index + 1);
+        }
+
+        // Retrieve the specific byte to modify
+        let len = self.0.len();
+        let byte = &mut self.0[len - byte_index - 1];
+
+        if value {
+            // Set the bit to 1
+            *byte |= 1 << bit_index;
+        } else {
+            // Set the bit to 0
+            *byte &= !(1 << bit_index);
+        }
+    }
+
+    /// Calculates the bit length of the [SpanBatchBits] bitfield.
+    pub fn bit_len(&self) -> usize {
+        // Iterate over the bytes from left to right to find the first non-zero byte
+        for (i, &byte) in self.0.iter().enumerate() {
+            if byte != 0 {
+                // Calculate the index of the most significant bit in the byte
+                let msb_index = 7 - byte.leading_zeros() as usize; // 0-based index
+
+                // Calculate the total bit length
+                let total_bit_length = msb_index + 1 + ((self.0.len() - i - 1) * 8);
+                return total_bit_length;
+            }
+        }
+
+        // If all bytes are zero, the bitlist is considered to have a length of 0
+        0
+    }
+
+    /// Resizes an array from the right. Useful for big-endian zero extension.
+    fn resize_from_right<T: Default + Clone>(vec: &mut Vec<T>, new_size: usize) {
+        let current_size = vec.len();
+        match new_size.cmp(&current_size) {
+            Ordering::Less => {
+                // Remove elements from the beginning.
+                let remove_count = current_size - new_size;
+                vec.drain(0..remove_count);
+            }
+            Ordering::Greater => {
+                // Calculate how many new elements to add.
+                let additional = new_size - current_size;
+                // Prepend new elements with default values.
+                let mut prepend_elements = vec![T::default(); additional];
+                prepend_elements.append(vec);
+                *vec = prepend_elements;
+            }
+            Ordering::Equal => { /* If new_size == current_size, do nothing. */ }
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use proptest::{collection::vec, prelude::any, proptest};
+
+    proptest! {
+        #[test]
+        fn test_encode_decode_roundtrip_span_bitlist(vec in vec(any::<u8>(), 0..5096)) {
+            let bits = SpanBatchBits(vec);
+            assert_eq!(SpanBatchBits::decode(&mut bits.as_ref(), bits.0.len() * 8).unwrap(), bits);
+            let mut encoded = Vec::new();
+            SpanBatchBits::encode(&mut encoded, bits.0.len() * 8, &bits).unwrap();
+            assert_eq!(encoded, bits.0);
+        }
+
+        #[test]
+        fn test_span_bitlist_bitlen(index in 0usize..65536) {
+            let mut bits = SpanBatchBits::default();
+            bits.set_bit(index, true);
+            assert_eq!(bits.0.len(), (index / 8) + 1);
+            assert_eq!(bits.bit_len(), index + 1);
+        }
+
+        #[test]
+        fn test_span_bitlist_bitlen_shrink(first_index in 8usize..65536) {
+            let second_index = first_index.clamp(0, first_index - 8);
+            let mut bits = SpanBatchBits::default();
+
+            // Set and clear first index.
+            bits.set_bit(first_index, true);
+            assert_eq!(bits.0.len(), (first_index / 8) + 1);
+            assert_eq!(bits.bit_len(), first_index + 1);
+            bits.set_bit(first_index, false);
+            assert_eq!(bits.0.len(), (first_index / 8) + 1);
+            assert_eq!(bits.bit_len(), 0);
+
+            // Set second bit. Even though the array is larger, as it was originally allocated with more words,
+            // the bitlength should still be lowered as the higher-order words are 0'd out.
+            bits.set_bit(second_index, true);
+            assert_eq!(bits.0.len(), (first_index / 8) + 1);
+            assert_eq!(bits.bit_len(), second_index + 1);
+        }
+    }
+
+    #[test]
+    fn bitlist_big_endian_zero_extended() {
+        let mut bits = SpanBatchBits::default();
+
+        bits.set_bit(1, true);
+        bits.set_bit(6, true);
+        bits.set_bit(8, true);
+        bits.set_bit(15, true);
+        assert_eq!(bits.0[0], 0b1000_0001);
+        assert_eq!(bits.0[1], 0b0100_0010);
+        assert_eq!(bits.0.len(), 2);
+        assert_eq!(bits.bit_len(), 16);
+    }
+
+    #[test]
+    fn test_static_set_get_bits_span_bitlist() {
+        let mut bits = SpanBatchBits::default();
+        assert!(bits.0.is_empty());
+
+        bits.set_bit(0, true);
+        bits.set_bit(1, true);
+        bits.set_bit(2, true);
+        bits.set_bit(4, true);
+        bits.set_bit(7, true);
+        assert_eq!(bits.0.len(), 1);
+        assert_eq!(bits.get_bit(0), Some(1));
+        assert_eq!(bits.get_bit(1), Some(1));
+        assert_eq!(bits.get_bit(2), Some(1));
+        assert_eq!(bits.get_bit(3), Some(0));
+        assert_eq!(bits.get_bit(4), Some(1));
+
+        bits.set_bit(17, true);
+        assert_eq!(bits.get_bit(17), Some(1));
+        assert_eq!(bits.get_bit(32), None);
+        assert_eq!(bits.0.len(), 3);
+    }
+}

crates/protocol/src/batch/mod.rs

@@ -6,6 +6,9 @@ pub use r#type::*;
 mod errors;
 pub use errors::{SpanBatchError, SpanDecodingError};
 
+mod bits;
+pub use bits::SpanBatchBits;
+
 mod element;
 pub use element::{SpanBatchElement, MAX_SPAN_BATCH_ELEMENTS};
 

crates/protocol/src/lib.rs

@@ -11,8 +11,9 @@ extern crate alloc;
 
 mod batch;
 pub use batch::{
-    BatchType, BatchValidationProvider, BatchValidity, SingleBatch, SpanBatchElement,
-    SpanBatchError, SpanDecodingError, MAX_SPAN_BATCH_ELEMENTS, SINGLE_BATCH_TYPE, SPAN_BATCH_TYPE,
+    BatchType, BatchValidationProvider, BatchValidity, SingleBatch, SpanBatchBits,
+    SpanBatchElement, SpanBatchError, SpanDecodingError, MAX_SPAN_BATCH_ELEMENTS,
+    SINGLE_BATCH_TYPE, SPAN_BATCH_TYPE,
 };
 
 mod block;