qdrant-edge 0.7.2

mod gaps;

use std::ops::Range;
use std::path::{Path, PathBuf};

use ahash::AHashSet;
use crate::common::bitvec::BitSlice;
use crate::common::mmap::{Advice, AdviceSetting, create_and_ensure_length};
use crate::common::stored_bitslice::StoredBitSlice;
use crate::common::universal_io::{MmapFile, OpenOptions, Populate, UniversalWrite};
use gaps::{BitmaskGaps, RegionGaps};
use itertools::Itertools;

use crate::gridstore::Result;
use crate::gridstore::config::StorageConfig;
use crate::gridstore::error::GridstoreError;
use crate::gridstore::tracker::{BlockOffset, PageId};

const BITMASK_NAME: &str = "bitmask.dat";

fn open_options() -> OpenOptions {
    OpenOptions {
        writeable: true,
        need_sequential: false,
        populate: Populate::No,
        advice: AdviceSetting::Advice(Advice::Random),
    }
}

type RegionId = u32;

/// Concrete bitmask type using memory-mapped storage.
pub type MmapBitmask = Bitmask<MmapFile>;

#[derive(Debug)]
pub struct Bitmask<S> {
    config: StorageConfig,

    /// A summary of every 1KB (8_192 bits) of contiguous zeros in the bitmask, or less if it is the last region.
    regions_gaps: BitmaskGaps<S>,

    /// The actual bitmask. Each bit represents a block. A 1 means the block is used, a 0 means it is free.
    bitslice: StoredBitSlice<S>,

    /// The path to the file containing the bitmask.
    path: PathBuf,
}

impl<S: UniversalWrite> Bitmask<S> {
    pub fn files(&self) -> Vec<PathBuf> {
        vec![self.path.clone(), self.regions_gaps.path()]
    }

    /// Calculate the amount of trailing free blocks in the bitmask.
    pub fn trailing_free_blocks(&self) -> Result<u32> {
        let trailing_gap = self.regions_gaps.trailing_free_blocks()?;
        #[cfg(debug_assertions)]
        {
            let all_bits = self.bitslice.read_all()?;
            let num_trailing_zeros = all_bits.trailing_zeros();
            debug_assert_eq!(num_trailing_zeros, trailing_gap as usize);
        }

        Ok(trailing_gap)
    }

    /// Calculate the amount of bytes needed for covering the blocks of a page.
    fn length_for_page(config: &StorageConfig) -> usize {
        assert_eq!(
            config.page_size_bytes % config.block_size_bytes,
            0,
            "Page size must be a multiple of block size"
        );

        // one bit per block
        let bits = config.page_size_bytes / config.block_size_bytes;

        // length in bytes
        bits / u8::BITS as usize
    }

    /// Create a bitmask for one page
    pub(crate) fn create(fs: &S::Fs, dir: &Path, config: StorageConfig) -> Result<Self> {
        debug_assert!(
            config.page_size_bytes % config.block_size_bytes * config.region_size_blocks == 0,
            "Page size must be a multiple of block size * region size"
        );

        let length = Self::length_for_page(&config);

        // create bitmask file
        let path = Self::bitmask_path(dir);
        create_and_ensure_length(&path, length)?;

        let bitslice = StoredBitSlice::open(fs, &path, open_options(), Default::default())?;

        let bit_len = bitslice.bit_len() as usize;
        assert_eq!(bit_len, length * 8, "Bitmask length mismatch");

        // create regions gaps
        let num_regions = bit_len / config.region_size_blocks;
        let region_gaps = vec![RegionGaps::all_free(config.region_size_blocks as u16); num_regions];

        let regions_gaps = BitmaskGaps::create(fs, dir, region_gaps.into_iter(), config.clone())?;

        Ok(Self {
            config,
            regions_gaps,
            bitslice,
            path,
        })
    }

    pub(crate) fn open(fs: &S::Fs, dir: &Path, config: StorageConfig) -> Result<Self> {
        debug_assert!(
            config
                .page_size_bytes
                .is_multiple_of(config.block_size_bytes),
            "Page size must be a multiple of block size"
        );

        let path = Self::bitmask_path(dir);
        if !path.exists() {
            return Err(GridstoreError::service_error(format!(
                "Bitmask file does not exist: {}",
                path.display()
            )));
        }

        let bitslice = StoredBitSlice::open(
            fs,
            &path,
            OpenOptions {
                writeable: true,
                need_sequential: false,
                populate: Populate::Auto,
                advice: AdviceSetting::Advice(Advice::Random),
            },
            Default::default(),
        )?;
        let regions_gaps = BitmaskGaps::open(fs, dir, config.clone())?;

        Ok(Self {
            config,
            regions_gaps,
            bitslice,
            path,
        })
    }

    fn bitmask_path(dir: &Path) -> PathBuf {
        dir.join(BITMASK_NAME)
    }

    pub fn flusher(&self) -> impl FnOnce() -> Result<()> + Send + use<S> {
        let bitslice_flusher = self.bitslice.flusher();
        let gaps_flusher = self.regions_gaps.flusher();
        move || {
            bitslice_flusher()?;
            gaps_flusher()?;
            Ok(())
        }
    }

    /// Compute the size of the storage in bytes.
    /// Does not include the metadata information (e.g. the regions gaps, bitmask...).
    pub fn get_storage_size_bytes(&self) -> Result<usize> {
        let mut size = 0;
        let region_size_blocks = self.config.region_size_blocks;
        let block_size_bytes = self.config.block_size_bytes;
        let region_size_bytes = region_size_blocks * block_size_bytes;
        let gaps = self.regions_gaps.read_all()?;
        let all_bits = self.bitslice.read_all()?;
        for (gap_id, gap) in gaps.iter().enumerate() {
            // skip empty regions
            if gap.is_empty(region_size_blocks as u16) {
                continue;
            }
            // fast path for full regions
            if gap.is_full() {
                size += region_size_bytes;
            } else {
                // compute the size of the occupied blocks for the region
                let gap_offset_start = gap_id * region_size_blocks;
                let gap_offset_end = gap_offset_start + region_size_blocks;
                let occupied_blocks = all_bits[gap_offset_start..gap_offset_end].count_ones();
                size += occupied_blocks * block_size_bytes
            }
        }
        Ok(size)
    }

    pub fn infer_num_pages(&self) -> usize {
        let bits = self.bitslice.bit_len() as usize;
        let covered_bytes = bits * self.config.block_size_bytes;
        covered_bytes.div_euclid(self.config.page_size_bytes)
    }

    /// Extend the bitslice to cover another page
    pub fn cover_new_page(&mut self) -> Result<()> {
        let extra_length = Self::length_for_page(&self.config);

        // flush outstanding changes
        self.bitslice.flusher()()?;

        // reopen the file with a larger size
        let previous_bit_len = self.bitslice.bit_len() as usize;
        let new_length = (previous_bit_len / u8::BITS as usize) + extra_length;
        create_and_ensure_length(&self.path, new_length)?;

        self.bitslice.reopen()?;

        let current_bit_len = self.bitslice.bit_len() as usize;

        // extend the region gaps
        let current_total_regions = self.regions_gaps.len()?;
        let expected_total_full_regions =
            current_bit_len.div_euclid(self.config.region_size_blocks);
        debug_assert!(
            current_bit_len.is_multiple_of(self.config.region_size_blocks),
            "Bitmask length must be a multiple of region size"
        );
        let new_regions = expected_total_full_regions.saturating_sub(current_total_regions);
        let new_gaps =
            vec![RegionGaps::all_free(self.config.region_size_blocks as u16); new_regions];
        self.regions_gaps.extend(new_gaps.into_iter())?;

        assert_eq!(
            self.regions_gaps.len()? * self.config.region_size_blocks,
            current_bit_len,
            "Bitmask length mismatch",
        );

        Ok(())
    }

    fn range_of_page(&self, page_id: PageId) -> Range<usize> {
        let page_blocks = self.config.page_size_bytes / self.config.block_size_bytes;
        let start = page_id as usize * page_blocks;
        let end = start + page_blocks;
        start..end
    }

    /// The amount of blocks that have never been used in the page.
    #[cfg(test)]
    pub(crate) fn free_blocks_for_page(&self, page_id: PageId) -> Result<usize> {
        let range_of_page = self.range_of_page(page_id);
        let all_bits = self.bitslice.read_all()?;
        Ok(all_bits[range_of_page].trailing_zeros())
    }

    pub(crate) fn find_available_blocks(
        &self,
        num_blocks: u32,
    ) -> Result<Option<(PageId, BlockOffset)>> {
        let Some(region_id_range) = self.regions_gaps.find_fitting_gap(num_blocks)? else {
            return Ok(None);
        };
        let regions_start_offset = region_id_range.start as usize * self.config.region_size_blocks;
        let regions_end_offset = region_id_range.end as usize * self.config.region_size_blocks;

        let translate_to_answer = |local_index: u32| {
            let page_size_in_blocks = self.config.page_size_bytes / self.config.block_size_bytes;

            let global_cursor_offset = local_index as usize + regions_start_offset;

            // Calculate the page id and the block offset within the page
            let page_id = global_cursor_offset.div_euclid(page_size_in_blocks);
            let page_block_offset = global_cursor_offset.rem_euclid(page_size_in_blocks);

            (page_id as PageId, page_block_offset as BlockOffset)
        };

        let all_bits = self.bitslice.read_all()?;
        let regions_bitslice = &all_bits[regions_start_offset..regions_end_offset];

        Ok(Self::find_available_blocks_in_slice(
            regions_bitslice,
            num_blocks,
            translate_to_answer,
        ))
    }

    pub fn find_available_blocks_in_slice<F>(
        bitslice: &BitSlice,
        num_blocks: u32,
        translate_local_index: F,
    ) -> Option<(PageId, BlockOffset)>
    where
        F: FnOnce(u32) -> (PageId, BlockOffset),
    {
        // Get raw memory region
        let (head, raw_region, tail) = bitslice
            .domain()
            .region()
            .expect("Regions cover more than one usize");

        // We expect the regions to not use partial usizes
        debug_assert!(head.is_none());
        debug_assert!(tail.is_none());

        let mut current_size: u32 = 0;
        let mut current_start: u32 = 0;
        let mut num_shifts = 0;
        // Iterate over the integers that compose the bitvec. So that we can perform bitwise operations.
        const BITS_IN_CHUNK: u32 = usize::BITS;
        for (chunk_idx, chunk) in raw_region.iter().enumerate() {
            let mut chunk = *chunk;

            // case of all zeros
            if chunk == 0 {
                current_size += BITS_IN_CHUNK;
                continue;
            }

            if chunk == !0 {
                // case of all ones
                if current_size >= num_blocks {
                    // bingo - we found a free cell of num_blocks
                    return Some(translate_local_index(current_start));
                }
                current_size = 0;
                current_start = (chunk_idx as u32 + 1) * BITS_IN_CHUNK;
                continue;
            }

            // At least one non-zero bit
            let leading = chunk.trailing_zeros();
            let trailing = chunk.leading_zeros();

            let max_possible_middle_gap = (BITS_IN_CHUNK - leading - trailing).saturating_sub(2);

            // Skip looking for local max if it won't improve global max
            if num_blocks > max_possible_middle_gap {
                current_size += leading;
                if current_size >= num_blocks {
                    // bingo - we found a free cell of num_blocks
                    return Some(translate_local_index(current_start));
                }
                current_size = trailing;
                current_start = (chunk_idx as u32) * BITS_IN_CHUNK + BITS_IN_CHUNK - trailing;
                continue;
            }

            while chunk != 0 {
                let num_zeros = chunk.trailing_zeros();
                current_size += num_zeros;
                if current_size >= num_blocks {
                    // bingo - we found a free cell of num_blocks
                    return Some(translate_local_index(current_start));
                }

                // shift by the number of zeros
                chunk >>= num_zeros as usize;
                num_shifts += num_zeros;

                // skip consecutive ones
                let num_ones = chunk.trailing_ones();
                if num_ones < BITS_IN_CHUNK {
                    chunk >>= num_ones;
                } else {
                    // all ones
                    debug_assert!(chunk == !0);
                    chunk = 0;
                }
                num_shifts += num_ones;

                current_size = 0;
                current_start = chunk_idx as u32 * BITS_IN_CHUNK + num_shifts;
            }
            // no more ones in the chunk
            current_size += BITS_IN_CHUNK - num_shifts;
            num_shifts = 0;
        }
        if current_size >= num_blocks {
            // bingo - we found a free cell of num_blocks
            return Some(translate_local_index(current_start));
        }

        None
    }

    pub(crate) fn mark_blocks(
        &mut self,
        page_id: PageId,
        block_offset: BlockOffset,
        num_blocks: u32,
        used: bool,
    ) -> Result<()> {
        let relative_range = block_offset as usize..(block_offset as usize + num_blocks as usize);
        self.mark_blocks_batch(page_id, std::iter::once(relative_range), used)
    }

    /// Marks blocks sharing the same page in batch. First updates all ranges in the bitmask, then updates the region gaps a single time.
    ///
    /// # Arguments
    /// * `page_id` - The ID of the page to mark blocks on.
    /// * `block_ranges` - An iterator over the ranges of blocks to mark, relative to the page start.
    /// * `used` - Whether the blocks should be marked as used or free.
    pub(crate) fn mark_blocks_batch(
        &mut self,
        page_id: PageId,
        local_block_ranges: impl Iterator<Item = Range<usize>>,
        used: bool,
    ) -> Result<()> {
        let page_start = self.range_of_page(page_id).start;

        let est_num_ranges = local_block_ranges.size_hint().1.unwrap_or(1);
        let mut dirty_regions = AHashSet::with_capacity(est_num_ranges);

        for range in local_block_ranges {
            let bitmask_range = (range.start + page_start)..(range.end + page_start);

            self.bitslice.set_ascending_bits_batch(
                (bitmask_range.start as u64..bitmask_range.end as u64).map(|i| (i, used)),
            )?;

            let start_region_id =
                (bitmask_range.start / self.config.region_size_blocks) as RegionId;
            let end_region_id =
                bitmask_range.end.div_ceil(self.config.region_size_blocks) as RegionId;

            dirty_regions.extend(start_region_id..end_region_id);
        }

        self.update_region_gaps(dirty_regions)
    }

    fn update_region_gaps(&mut self, dirty_regions: AHashSet<RegionId>) -> Result<()> {
        for region_id in dirty_regions {
            let region_id = region_id as usize;
            let region_start = (region_id * self.config.region_size_blocks) as u64;
            let region_end = region_start + self.config.region_size_blocks as u64;

            let bitslice = &self.bitslice.read_bit_range(region_start..region_end)?;

            let gaps = Self::calculate_gaps(bitslice, self.config.region_size_blocks);

            self.regions_gaps.set(region_id, gaps)?;
        }
        Ok(())
    }

    pub fn calculate_gaps(region: &BitSlice, region_size_blocks: usize) -> RegionGaps {
        debug_assert_eq!(region.len(), region_size_blocks, "Unexpected region size");
        // Get raw memory region
        let (head, raw_region, tail) = region
            .domain()
            .region()
            .expect("Region covers more than one usize");

        // We expect the region to not use partial usizes
        debug_assert!(head.is_none());
        debug_assert!(tail.is_none());

        // Iterate over the integers that compose the bitslice. So that we can perform bitwise operations.
        let mut max = 0;
        let mut current = 0;
        const BITS_IN_CHUNK: u32 = usize::BITS;
        let mut num_shifts = 0;
        // In reverse, because we expect the regions to be filled start to end.
        // So starting from the end should give us bigger `max` earlier.
        for chunk in raw_region.iter().rev() {
            // Ensure that the chunk is little-endian.
            let mut chunk = chunk.to_le();
            // case of all zeros
            if chunk == 0 {
                current += BITS_IN_CHUNK;
                continue;
            }

            if chunk == !0 {
                // case of all ones
                max = max.max(current);
                current = 0;
                continue;
            }

            // At least one non-zero bit
            let leading = chunk.leading_zeros();
            let trailing = chunk.trailing_zeros();

            let max_possible_middle_gap = (BITS_IN_CHUNK - leading - trailing).saturating_sub(2);

            // Skip looking for local max if it won't improve global max
            if max > max_possible_middle_gap {
                current += leading;
                max = max.max(current);
                current = trailing;
                continue;
            }

            // Otherwise, look for the actual maximum in the chunk
            while chunk != 0 {
                // count consecutive zeros
                let num_zeros = chunk.leading_zeros();
                current += num_zeros;
                max = max.max(current);
                current = 0;

                // shift by the number of zeros
                chunk <<= num_zeros as usize;
                num_shifts += num_zeros;

                // skip consecutive ones
                let num_ones = chunk.leading_ones();
                if num_ones < BITS_IN_CHUNK {
                    chunk <<= num_ones;
                } else {
                    // all ones
                    debug_assert!(chunk == !0);
                    chunk = 0;
                }
                num_shifts += num_ones;
            }

            // no more ones in the chunk
            current += BITS_IN_CHUNK - num_shifts;
            num_shifts = 0;
        }

        max = max.max(current);

        let leading;
        let trailing;
        if max == region_size_blocks as u32 {
            leading = max;
            trailing = max;
        } else {
            leading = raw_region
                .iter()
                .take_while_inclusive(|chunk| chunk == &&0)
                .map(|chunk| chunk.trailing_zeros())
                .sum::<u32>();
            trailing = raw_region
                .iter()
                .rev()
                .take_while_inclusive(|chunk| chunk == &&0)
                .map(|chunk| chunk.leading_zeros())
                .sum::<u32>();
        }

        #[cfg(debug_assertions)]
        {
            RegionGaps::new(
                leading as u16,
                trailing as u16,
                max as u16,
                region_size_blocks as u16,
            )
        }

        #[cfg(not(debug_assertions))]
        {
            RegionGaps::new(leading as u16, trailing as u16, max as u16)
        }
    }

    /// Populate all pages in the mmap.
    /// Block until all pages are populated.
    pub fn populate(&self) -> Result<()> {
        self.bitslice.populate()?;
        self.regions_gaps.populate()?;
        Ok(())
    }

    /// Drop disk cache.
    pub fn clear_cache(&self) -> Result<()> {
        self.bitslice.clear_ram_cache()?;
        self.regions_gaps.clear_cache()?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {

    use bitvec::bits;
    use crate::common::bitvec::BitVec;
    use crate::common::universal_io::MmapFs;
    use proptest::prelude::*;
    use rand::{RngExt, rng};

    use crate::gridstore::bitmask::MmapBitmask;
    use crate::gridstore::config::{DEFAULT_BLOCK_SIZE_BYTES, DEFAULT_REGION_SIZE_BLOCKS, StorageOptions};

    #[test]
    fn test_length_for_page() {
        let config = &StorageOptions {
            page_size_bytes: Some(8192),
            region_size_blocks: Some(1),
            ..Default::default()
        }
        .try_into()
        .unwrap();
        assert_eq!(MmapBitmask::length_for_page(config), 8);
    }

    #[test]
    fn test_find_available_blocks() {
        let page_size = DEFAULT_BLOCK_SIZE_BYTES * DEFAULT_REGION_SIZE_BLOCKS;

        let blocks_per_page = (page_size / DEFAULT_BLOCK_SIZE_BYTES) as u32;

        let dir = tempfile::tempdir().unwrap();

        let options = StorageOptions {
            page_size_bytes: Some(page_size),
            ..Default::default()
        };

        let mut bitmask: MmapBitmask =
            super::Bitmask::create(&MmapFs, dir.path(), options.try_into().unwrap()).unwrap();
        bitmask.cover_new_page().unwrap();

        assert_eq!(bitmask.bitslice.bit_len() as u32, blocks_per_page * 2);

        // 1..10
        bitmask.mark_blocks(0, 1, 9, true).unwrap();

        // 15..20
        bitmask.mark_blocks(0, 15, 5, true).unwrap();

        // 30..blocks_per_page
        bitmask
            .mark_blocks(0, 30, blocks_per_page - 30, true)
            .unwrap();

        // blocks_per_page..blocks_per_page + 1
        bitmask.mark_blocks(1, 0, 1, true).unwrap();

        let (page_id, block_offset) = bitmask.find_available_blocks(1).unwrap().unwrap();
        assert_eq!(block_offset, 0);
        assert_eq!(page_id, 0);

        let (page_id, block_offset) = bitmask.find_available_blocks(2).unwrap().unwrap();
        assert_eq!(block_offset, 10);
        assert_eq!(page_id, 0);

        let (page_id, block_offset) = bitmask.find_available_blocks(5).unwrap().unwrap();
        assert_eq!(block_offset, 10);
        assert_eq!(page_id, 0);

        let (page_id, block_offset) = bitmask.find_available_blocks(6).unwrap().unwrap();
        assert_eq!(block_offset, 20);
        assert_eq!(page_id, 0);

        // first free block of the next page
        let (page_id, block_offset) = bitmask.find_available_blocks(30).unwrap().unwrap();
        assert_eq!(block_offset, 1);
        assert_eq!(page_id, 1);

        // not fitting cell
        let found_large = bitmask.find_available_blocks(blocks_per_page).unwrap();
        assert_eq!(found_large, None);
    }

    #[test]
    fn test_raw_bitvec() {
        use bitvec::prelude::Lsb0;
        let bits = bits![
            0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
            1, 1, 1, 1, 1, 1
        ];

        let mut bitvec = BitVec::new();
        bitvec.extend_from_bitslice(bits);

        assert_eq!(bitvec.len(), 64);

        let raw = bitvec.as_raw_slice();
        assert_eq!(raw.len() as u32, 64 / usize::BITS);

        assert_eq!(raw[0].trailing_zeros(), 4);
        assert_eq!(raw[0].leading_zeros(), 0);
        assert_eq!((raw[0] >> 1).trailing_zeros(), 3)
    }

    prop_compose! {
        /// Creates a fixture bitvec which has gaps of a specific size
        fn regions_bitvec_with_max_gap(max_gap_size: usize) (len in 0..DEFAULT_REGION_SIZE_BLOCKS*4) -> (BitVec, usize) {
            assert!(max_gap_size > 0);
            let len = len.next_multiple_of(DEFAULT_REGION_SIZE_BLOCKS);

            let mut bitvec = BitVec::new();
            bitvec.resize(len, true);

            let mut rng = rng();

            let mut i = 0;
            let mut max_gap = 0;
            while i < len {
                let run = rng.random_range(1..max_gap_size).min(len - i);
                let skip = rng.random_range(1..max_gap_size);

                for j in 0..run {
                    bitvec.set(i + j, false);
                }

                if run > max_gap {
                    max_gap = run;
                }

                i += run + skip;
            }

            (bitvec, max_gap)
        }
    }
    proptest! {
        #![proptest_config(ProptestConfig::with_cases(1000))]

        #[test]
        fn test_find_available_blocks_properties((bitvec, max_gap) in regions_bitvec_with_max_gap(120)) {
            let bitslice = bitvec.as_bitslice();

            // Helper to check if a range is all zeros
            let is_free_range = |start: usize, len: usize| {
                let range = start..(start + len);
                bitslice.get(range)
                    .map(|slice| slice.not_any())
                    .unwrap_or(false)
            };

            // For different requested block sizes
            for req_blocks in 1..=max_gap {
                if let Some((_, block_offset)) = MmapBitmask::find_available_blocks_in_slice(
                    bitslice,
                    req_blocks as u32,
                    |idx| (0, idx),
                ) {
                    // The found position should have enough free blocks
                    prop_assert!(is_free_range(block_offset as usize, req_blocks));
                } else {
                    prop_assert!(false, "Should've found a free range")
                }
            }

            // For a block size that doesn't fit
            let req_blocks = max_gap + 1;
            prop_assert!(MmapBitmask::find_available_blocks_in_slice(
                bitslice,
                req_blocks as u32,
                |idx| (0, idx),
            ).is_none());
        }
    }
}