use alloc::vec::Vec;
use core::ops::Bound;
use super::blocks::storage_dimensions;
use super::key::checked_ordered_word;
use super::{RangeIndex, RangeIndexBounds, RangeIndexInconsistency, RangeIndexKey};
use crate::storage::warp_table::secondary_index::SlotSet;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RangeIndexLookupError<Key> {
OrderedWordOutsideWidth {
key: Key,
ordered_word: u64,
bits: u32,
},
ReversedBounds {
bounds: RangeIndexBounds<Key>,
},
Inconsistent {
reason: RangeIndexInconsistency,
},
}
#[doc(hidden)]
#[must_use]
pub struct RangeIndexMembership<Key> {
bounds: RangeIndexBounds<Key>,
slots: Option<SlotSet>,
}
impl<Key> RangeIndexMembership<Key> {
#[doc(hidden)]
pub fn into_parts(self) -> (RangeIndexBounds<Key>, Option<SlotSet>) {
(self.bounds, self.slots)
}
}
impl<Key: RangeIndexKey> RangeIndex<Key> {
pub fn find(
&self,
range: impl core::ops::RangeBounds<Key>,
) -> Result<RangeIndexMembership<Key>, RangeIndexLookupError<Key>> {
self.validate_storage()
.map_err(|reason| RangeIndexLookupError::Inconsistent { reason })?;
let bounds = RangeIndexBounds::from_range(range);
let start = bound_word(bounds.start())?;
let end = bound_word(bounds.end())?;
if matches!((start, end), (Some((left, _)), Some((right, _))) if left > right) {
return Err(RangeIndexLookupError::ReversedBounds { bounds });
}
let slots = self.find_between(start, end)?;
let slots = (!slots.is_empty()).then_some(slots);
Ok(RangeIndexMembership { bounds, slots })
}
#[allow(
clippy::arithmetic_side_effects,
reason = "validated fixed block dimensions bound every offset and allocation"
)]
fn find_between(
&self,
start: Option<(u64, bool)>,
end: Option<(u64, bool)>,
) -> Result<SlotSet, RangeIndexLookupError<Key>> {
let Some((words_per_block, _)) = storage_dimensions(self.maximum_rows, Key::BITS) else {
return Err(inconsistent(
RangeIndexInconsistency::StorageCapacityOverflow {
maximum_rows: self.maximum_rows,
bits: Key::BITS,
},
));
};
match RangePlan::new(start, end) {
RangePlan::Empty => Ok(SlotSet::empty()),
RangePlan::Exact { word } => {
collect_exact_blocks::<Key>(&self.blocks, words_per_block, word)
}
RangePlan::Split {
lower,
lower_inclusive,
upper,
upper_inclusive,
split_bit,
} => {
collect_matching_blocks::<Key>(&self.blocks, words_per_block, |members, planes| {
split_matches::<Key>(
members,
planes,
words_per_block,
lower,
lower_inclusive,
upper,
upper_inclusive,
split_bit,
)
})
}
RangePlan::Separate { start, end } => {
collect_matching_blocks::<Key>(&self.blocks, words_per_block, |members, planes| {
separate_matches::<Key>(members, planes, words_per_block, start, end)
})
}
}
}
}
#[derive(Clone, Copy)]
enum RangePlan {
Empty,
Exact {
word: u64,
},
Split {
lower: u64,
lower_inclusive: bool,
upper: u64,
upper_inclusive: bool,
split_bit: u32,
},
Separate {
start: Option<(u64, bool)>,
end: Option<(u64, bool)>,
},
}
impl RangePlan {
fn new(start: Option<(u64, bool)>, end: Option<(u64, bool)>) -> Self {
match (start, end) {
(Some((lower, lower_inclusive)), Some((upper, upper_inclusive))) if lower == upper => {
if lower_inclusive && upper_inclusive {
Self::Exact { word: lower }
} else {
Self::Empty
}
}
(Some((lower, lower_inclusive)), Some((upper, upper_inclusive))) => {
let differing = lower ^ upper;
let split_bit = differing.ilog2();
Self::Split {
lower,
lower_inclusive,
upper,
upper_inclusive,
split_bit,
}
}
_ => Self::Separate { start, end },
}
}
}
#[allow(
clippy::arithmetic_side_effects,
reason = "validated fixed block dimensions bound every allocation"
)]
fn collect_exact_blocks<Key: RangeIndexKey>(
blocks: &[u64],
words_per_block: usize,
word: u64,
) -> Result<SlotSet, RangeIndexLookupError<Key>> {
let block_count = blocks.len() / words_per_block;
let mut bytes = Vec::with_capacity(block_count * size_of::<u64>());
for block in blocks.chunks_exact(words_per_block) {
let Some((members, planes)) = block.split_first() else {
return Err(inconsistent(
RangeIndexInconsistency::BlockStorageLengthMismatch {
expected: words_per_block,
actual: block.len(),
},
));
};
let mut equal = *members;
for bit in (0..Key::BITS).rev() {
if equal == 0 {
break;
}
let ones = plane::<Key>(planes, words_per_block, bit)?;
let mask = bit_mask::<Key>(bit)?;
if word & mask == 0 {
equal &= !ones;
} else {
equal &= ones;
}
}
bytes.extend_from_slice(&equal.to_le_bytes());
}
Ok(SlotSet::from_lsb0_bytes(&bytes))
}
#[allow(
clippy::arithmetic_side_effects,
reason = "validated fixed block dimensions bound every allocation"
)]
fn collect_matching_blocks<Key>(
blocks: &[u64],
words_per_block: usize,
mut matching: impl FnMut(u64, &[u64]) -> Result<u64, RangeIndexLookupError<Key>>,
) -> Result<SlotSet, RangeIndexLookupError<Key>> {
let block_count = blocks.len() / words_per_block;
let mut bytes = Vec::with_capacity(block_count * size_of::<u64>());
for block in blocks.chunks_exact(words_per_block) {
let Some((members, planes)) = block.split_first() else {
return Err(inconsistent(
RangeIndexInconsistency::BlockStorageLengthMismatch {
expected: words_per_block,
actual: block.len(),
},
));
};
bytes.extend_from_slice(&matching(*members, planes)?.to_le_bytes());
}
Ok(SlotSet::from_lsb0_bytes(&bytes))
}
#[allow(clippy::too_many_arguments)]
fn split_matches<Key: RangeIndexKey>(
members: u64,
planes: &[u64],
words_per_block: usize,
lower: u64,
lower_inclusive: bool,
upper: u64,
upper_inclusive: bool,
split_bit: u32,
) -> Result<u64, RangeIndexLookupError<Key>> {
let mut common = members;
let prefix_start = split_bit.checked_add(1).ok_or_else(|| {
inconsistent(RangeIndexInconsistency::BitPositionOutsideWord { bit: split_bit })
})?;
for bit in (prefix_start..Key::BITS).rev() {
if common == 0 {
return Ok(0);
}
let ones = plane::<Key>(planes, words_per_block, bit)?;
let mask = bit_mask::<Key>(bit)?;
if lower & mask == 0 {
common &= !ones;
} else {
common &= ones;
}
}
let split_ones = plane::<Key>(planes, words_per_block, split_bit)?;
let mut lower_equal = common & !split_ones;
let mut upper_equal = common & split_ones;
let mut lower_above = 0_u64;
let mut upper_below = 0_u64;
for bit in (0..split_bit).rev() {
if lower_equal == 0 && upper_equal == 0 {
break;
}
let ones = plane::<Key>(planes, words_per_block, bit)?;
let mask = bit_mask::<Key>(bit)?;
if lower & mask == 0 {
lower_above |= lower_equal & ones;
lower_equal &= !ones;
} else {
lower_equal &= ones;
}
if upper & mask == 0 {
upper_equal &= !ones;
} else {
upper_below |= upper_equal & !ones;
upper_equal &= ones;
}
}
if !lower_inclusive {
lower_equal = 0;
}
if !upper_inclusive {
upper_equal = 0;
}
Ok(lower_above | lower_equal | upper_below | upper_equal)
}
fn separate_matches<Key: RangeIndexKey>(
members: u64,
planes: &[u64],
words_per_block: usize,
start: Option<(u64, bool)>,
end: Option<(u64, bool)>,
) -> Result<u64, RangeIndexLookupError<Key>> {
let (lower, lower_inclusive, mut lower_equal, mut lower_above) = match start {
Some((word, inclusive)) => (word, inclusive, members, 0_u64),
None => (0, false, 0, members),
};
let (upper, upper_inclusive, mut upper_equal, mut upper_below) = match end {
Some((word, inclusive)) => (word, inclusive, members, 0_u64),
None => (0, false, 0, members),
};
for bit in (0..Key::BITS).rev() {
if lower_equal == 0 && upper_equal == 0 {
break;
}
let ones = plane::<Key>(planes, words_per_block, bit)?;
let mask = bit_mask::<Key>(bit)?;
if lower & mask == 0 {
lower_above |= lower_equal & ones;
lower_equal &= !ones;
} else {
lower_equal &= ones;
}
if upper & mask == 0 {
upper_equal &= !ones;
} else {
upper_below |= upper_equal & !ones;
upper_equal &= ones;
}
}
if !lower_inclusive {
lower_equal = 0;
}
if !upper_inclusive {
upper_equal = 0;
}
Ok((lower_above | lower_equal) & (upper_below | upper_equal))
}
fn plane<Key>(
planes: &[u64],
words_per_block: usize,
bit: u32,
) -> Result<u64, RangeIndexLookupError<Key>> {
planes.get(bit as usize).copied().ok_or_else(|| {
inconsistent(RangeIndexInconsistency::BlockStorageLengthMismatch {
expected: words_per_block,
actual: planes.len().saturating_add(1),
})
})
}
fn bit_mask<Key>(bit: u32) -> Result<u64, RangeIndexLookupError<Key>> {
1_u64
.checked_shl(bit)
.ok_or_else(|| inconsistent(RangeIndexInconsistency::BitPositionOutsideWord { bit }))
}
fn bound_word<Key: RangeIndexKey>(
bound: Bound<&Key>,
) -> Result<Option<(u64, bool)>, RangeIndexLookupError<Key>> {
match bound {
Bound::Included(key) => checked_word(*key).map(|word| Some((word, true))),
Bound::Excluded(key) => checked_word(*key).map(|word| Some((word, false))),
Bound::Unbounded => Ok(None),
}
}
fn checked_word<Key: RangeIndexKey>(key: Key) -> Result<u64, RangeIndexLookupError<Key>> {
checked_ordered_word(key).map_err(|outside| RangeIndexLookupError::OrderedWordOutsideWidth {
key: outside.key,
ordered_word: outside.ordered_word,
bits: outside.bits,
})
}
fn inconsistent<Key>(reason: RangeIndexInconsistency) -> RangeIndexLookupError<Key> {
RangeIndexLookupError::Inconsistent { reason }
}