use alloc::vec::Vec;
use core::cmp::Ordering;
use super::{
OrderedCandidates, OrderedRowLimit, OrderedSelection, OrderedSelectionComparingRows,
OrderedSelectionFirstCuttingTiesOutcome, OrderedSelectionFirstKeepingTiesOutcome,
};
use crate::storage::warp_table::inconsistency::invariant_violation;
use crate::storage::warp_table::query::rows::resolve;
use crate::storage::warp_table::query::selection::SelectionMembership;
use crate::storage::warp_table::secondary_index::{
RangeIndexOrderingError, RangeIndexRanking, RangeIndexRankingCandidates, RankedSlots, SlotSet,
StableSlot, StableSlotResolver, StraddlingRangeIndexTies, rank,
};
use crate::storage::warp_table::{WarpTable, WarpTableRow};
enum RankedRows<RowRef> {
Exact {
rows: Vec<RowRef>,
},
Straddling {
leaders: Vec<RowRef>,
ties: StraddlingRangeIndexTies,
},
}
enum EqualRowCollection {
FirstNeeded,
Every,
}
struct BrokenTies<RowRef> {
better: Vec<RowRef>,
equal: Vec<RowRef>,
equal_count: usize,
needed_from_equal: usize,
}
struct DirectCandidate {
slot: StableSlot,
sortable_word: u64,
}
enum DirectRankingInconsistency {
PivotUnavailable {
pivot: usize,
candidates: usize,
},
BoundaryUnavailable {
limit: usize,
leaders: usize,
tied: usize,
},
}
fn first_row<RowRef>(outcome: OrderedSelectionFirstCuttingTiesOutcome<RowRef>) -> Option<RowRef> {
match outcome {
OrderedSelectionFirstCuttingTiesOutcome::Exact { rows }
| OrderedSelectionFirstCuttingTiesOutcome::TiesCut { rows, .. } => rows.into_iter().next(),
}
}
impl core::fmt::Debug for DirectRankingInconsistency {
fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::PivotUnavailable { pivot, candidates } => formatter
.debug_struct("PivotUnavailable")
.field("pivot", pivot)
.field("candidates", candidates)
.finish(),
Self::BoundaryUnavailable {
limit,
leaders,
tied,
} => formatter
.debug_struct("BoundaryUnavailable")
.field("limit", limit)
.field("leaders", leaders)
.field("tied", tied)
.finish(),
}
}
}
const MAXIMUM_DIRECT_CANDIDATES_PER_BLOCK: usize = 2;
const MAXIMUM_DIRECT_CANDIDATES: usize = 1_024;
impl<'table, Row, Brand> OrderedSelection<'table, Row, Brand>
where
Row: WarpTableRow,
Row::Indexes: StableSlotResolver,
{
pub fn first(self) -> Result<Option<Row::Ref<'table>>, RangeIndexOrderingError> {
if self.ordering.later_levels.is_empty()
&& let OrderedCandidates::EveryRow { row_count } = self.candidates
{
return self
.ordering
.first_level
.first_member(row_count)
.map(|slot| slot.map(|slot| resolve(self.table, slot)));
}
self.first_cutting_ties(OrderedRowLimit::new(core::num::NonZeroU32::MIN))
.map(first_row)
}
pub fn first_cutting_ties(
self,
limit: OrderedRowLimit,
) -> Result<OrderedSelectionFirstCuttingTiesOutcome<Row::Ref<'table>>, RangeIndexOrderingError>
{
match self.ranked_rows(limit, &mut keep_slot_order)? {
RankedRows::Exact { rows } => {
Ok(OrderedSelectionFirstCuttingTiesOutcome::Exact { rows })
}
RankedRows::Straddling { mut leaders, ties } => {
let needed = ties.needed().get();
leaders.extend(
ties.tied()
.slots()
.take(needed)
.map(|slot| resolve(self.table, slot)),
);
Ok(OrderedSelectionFirstCuttingTiesOutcome::TiesCut {
rows: leaders,
tied_rows_left_out: ties.tied().len().saturating_sub(needed),
})
}
}
}
pub fn first_keeping_ties(
self,
limit: OrderedRowLimit,
) -> Result<OrderedSelectionFirstKeepingTiesOutcome<Row::Ref<'table>>, RangeIndexOrderingError>
{
match self.ranked_rows(limit, &mut keep_slot_order)? {
RankedRows::Exact { rows } => {
Ok(OrderedSelectionFirstKeepingTiesOutcome::Exact { rows })
}
RankedRows::Straddling { leaders, ties } => {
Ok(OrderedSelectionFirstKeepingTiesOutcome::TiesKept {
leaders,
tied: ties
.tied()
.slots()
.map(|slot| resolve(self.table, slot))
.collect(),
})
}
}
}
fn ranked_rows<Compare>(
&self,
limit: OrderedRowLimit,
compare: &mut Compare,
) -> Result<RankedRows<Row::Ref<'table>>, RangeIndexOrderingError>
where
Compare: FnMut(&Row::Ref<'table>, &Row::Ref<'table>) -> Ordering,
{
let limit = limit.row_count();
let ranking = match &self.candidates {
OrderedCandidates::EveryRow { row_count } => rank(
&self.ordering.first_level,
&self.ordering.later_levels,
RangeIndexRankingCandidates::EveryMember {
expected_members: *row_count,
},
limit,
)?,
OrderedCandidates::Membership(SelectionMembership::Empty) => {
return Ok(RankedRows::Exact { rows: Vec::new() });
}
OrderedCandidates::Membership(SelectionMembership::Borrowed(slots)) => {
if prefers_direct_ranking(self.ordering.level_count().get(), slots) {
return self.rank_directly(slots, limit, compare);
}
self.rank_membership(slots, limit)?
}
OrderedCandidates::Membership(SelectionMembership::Owned(slots)) => {
if prefers_direct_ranking(self.ordering.level_count().get(), slots) {
return self.rank_directly(slots, limit, compare);
}
self.rank_membership(slots, limit)?
}
};
match ranking {
RangeIndexRanking::Exact { winners } => Ok(RankedRows::Exact {
rows: self.sorted_rows(&winners, compare)?,
}),
RangeIndexRanking::Straddling { leaders, ties } => Ok(RankedRows::Straddling {
leaders: self.sorted_rows(&leaders, compare)?,
ties,
}),
}
}
fn rank_membership(
&self,
slots: &SlotSet,
limit: core::num::NonZeroUsize,
) -> Result<RangeIndexRanking, RangeIndexOrderingError> {
rank(
&self.ordering.first_level,
&self.ordering.later_levels,
RangeIndexRankingCandidates::Slots(slots),
limit,
)
}
fn rank_directly<Compare>(
&self,
slots: &SlotSet,
limit: core::num::NonZeroUsize,
compare: &mut Compare,
) -> Result<RankedRows<Row::Ref<'table>>, RangeIndexOrderingError>
where
Compare: FnMut(&Row::Ref<'table>, &Row::Ref<'table>) -> Ordering,
{
let candidate_count = usize::try_from(slots.len()).unwrap_or(usize::MAX);
if candidate_count <= limit.get() {
let mut rows = Vec::with_capacity(candidate_count);
for slot in slots.slots() {
self.ordering.first_level.validate_slot(slot)?;
rows.push(resolve(self.table, slot));
}
self.sort_direct_rows(&mut rows, compare);
return Ok(RankedRows::Exact { rows });
}
let sortable_word = self.ordering.first_sortable_word;
let mut candidates = Vec::with_capacity(candidate_count);
for slot in slots.slots() {
self.ordering.first_level.validate_slot(slot)?;
let row = resolve(self.table, slot);
candidates.push(DirectCandidate {
slot,
sortable_word: sortable_word(&row),
});
}
let pivot_index = limit.get().saturating_sub(1);
candidates.select_nth_unstable_by_key(pivot_index, |candidate| candidate.sortable_word);
let pivot = candidates.get(pivot_index).unwrap_or_else(|| {
invariant_violation(DirectRankingInconsistency::PivotUnavailable {
pivot: pivot_index,
candidates: candidates.len(),
})
});
let pivot_word = pivot.sortable_word;
let leader_count = candidates
.iter()
.filter(|candidate| candidate.sortable_word < pivot_word)
.count();
let tied_count = candidates
.iter()
.filter(|candidate| candidate.sortable_word == pivot_word)
.count();
let needed = core::num::NonZeroUsize::new(limit.get().saturating_sub(leader_count))
.filter(|needed| needed.get() <= tied_count)
.unwrap_or_else(|| {
invariant_violation(DirectRankingInconsistency::BoundaryUnavailable {
limit: limit.get(),
leaders: leader_count,
tied: tied_count,
})
});
if tied_count == needed.get() {
candidates.retain(|candidate| candidate.sortable_word <= pivot_word);
candidates.sort_unstable_by_key(|candidate| (candidate.sortable_word, candidate.slot));
let mut rows = candidates
.into_iter()
.map(|candidate| resolve(self.table, candidate.slot))
.collect::<Vec<_>>();
self.sort_direct_rows(&mut rows, compare);
return Ok(RankedRows::Exact { rows });
}
let tied = candidates
.iter()
.filter(|candidate| candidate.sortable_word == pivot_word)
.map(|candidate| candidate.slot)
.collect();
candidates.retain(|candidate| candidate.sortable_word < pivot_word);
candidates.sort_unstable_by_key(|candidate| (candidate.sortable_word, candidate.slot));
let mut leaders = candidates
.into_iter()
.map(|candidate| resolve(self.table, candidate.slot))
.collect::<Vec<_>>();
self.sort_direct_rows(&mut leaders, compare);
Ok(RankedRows::Straddling {
leaders,
ties: StraddlingRangeIndexTies::new(RankedSlots::from_slots(tied), needed),
})
}
fn sort_direct_rows<Compare>(&self, rows: &mut [Row::Ref<'table>], compare: &mut Compare)
where
Compare: FnMut(&Row::Ref<'table>, &Row::Ref<'table>) -> Ordering,
{
let sortable_word = self.ordering.first_sortable_word;
rows.sort_by(|left, right| {
sortable_word(left)
.cmp(&sortable_word(right))
.then_with(|| compare(left, right))
});
}
fn sorted_rows<Compare>(
&self,
slots: &RankedSlots,
compare: &mut Compare,
) -> Result<Vec<Row::Ref<'table>>, RangeIndexOrderingError>
where
Compare: FnMut(&Row::Ref<'table>, &Row::Ref<'table>) -> Ordering,
{
let level_count = self.ordering.level_count().get();
let mut sortable_words = Vec::with_capacity(slots.len().saturating_mul(level_count));
let mut resolved_rows = Vec::with_capacity(slots.len());
for slot in slots.slots() {
for (level, _) in self.ordering.levels() {
level.validate_slot(slot)?;
}
let row = resolve(self.table, slot);
for (_, sortable_word) in self.ordering.levels() {
sortable_words.push(sortable_word(&row));
}
resolved_rows.push(row);
}
let mut rows = sortable_words
.chunks_exact(level_count)
.zip(resolved_rows)
.collect::<Vec<_>>();
rows.sort_by(|(left_key, left_row), (right_key, right_row)| {
left_key
.cmp(right_key)
.then_with(|| compare(left_row, right_row))
});
Ok(rows.into_iter().map(|(_, row)| row).collect())
}
}
impl<'table, Row, Brand, Compare> OrderedSelectionComparingRows<'table, Row, Brand, Compare>
where
Row: WarpTableRow,
Row::Indexes: StableSlotResolver,
Compare: FnMut(&Row::Ref<'table>, &Row::Ref<'table>) -> Ordering,
{
pub fn first(self) -> Result<Option<Row::Ref<'table>>, RangeIndexOrderingError> {
self.first_cutting_ties(OrderedRowLimit::new(core::num::NonZeroU32::MIN))
.map(first_row)
}
pub fn first_cutting_ties(
mut self,
limit: OrderedRowLimit,
) -> Result<OrderedSelectionFirstCuttingTiesOutcome<Row::Ref<'table>>, RangeIndexOrderingError>
{
match self.ordered.ranked_rows(limit, &mut self.compare)? {
RankedRows::Exact { rows } => {
Ok(OrderedSelectionFirstCuttingTiesOutcome::Exact { rows })
}
RankedRows::Straddling { mut leaders, ties } => {
let broken = break_ties(
self.ordered.table,
&ties,
&mut self.compare,
&EqualRowCollection::FirstNeeded,
);
let tied_rows_left_out =
broken.equal_count.saturating_sub(broken.needed_from_equal);
leaders.extend(broken.better);
leaders.extend(broken.equal);
if tied_rows_left_out == 0 {
return Ok(OrderedSelectionFirstCuttingTiesOutcome::Exact { rows: leaders });
}
Ok(OrderedSelectionFirstCuttingTiesOutcome::TiesCut {
rows: leaders,
tied_rows_left_out,
})
}
}
}
pub fn first_keeping_ties(
mut self,
limit: OrderedRowLimit,
) -> Result<OrderedSelectionFirstKeepingTiesOutcome<Row::Ref<'table>>, RangeIndexOrderingError>
{
match self.ordered.ranked_rows(limit, &mut self.compare)? {
RankedRows::Exact { rows } => {
Ok(OrderedSelectionFirstKeepingTiesOutcome::Exact { rows })
}
RankedRows::Straddling { mut leaders, ties } => {
let broken = break_ties(
self.ordered.table,
&ties,
&mut self.compare,
&EqualRowCollection::Every,
);
leaders.extend(broken.better);
if broken.equal_count == broken.needed_from_equal {
leaders.extend(broken.equal);
return Ok(OrderedSelectionFirstKeepingTiesOutcome::Exact { rows: leaders });
}
Ok(OrderedSelectionFirstKeepingTiesOutcome::TiesKept {
leaders,
tied: broken.equal,
})
}
}
}
}
fn break_ties<'table, Row, Brand, Compare>(
table: &'table WarpTable<Row, Brand>,
ties: &StraddlingRangeIndexTies,
compare: &mut Compare,
collection: &EqualRowCollection,
) -> BrokenTies<Row::Ref<'table>>
where
Row: WarpTableRow,
Row::Indexes: StableSlotResolver,
Compare: FnMut(&Row::Ref<'table>, &Row::Ref<'table>) -> Ordering,
{
let needed = ties.needed().get();
let pivot_index = needed.saturating_sub(1);
let buffer_limit = needed.saturating_mul(2);
let mut best = Vec::with_capacity(buffer_limit);
for slot in ties.tied().slots() {
best.push(resolve(table, slot));
if best.len() == buffer_limit {
best.select_nth_unstable_by(pivot_index, &mut *compare);
best.truncate(needed);
}
}
let (better, pivot, _) = best.select_nth_unstable_by(pivot_index, &mut *compare);
better.sort_by(&mut *compare);
let better_count = better.partition_point(|row| compare(row, pivot) == Ordering::Less);
let needed_from_equal = needed.saturating_sub(better_count);
let mut equal = Vec::new();
let mut equal_count = 0_usize;
for slot in ties.tied().slots() {
let row = resolve(table, slot);
if compare(&row, pivot) != Ordering::Equal {
continue;
}
equal_count = equal_count.saturating_add(1);
let collecting = match collection {
EqualRowCollection::Every => true,
EqualRowCollection::FirstNeeded => equal.len() < needed_from_equal,
};
if collecting {
equal.push(row);
}
}
best.truncate(better_count);
BrokenTies {
better: best,
equal,
equal_count,
needed_from_equal,
}
}
fn keep_slot_order<RowRef>(_: &RowRef, _: &RowRef) -> Ordering {
Ordering::Equal
}
pub(super) fn prefers_direct_ranking(level_count: usize, slots: &SlotSet) -> bool {
let candidate_count = usize::try_from(slots.len()).unwrap_or(usize::MAX);
if level_count != 1 || candidate_count > MAXIMUM_DIRECT_CANDIDATES {
return false;
}
let mut previous_block = None;
let mut occupied_blocks = 0_usize;
for slot in slots.slots() {
let block = slot.value() / u64::BITS;
if previous_block != Some(block) {
occupied_blocks = occupied_blocks.saturating_add(1);
previous_block = Some(block);
}
}
candidate_count <= occupied_blocks.saturating_mul(MAXIMUM_DIRECT_CANDIDATES_PER_BLOCK)
}