pipecircuit 0.0.25

Correct. Robust. Fast. Portable. Pick four.
Documentation
#[cfg(test)]
use alloc::vec::Vec;

use crate::storage::StableSlot;

#[cfg(test)]
pub(in crate::storage::warp_table) struct RankedSlots {
    blocks: Vec<RankedBlock>,
    len: usize,
}

pub(super) struct RankedBlock {
    pub(super) slot_base: u32,
    pub(super) slots: u64,
}

impl RankedBlock {
    pub(super) const fn new(slot_base: u32, slots: u64) -> Self {
        Self { slot_base, slots }
    }
}

#[cfg(test)]
impl RankedSlots {
    pub(super) fn new(mut blocks: Vec<RankedBlock>) -> Self {
        blocks.retain(|block| block.slots != 0);
        blocks.sort_unstable_by_key(|block| block.slot_base);
        let len = blocks.iter().fold(0_usize, |len, block| {
            len.saturating_add(block.slots.count_ones() as usize)
        });
        Self { blocks, len }
    }

    pub(in crate::storage::warp_table) const fn len(&self) -> usize {
        self.len
    }

    pub(in crate::storage::warp_table) fn from_slots(mut slots: Vec<StableSlot>) -> Self {
        slots.sort_unstable();
        let mut blocks: Vec<RankedBlock> = Vec::new();
        for slot in slots {
            let slot = slot.value();
            let slot_base = slot & !super::SLOT_OFFSET_MASK;
            let mask = 1_u64.wrapping_shl(slot & super::SLOT_OFFSET_MASK);
            match blocks.last_mut() {
                Some(block) if block.slot_base == slot_base => block.slots |= mask,
                Some(_) | None => blocks.push(RankedBlock::new(slot_base, mask)),
            }
        }
        Self::new(blocks)
    }

    pub(in crate::storage::warp_table) fn slots(&self) -> impl Iterator<Item = StableSlot> + '_ {
        block_slots(&self.blocks)
    }
}

pub(super) fn block_slots(blocks: &[RankedBlock]) -> impl Iterator<Item = StableSlot> + '_ {
    blocks.iter().flat_map(|block| {
        let offsets = SlotOffsets(block.slots);
        #[cfg(test)]
        let offsets = offsets.inspect({
            let mut remaining = block.slots.count_ones();
            move |_| {
                assert_ne!(
                    remaining, 0,
                    "ranked slot iteration exceeded its member count"
                );
                remaining = remaining.saturating_sub(1);
            }
        });
        offsets.map(|offset| StableSlot::new(block.slot_base | offset))
    })
}

struct SlotOffsets(u64);

impl Iterator for SlotOffsets {
    type Item = u32;

    fn next(&mut self) -> Option<Self::Item> {
        if self.0 == 0 {
            return None;
        }
        let offset = self.0.trailing_zeros();
        self.0 &= self.0.wrapping_sub(1);
        Some(offset)
    }
}

#[cfg(test)]
mod tests {
    use alloc::vec::Vec;

    use super::RankedSlots;
    use crate::storage::StableSlot;

    #[test]
    fn direct_slots_are_grouped_and_restored_across_blocks() {
        let ranked = RankedSlots::from_slots(Vec::from([
            StableSlot::new(129),
            StableSlot::new(63),
            StableSlot::new(1),
            StableSlot::new(64),
        ]));

        assert_eq!(
            ranked.slots().map(StableSlot::value).collect::<Vec<_>>(),
            Vec::from([1, 63, 64, 129])
        );
        assert_eq!(ranked.blocks.len(), 3);
    }
}