pipecircuit 0.0.18

use super::{RangeIndex, RangeIndexInconsistency, RangeIndexKey};
use crate::storage::warp_table::secondary_index::StableSlot;

const SLOTS_PER_BLOCK: usize = u64::BITS as usize;

#[derive(Clone, Copy)]
pub(super) struct SlotLocation {
    start: usize,
    end: usize,
    mask: u64,
}

pub(super) fn storage_dimensions(maximum_rows: usize, bits: u32) -> Option<(usize, usize)> {
    let bit_count = usize::try_from(bits).ok()?;
    let words_per_block = bit_count.checked_add(1)?;
    let block_count = maximum_rows.div_ceil(SLOTS_PER_BLOCK);
    let storage_words = block_count.checked_mul(words_per_block)?;
    Some((words_per_block, storage_words))
}

impl<Key: RangeIndexKey> RangeIndex<Key> {
    pub(super) fn validate_storage(&self) -> Result<(), RangeIndexInconsistency> {
        self.validated_words_per_block().map(|_| ())
    }

    pub(super) fn validated_words_per_block(&self) -> Result<usize, RangeIndexInconsistency> {
        let Some((words_per_block, expected)) = storage_dimensions(self.maximum_rows, Key::BITS)
        else {
            return Err(RangeIndexInconsistency::StorageCapacityOverflow {
                maximum_rows: self.maximum_rows,
                bits: Key::BITS,
            });
        };
        if self.blocks.len() != expected {
            return Err(RangeIndexInconsistency::BlockStorageLengthMismatch {
                expected,
                actual: self.blocks.len(),
            });
        }
        Ok(words_per_block)
    }

    pub(super) fn locate_slot(
        &self,
        slot: StableSlot,
    ) -> Result<SlotLocation, RangeIndexInconsistency> {
        let slot_index = usize::try_from(slot.value()).map_err(|_| {
            RangeIndexInconsistency::SlotOutsideCapacity {
                maximum_rows: self.maximum_rows,
            }
        })?;
        if slot_index >= self.maximum_rows {
            return Err(RangeIndexInconsistency::SlotOutsideCapacity {
                maximum_rows: self.maximum_rows,
            });
        }
        let Some((words_per_block, _)) = storage_dimensions(self.maximum_rows, Key::BITS) else {
            return Err(RangeIndexInconsistency::StorageCapacityOverflow {
                maximum_rows: self.maximum_rows,
                bits: Key::BITS,
            });
        };
        let block = slot_index / SLOTS_PER_BLOCK;
        let Some(start) = block.checked_mul(words_per_block) else {
            return Err(RangeIndexInconsistency::StorageCapacityOverflow {
                maximum_rows: self.maximum_rows,
                bits: Key::BITS,
            });
        };
        let Some(end) = start.checked_add(words_per_block) else {
            return Err(RangeIndexInconsistency::StorageCapacityOverflow {
                maximum_rows: self.maximum_rows,
                bits: Key::BITS,
            });
        };
        let bit = u32::try_from(slot_index % SLOTS_PER_BLOCK)
            .map_err(|_| RangeIndexInconsistency::BitPositionOutsideWord { bit: u32::MAX })?;
        let mask = 1_u64
            .checked_shl(bit)
            .ok_or(RangeIndexInconsistency::BitPositionOutsideWord { bit })?;
        Ok(SlotLocation { start, end, mask })
    }

    pub(super) fn member_at(
        &self,
        location: SlotLocation,
    ) -> Result<bool, RangeIndexInconsistency> {
        let (members, _) = self.block_at(location)?;
        Ok(*members & location.mask != 0)
    }

    pub(super) fn word_at(&self, location: SlotLocation) -> Result<u64, RangeIndexInconsistency> {
        let (_, planes) = self.block_at(location)?;
        let mut word = 0_u64;
        for (bit, plane) in planes.iter().enumerate() {
            if *plane & location.mask != 0 {
                let bit = u32::try_from(bit).map_err(|_| {
                    RangeIndexInconsistency::BitPositionOutsideWord { bit: u32::MAX }
                })?;
                let mask = 1_u64
                    .checked_shl(bit)
                    .ok_or(RangeIndexInconsistency::BitPositionOutsideWord { bit })?;
                word |= mask;
            }
        }
        Ok(word)
    }

    pub(super) fn write_slot(
        &mut self,
        location: SlotLocation,
        present: bool,
        word: u64,
    ) -> Result<(), RangeIndexInconsistency> {
        let actual = self.blocks.len();
        let Some(block) = self.blocks.get_mut(location.start..location.end) else {
            return Err(RangeIndexInconsistency::BlockStorageLengthMismatch {
                expected: location.end,
                actual,
            });
        };
        let Some((members, planes)) = block.split_first_mut() else {
            return Err(RangeIndexInconsistency::BlockStorageLengthMismatch {
                expected: location.end,
                actual,
            });
        };
        if present {
            *members |= location.mask;
        } else {
            *members &= !location.mask;
        }
        for (bit, plane) in planes.iter_mut().enumerate() {
            let bit = u32::try_from(bit)
                .map_err(|_| RangeIndexInconsistency::BitPositionOutsideWord { bit: u32::MAX })?;
            let bit_mask = 1_u64
                .checked_shl(bit)
                .ok_or(RangeIndexInconsistency::BitPositionOutsideWord { bit })?;
            if word & bit_mask == 0 {
                *plane &= !location.mask;
            } else {
                *plane |= location.mask;
            }
        }
        Ok(())
    }

    fn block_at(&self, location: SlotLocation) -> Result<(&u64, &[u64]), RangeIndexInconsistency> {
        let Some(block) = self.blocks.get(location.start..location.end) else {
            return Err(RangeIndexInconsistency::BlockStorageLengthMismatch {
                expected: location.end,
                actual: self.blocks.len(),
            });
        };
        block
            .split_first()
            .ok_or(RangeIndexInconsistency::BlockStorageLengthMismatch {
                expected: location.end,
                actual: self.blocks.len(),
            })
    }
}