Skip to main content

data_anchor_blober/state/
bitmap.rs

1use anchor_lang::{prelude::*, AnchorDeserialize, AnchorSerialize, InitSpace};
2
3use crate::{constants::CHUNKS_BITMAP_SIZE, error::ErrorCode};
4
5#[derive(
6    Clone, Copy, PartialEq, Eq, PartialOrd, Ord, InitSpace, AnchorSerialize, AnchorDeserialize,
7)]
8pub struct Bitmap {
9    pub num_chunks: u16,
10    pub map: [u8; CHUNKS_BITMAP_SIZE as usize],
11}
12
13fn byte_containing_idx(idx: u16) -> usize {
14    (idx / 8) as usize
15}
16
17fn bit_offset_for_idx(idx: u16) -> usize {
18    (idx % 8) as usize
19}
20
21impl Bitmap {
22    pub fn new(number_of_chunks: u16) -> Self {
23        Self {
24            num_chunks: number_of_chunks,
25            map: [0; CHUNKS_BITMAP_SIZE as usize],
26        }
27    }
28
29    /// Mark the bit corresponding to the given index and return whether it was already set, panicking if it is out of bounds
30    ///
31    /// Panicking is the "correct" behavior since it will cause the Solana tx to revert
32    pub fn test_and_set(&mut self, idx: u16) -> std::result::Result<(), ErrorCode> {
33        if idx >= self.num_chunks || byte_containing_idx(idx) > self.map.len() {
34            panic!("chunk {idx} out of bounds");
35        }
36
37        let byte_offset = byte_containing_idx(idx);
38        let byte = self.map[byte_offset];
39
40        let bit_offset = bit_offset_for_idx(idx);
41        let bit_mask = 1 << bit_offset;
42
43        if byte & bit_mask != 0 {
44            return Err(ErrorCode::DuplicateChunk);
45        }
46
47        self.map[byte_offset] = byte | bit_mask;
48
49        Ok(())
50    }
51
52    /// Check if all bits are set to 1.
53    pub fn is_complete(&self) -> bool {
54        let limit = byte_containing_idx(self.num_chunks);
55        // Every byte except the last.
56        for i in 0..limit {
57            if self.map[i] != 0b11111111 {
58                return false;
59            }
60        }
61        // Then check the last one separately, which is not a full byte of ones.
62        self.map[limit] == (1 << bit_offset_for_idx(self.num_chunks) as u8) - 1
63    }
64}