#[cfg(test)]
use alloc::vec::Vec;
use core::num::NonZeroUsize;
use super::ranked_slots::RankedBlock;
#[cfg(test)]
use super::ranked_slots::RankedSlots;
use super::{
RangeIndexInconsistency, RangeIndexOrderingDirection, RangeIndexOrderingError,
RangeIndexOrderingLevel, RangeIndexPlanes, SLOT_OFFSET_MASK,
};
use crate::storage::query_workspace::{
WorkspaceMembership,
chunk::read::{ChunkRead, ReadPayload},
};
use crate::storage::warp_table::inconsistency::invariant_violation;
#[cfg(feature = "secondary-indexes")]
use crate::storage::warp_table::secondary_index::SlotSetRead;
#[cfg(test)]
mod tests;
mod workspace;
pub use workspace::RankingWorkspaceError;
pub(in crate::storage::warp_table) use workspace::{
WorkspaceRanking, rank_with, ranking_storage_bytes,
};
pub(in crate::storage::warp_table) enum RangeIndexRankingCandidates<'slots> {
#[cfg(feature = "secondary-indexes")]
EveryMember {
expected_members: usize,
},
#[cfg(feature = "secondary-indexes")]
Slots(SlotSetRead<'slots>),
Workspace(&'slots WorkspaceMembership<'slots>),
}
#[cfg(test)]
pub(in crate::storage::warp_table) enum RangeIndexRanking {
Exact {
winners: RankedSlots,
},
Straddling {
leaders: RankedSlots,
ties: StraddlingRangeIndexTies,
},
}
#[cfg(test)]
pub(in crate::storage::warp_table) struct StraddlingRangeIndexTies {
tied: RankedSlots,
needed: NonZeroUsize,
}
#[cfg(test)]
impl StraddlingRangeIndexTies {
pub(in crate::storage::warp_table) const fn tied(&self) -> &RankedSlots {
&self.tied
}
pub(in crate::storage::warp_table) const fn needed(&self) -> NonZeroUsize {
self.needed
}
}
#[derive(Clone, Copy)]
struct CandidateBlock<'index> {
slot_base: u32,
boundary: u64,
leaders: u64,
planes: &'index [u64],
}
impl CandidateBlock<'_> {
const fn new(slot_base: u32, boundary: u64) -> Self {
Self {
slot_base,
boundary,
leaders: 0,
planes: &[],
}
}
}
enum PreferredRows {
Promote { count: usize },
Narrow,
Unchanged,
}
#[cfg(test)]
pub(in crate::storage::warp_table) fn rank<'index: 'level, 'level>(
first_level: &'level RangeIndexOrderingLevel<'index>,
later_levels: impl IntoIterator<Item = &'level RangeIndexOrderingLevel<'index>>,
candidates: RangeIndexRankingCandidates<'_>,
limit: NonZeroUsize,
) -> Result<RangeIndexRanking, RangeIndexOrderingError> {
let first_planes = first_level.planes()?;
let mut blocks = match candidates {
RangeIndexRankingCandidates::EveryMember { expected_members } => {
every_member_blocks(first_planes, expected_members)?
}
RangeIndexRankingCandidates::Slots(slots) => slot_blocks(slots),
RangeIndexRankingCandidates::Workspace(membership) => {
let mut blocks = Vec::with_capacity(workspace_block_count(membership)?);
visit_workspace_blocks(membership, |block| blocks.push(block))?;
blocks
}
};
attach_planes(first_planes, &mut blocks)?;
if slot_count(blocks.iter().map(|block| block.boundary)) <= limit.get() {
return Ok(classify(Vec::new(), &blocks, 0, limit));
}
let mut settled = Vec::with_capacity(blocks.len().min(limit.get()));
let (active, leader_count) =
refine_candidates(first_level, later_levels, &mut blocks, limit, |block| {
settled.push(block)
})?;
blocks.truncate(active);
Ok(classify(settled, &blocks, leader_count, limit))
}
fn refine_candidates<'index: 'level, 'level>(
first_level: &'level RangeIndexOrderingLevel<'index>,
later_levels: impl IntoIterator<Item = &'level RangeIndexOrderingLevel<'index>>,
blocks: &mut [CandidateBlock<'index>],
limit: NonZeroUsize,
mut settled: impl FnMut(RankedBlock),
) -> Result<(usize, usize), RangeIndexOrderingError> {
let mut active = blocks.len();
let mut leader_count = 0_usize;
'levels: for (level_index, level) in core::iter::once(first_level)
.chain(later_levels)
.enumerate()
{
let planes = level.planes()?;
let blocks = blocks
.get_mut(..active)
.unwrap_or_else(|| invariant_violation("ranking active prefix exceeded storage"));
if level_index != 0 {
attach_planes(planes, blocks)?;
}
let direction = level.direction;
for bit in (0..distinguishing_bits(blocks)).rev() {
let blocks = blocks
.get_mut(..active)
.unwrap_or_else(|| invariant_violation("ranking active prefix exceeded storage"));
let remaining = limit.get().saturating_sub(leader_count);
let preferred_rows = preferred_partition(planes, blocks, bit, direction, remaining)?;
if let PreferredRows::Promote { count } = preferred_rows {
leader_count = leader_count.saturating_add(count);
}
active = settle(
planes,
blocks,
bit,
direction,
&mut settled,
&preferred_rows,
)?;
if leader_count == limit.get() {
break 'levels;
}
}
}
Ok((active, leader_count))
}
fn distinguishing_bits(blocks: &[CandidateBlock<'_>]) -> usize {
let Some((first, remaining)) = blocks.split_first() else {
return 0;
};
let mut varying_bits = first
.planes
.iter()
.rposition(|plane| {
let ones = plane & first.boundary;
ones != 0 && ones != first.boundary
})
.map_or(0, |bit| bit.saturating_add(1));
for block in remaining {
if varying_bits == first.planes.len() {
break;
}
let prefixes = first.planes.iter().zip(block.planes).skip(varying_bits);
let differing_rows = prefixes.clone().fold(0, |differing, (reference, plane)| {
differing
| if reference & first.boundary == 0 {
*plane
} else {
!plane
}
});
if differing_rows & block.boundary != 0 {
for (bit, (reference, plane)) in prefixes.enumerate().rev() {
let expected = if reference & first.boundary == 0 {
0
} else {
block.boundary
};
if plane & block.boundary != expected {
varying_bits = varying_bits.saturating_add(bit).saturating_add(1);
break;
}
}
}
}
varying_bits
}
#[cfg(test)]
fn every_member_blocks<'index>(
planes: &RangeIndexPlanes<'index>,
expected_members: usize,
) -> Result<Vec<CandidateBlock<'index>>, RangeIndexOrderingError> {
let maximum_blocks = planes.blocks.chunks_exact(planes.words_per_block).len();
let mut blocks = Vec::with_capacity(maximum_blocks.min(expected_members));
visit_every_member(planes, expected_members, |block| blocks.push(block))?;
Ok(blocks)
}
#[cfg(feature = "secondary-indexes")]
fn visit_every_member<'index>(
planes: &RangeIndexPlanes<'index>,
expected_members: usize,
mut emit: impl FnMut(CandidateBlock<'index>),
) -> Result<(), RangeIndexOrderingError> {
let mut index_members = 0_usize;
for (block_index, members) in planes.members().enumerate() {
if members != 0 {
emit(CandidateBlock::new(planes.slot_base(block_index)?, members));
index_members = index_members.saturating_add(members.count_ones() as usize);
}
}
if index_members != expected_members {
return Err(RangeIndexOrderingError::MemberCountMismatch {
field: planes.field,
table_rows: expected_members,
index_members,
});
}
Ok(())
}
#[cfg(test)]
fn slot_blocks<'index>(slots: SlotSetRead<'_>) -> Vec<CandidateBlock<'index>> {
let mut blocks = Vec::with_capacity(slot_block_count(slots));
visit_slot_blocks(slots, |block| blocks.push(block));
blocks
}
#[cfg(feature = "secondary-indexes")]
fn slot_block_count(slots: SlotSetRead<'_>) -> usize {
let word_capacity = slot_word_capacity(slots);
slots.chunks().fold(0_usize, |count, (high, chunk)| {
count.saturating_add(chunk_block_count(
chunk,
chunk_word_count(high, word_capacity),
))
})
}
fn workspace_block_count(
membership: &WorkspaceMembership<'_>,
) -> Result<usize, RangeIndexOrderingError> {
membership.read_chunks().try_fold(0_usize, |count, chunk| {
let (high, chunk) = chunk.map_err(RangeIndexOrderingError::Membership)?;
let word_count = chunk_word_count(high, membership.capacity().div_ceil(64));
Ok(count.saturating_add(chunk_block_count(chunk, word_count)))
})
}
#[cfg(feature = "secondary-indexes")]
fn slot_word_capacity(slots: SlotSetRead<'_>) -> usize {
match slots {
SlotSetRead::Workspace { capacity, .. } => capacity.div_ceil(64),
SlotSetRead::Empty | SlotSetRead::Owned(_) | SlotSetRead::Indexed(_) => usize::MAX,
}
}
fn chunk_word_count(high: u16, word_capacity: usize) -> usize {
word_capacity
.saturating_sub(usize::from(high) << 10)
.min(1024)
}
fn chunk_block_count(chunk: ChunkRead<'_>, word_count: usize) -> usize {
match chunk.payload {
ReadPayload::Sparse(values) => values
.chunk_by(|left, right| left >> 6 == right >> 6)
.count(),
ReadPayload::Dense(words) => words
.iter()
.take(word_count)
.filter(|word| **word != 0)
.count(),
}
}
#[cfg(feature = "secondary-indexes")]
fn visit_slot_blocks<'index>(slots: SlotSetRead<'_>, mut emit: impl FnMut(CandidateBlock<'index>)) {
let word_capacity = slot_word_capacity(slots);
for (high, chunk) in slots.chunks() {
visit_chunk_blocks(
high,
chunk,
chunk_word_count(high, word_capacity),
&mut emit,
);
}
}
fn visit_workspace_blocks<'index>(
membership: &WorkspaceMembership<'_>,
mut emit: impl FnMut(CandidateBlock<'index>),
) -> Result<(), RangeIndexOrderingError> {
for chunk in membership.read_chunks() {
let (high, chunk) = chunk.map_err(RangeIndexOrderingError::Membership)?;
let word_count = chunk_word_count(high, membership.capacity().div_ceil(64));
visit_chunk_blocks(high, chunk, word_count, &mut emit);
}
Ok(())
}
fn visit_chunk_blocks<'index>(
high: u16,
chunk: ChunkRead<'_>,
word_count: usize,
emit: &mut impl FnMut(CandidateBlock<'index>),
) {
let base = u32::from(high) << 16;
match chunk.payload {
ReadPayload::Sparse(values) => {
for group in values.chunk_by(|left, right| left >> 6 == right >> 6) {
let first = group
.first()
.unwrap_or_else(|| invariant_violation("empty group"));
let mask = group.iter().fold(0_u64, |mask, low| {
mask | 1_u64.wrapping_shl(u32::from(*low) & SLOT_OFFSET_MASK)
});
emit(CandidateBlock::new(
base | (u32::from(*first) & !SLOT_OFFSET_MASK),
mask,
));
}
}
ReadPayload::Dense(words) => {
for (ordinal, word) in words.iter().take(word_count).enumerate() {
if *word != 0 {
emit(CandidateBlock::new(base | ((ordinal as u32) << 6), *word));
}
}
}
}
}
fn attach_planes<'index>(
planes: &RangeIndexPlanes<'index>,
blocks: &mut [CandidateBlock<'index>],
) -> Result<(), RangeIndexOrderingError> {
for block in blocks {
let range_block = planes.block(block.slot_base)?;
if block.boundary & !range_block.members != 0 {
return Err(planes.inconsistent(RangeIndexInconsistency::MemberAbsent));
}
block.planes = range_block.planes;
}
Ok(())
}
fn preferred_partition(
planes: &RangeIndexPlanes<'_>,
blocks: &[CandidateBlock<'_>],
bit: usize,
direction: RangeIndexOrderingDirection,
remaining_limit: usize,
) -> Result<PreferredRows, RangeIndexOrderingError> {
let mut preferred_count = 0_usize;
let mut removed = 0_u64;
let mut remaining = blocks.iter();
while let Some(block) = remaining.next() {
let plane = planes.plane_at(block.planes, bit)?;
let preferred = block.boundary & direction.preferred_rows(plane);
removed |= block.boundary ^ preferred;
let count = preferred.count_ones() as usize;
if count > remaining_limit.saturating_sub(preferred_count) {
while removed == 0 {
let Some(block) = remaining.next() else {
return Ok(PreferredRows::Unchanged);
};
let plane = planes.plane_at(block.planes, bit)?;
removed = block.boundary & !direction.preferred_rows(plane);
}
return Ok(PreferredRows::Narrow);
}
preferred_count = preferred_count.saturating_add(count);
}
Ok(if preferred_count == 0 {
PreferredRows::Unchanged
} else {
PreferredRows::Promote {
count: preferred_count,
}
})
}
fn settle(
planes: &RangeIndexPlanes<'_>,
blocks: &mut [CandidateBlock<'_>],
bit: usize,
direction: RangeIndexOrderingDirection,
settled: &mut impl FnMut(RankedBlock),
preferred_rows: &PreferredRows,
) -> Result<usize, RangeIndexOrderingError> {
if let PreferredRows::Unchanged = preferred_rows {
return Ok(blocks.len());
}
let mut retained = 0_usize;
for position in 0..blocks.len() {
let block = blocks
.get_mut(position)
.unwrap_or_else(|| invariant_violation("ranking block missing"));
let plane = planes.plane_at(block.planes, bit)?;
let preferred = block.boundary & direction.preferred_rows(plane);
if let PreferredRows::Promote { .. } = preferred_rows {
block.leaders |= preferred;
block.boundary &= !preferred;
} else {
block.boundary = preferred;
}
if block.boundary != 0 {
if position != retained {
let kept = *block;
*blocks
.get_mut(retained)
.unwrap_or_else(|| invariant_violation("ranking retained block missing")) =
kept;
}
retained = retained.saturating_add(1);
} else if block.leaders != 0 {
settled(RankedBlock::new(block.slot_base, block.leaders));
}
}
Ok(retained)
}
#[cfg(test)]
fn classify(
settled: Vec<RankedBlock>,
blocks: &[CandidateBlock<'_>],
leader_count: usize,
limit: NonZeroUsize,
) -> RangeIndexRanking {
let tied_count = slot_count(blocks.iter().map(|block| block.boundary));
let leaders_with = |word_of: fn(&CandidateBlock<'_>) -> u64| {
let mut leaders = settled;
let additional = blocks.iter().filter(|block| word_of(block) != 0).count();
leaders.reserve(additional);
leaders.extend(blocks.iter().filter_map(|block| {
let slots = word_of(block);
(slots != 0).then_some(RankedBlock::new(block.slot_base, slots))
}));
RankedSlots::new(leaders)
};
match NonZeroUsize::new(limit.get().saturating_sub(leader_count)) {
None => RangeIndexRanking::Exact {
winners: leaders_with(|block| block.leaders),
},
Some(needed) if tied_count > needed.get() => RangeIndexRanking::Straddling {
leaders: leaders_with(|block| block.leaders),
ties: StraddlingRangeIndexTies {
tied: RankedSlots::new(
blocks
.iter()
.map(|block| RankedBlock::new(block.slot_base, block.boundary))
.collect(),
),
needed,
},
},
Some(_) => RangeIndexRanking::Exact {
winners: leaders_with(|block| block.leaders | block.boundary),
},
}
}
fn slot_count(words: impl Iterator<Item = u64>) -> usize {
words.fold(0_usize, |count, word| {
count.saturating_add(word.count_ones() as usize)
})
}