near-primitives 0.36.0

use bitvec::prelude::*;
use borsh::{BorshDeserialize, BorshSerialize};
use near_schema_checker_lib::ProtocolSchema;

/// Represents a collection of bitmaps, one per shard, to store whether the endorsements from the chunk validators has been received.
///
/// For each shard, the endorsements are encoded as a sequence of bits: 1 means endorsement received and 0 means not received.
/// While the number of chunk validator seats is fixed, the number of chunk-validator assignments may be smaller and may change,
/// since the seats are assigned to validators weighted by their stake. Thus, we represent the bits as a vector of bytes.
/// The number of assignments may be less or equal to the number of total bytes. This representation allows increasing
/// the chunk validator seats in the future (which will be represented by a vector of greater length).
#[derive(
    BorshSerialize,
    BorshDeserialize,
    serde::Serialize,
    serde::Deserialize,
    Debug,
    Clone,
    Eq,
    PartialEq,
    Default,
    ProtocolSchema,
)]
pub struct ChunkEndorsementsBitmap {
    inner: Vec<Vec<u8>>,
}

/// Type of BitVec used in the internal implementation.
type BitVecType = BitVec<u8, Lsb0>;

impl ChunkEndorsementsBitmap {
    /// Creates an empty endorsement bitmap for each shard.
    pub fn new(num_shards: usize) -> Self {
        Self { inner: vec![Default::default(); num_shards] }
    }

    /// Returns the bitmap specifically for the genesis block.
    /// This matches the chunk endorsement signatures in the genesis block body,
    /// which is an empty vector (no shards).
    pub fn genesis() -> Self {
        Self::new(0)
    }

    /// Creates an endorsement bitmap from the provided bytes.
    pub fn from_bytes(bytes: Vec<Vec<u8>>) -> Self {
        Self { inner: bytes }
    }

    /// Returns the endorsements bits as a vector of bytes.
    /// The struct can be reconstructed by providing the output of this function to `new_from_bits`.
    pub fn bytes(&self) -> Vec<Vec<u8>> {
        self.inner.clone()
    }

    // Creates an endorsement bitmap for all the shards.
    pub fn from_endorsements(shards_to_endorsements: Vec<Vec<bool>>) -> Self {
        let mut bitmap = ChunkEndorsementsBitmap::new(shards_to_endorsements.len());
        for (shard_index, endorsements) in shards_to_endorsements.into_iter().enumerate() {
            bitmap.add_endorsements(shard_index, endorsements);
        }
        bitmap
    }

    /// Adds the provided endorsements to the bitmap for the specified shard.
    pub fn add_endorsements(&mut self, shard_index: usize, endorsements: Vec<bool>) {
        let bitvec: BitVecType = endorsements.iter().collect();
        self.inner[shard_index] = bitvec.into();
    }

    /// Returns an iterator over the endorsements (yields true if the endorsement for the respective position was received).
    pub fn iter(&self, shard_index: usize) -> Box<dyn Iterator<Item = bool>> {
        let bitvec = BitVecType::from_vec(self.inner[shard_index].clone());
        Box::new(bitvec.into_iter())
    }

    /// Returns the number of shards in the endorsements.
    pub fn num_shards(&self) -> usize {
        self.inner.len()
    }

    /// Returns the full length of the bitmap for a given shard.
    /// Note that the size may be greater than the number of validator assignments.
    pub fn len(&self, shard_index: usize) -> Option<usize> {
        self.inner.get(shard_index).map(|v| v.len() * 8)
    }
}

#[cfg(test)]
mod tests {
    use super::ChunkEndorsementsBitmap;
    use itertools::Itertools;
    use rand::Rng;

    const NUM_SHARDS: usize = 4;

    /// Asserts that the bitmap encodes the expected endorsements and zeros for the other bits.
    fn assert_bitmap(
        bitmap: &ChunkEndorsementsBitmap,
        num_assignments: usize,
        expected_endorsements: &Vec<Vec<bool>>,
    ) {
        // Endorsements from the bitmap iterator must match the endorsements given previously.
        for (shard_index, endorsements) in expected_endorsements.iter().enumerate() {
            let num_bits = bitmap.len(shard_index).unwrap();
            let bits = bitmap.iter(shard_index).collect_vec();
            // Number of bits must be equal to the size of the bit iterator for the corresponding shard.
            assert_eq!(num_bits, bits.len());
            // Bitmap must contain the minimal number of bits to represent the endorsements.
            assert_eq!(num_bits, num_assignments.div_ceil(8) * 8);
            // Endorsements from the bitmap iterator must match the endorsements given previously.
            let mut expected_bits = endorsements.clone();
            // All remaining bits after the endorsement bits must be false.
            expected_bits.extend(std::iter::repeat(false).take(num_bits - num_assignments));
            assert_eq!(bits, expected_bits);
        }
    }

    fn run_bitmap_test(num_assignments: usize, num_produced: usize) {
        let mut rng = rand::thread_rng();
        let mut bitmap = ChunkEndorsementsBitmap::new(NUM_SHARDS);
        let mut expected_endorsements = vec![];
        for shard_index in 0..NUM_SHARDS {
            let mut endorsements = vec![false; num_assignments];
            for _ in 0..num_produced {
                endorsements[rng.gen_range(0..num_assignments)] = true;
            }
            expected_endorsements.push(endorsements.clone());
            bitmap.add_endorsements(shard_index, endorsements);
        }
        // Check before serialization.
        assert_bitmap(&bitmap, num_assignments, &expected_endorsements);

        // Check after serialization.
        let serialized = borsh::to_vec(&bitmap).unwrap();
        let bitmap = borsh::from_slice::<ChunkEndorsementsBitmap>(&serialized).unwrap();
        assert_bitmap(&bitmap, num_assignments, &expected_endorsements);
    }

    #[test]
    fn test_chunk_endorsements_bitmap() {
        run_bitmap_test(0, 0);
        run_bitmap_test(60, 20);
        run_bitmap_test(68, 20);
        run_bitmap_test(125, 30);
        run_bitmap_test(130, 50);
        run_bitmap_test(300, 100);
    }
}