use core::{iter::FusedIterator, slice};
use super::{RangeIndexLookupOutcome, SecondaryIndexLookupOutcome, Selection};
use crate::storage::warp_table::inconsistency::invariant_violation;
use crate::storage::warp_table::secondary_index::{
StableSlot, StableSlotResolveOutcome, StableSlotResolver,
};
use crate::storage::warp_table::{
WarpTable, WarpTableGetAtOutcome, WarpTableInconsistency, WarpTableRow,
};
#[derive(
::core::clone::Clone,
::core::marker::Copy,
::core::fmt::Debug,
::core::cmp::Eq,
::core::cmp::PartialEq,
)]
pub struct DenseSelectionStorageTooSmall {
pub required: u64,
pub available: usize,
}
#[must_use]
pub struct DenseSelectionRows<'table, 'workspace, Row: WarpTableRow, Brand = Row> {
table: &'table WarpTable<Row, Brand>,
rows: slice::Iter<'workspace, usize>,
}
pub(super) fn dense_rows<'workspace, Row: WarpTableRow, Brand>(
table: &WarpTable<Row, Brand>,
count: u64,
slots: impl Iterator<Item = StableSlot>,
storage: &'workspace mut [usize],
) -> Result<&'workspace [usize], DenseSelectionStorageTooSmall>
where
Row::Indexes: StableSlotResolver,
{
let refusal = DenseSelectionStorageTooSmall {
required: count,
available: storage.len(),
};
let count = usize::try_from(count).map_err(|_| refusal)?;
let rows = storage.get_mut(..count).ok_or(refusal)?;
let mut filled = 0_usize;
for slot in slots {
let row = rows.get_mut(filled).unwrap_or_else(|| {
invariant_violation("selection contains more rows than its cardinality")
});
let StableSlotResolveOutcome::Resolved { dense_row, .. } = table
.__indexes()
.resolve_stable_slot(slot)
.unwrap_or_else(|reason| invariant_violation(reason));
*row = dense_row;
filled = filled.saturating_add(1);
}
if filled != count {
invariant_violation("selection contains fewer rows than its cardinality");
}
rows.sort_unstable();
Ok(rows)
}
impl<'table, 'workspace, Row: WarpTableRow, Brand>
DenseSelectionRows<'table, 'workspace, Row, Brand>
{
pub(super) fn new(table: &'table WarpTable<Row, Brand>, rows: &'workspace [usize]) -> Self {
Self {
table,
rows: rows.iter(),
}
}
fn row(&self, dense_row: usize) -> Row::Ref<'table> {
match self.table.get_at(dense_row) {
WarpTableGetAtOutcome::Present { row, .. } => row,
WarpTableGetAtOutcome::Absent { dense_row } => {
invariant_violation(WarpTableInconsistency::DenseRowUnavailable {
row: dense_row,
rows: self.table.len(),
})
}
}
}
}
impl<'table, Row: WarpTableRow, Brand> Iterator for DenseSelectionRows<'table, '_, Row, Brand> {
type Item = Row::Ref<'table>;
fn next(&mut self) -> Option<Self::Item> {
let row = *self.rows.next()?;
Some(self.row(row))
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.rows.size_hint()
}
fn fold<Acc, Fold: FnMut(Acc, Self::Item) -> Acc>(
self,
mut initial: Acc,
mut fold: Fold,
) -> Acc {
for run in self
.rows
.as_slice()
.chunk_by(|left, right| left.checked_add(1) == Some(*right))
{
let (first, last) = run
.first()
.zip(run.last())
.unwrap_or_else(|| invariant_violation("chunk_by emitted an empty run"));
let ids = self.table.ids().get(*first..=*last).unwrap_or_else(|| {
invariant_violation(WarpTableInconsistency::DenseRowUnavailable {
row: *last,
rows: self.table.len(),
})
});
initial = Row::fold_rows(ids, self.table.columns(), *first, initial, &mut fold)
.unwrap_or_else(|reason| invariant_violation(WarpTableInconsistency::from(reason)));
}
initial
}
}
impl<Row: WarpTableRow, Brand> DoubleEndedIterator for DenseSelectionRows<'_, '_, Row, Brand> {
fn next_back(&mut self) -> Option<Self::Item> {
let row = *self.rows.next_back()?;
Some(self.row(row))
}
}
impl<Row: WarpTableRow, Brand> ExactSizeIterator for DenseSelectionRows<'_, '_, Row, Brand> {}
impl<Row: WarpTableRow, Brand> FusedIterator for DenseSelectionRows<'_, '_, Row, Brand> {}
impl<'table, Row: WarpTableRow, Brand> Selection<'table, Row, Brand> {
pub fn rows_in_dense_order<'workspace>(
self,
storage: &'workspace mut [usize],
) -> Result<DenseSelectionRows<'table, 'workspace, Row, Brand>, DenseSelectionStorageTooSmall>
where
Row::Indexes: StableSlotResolver,
{
let membership = self.membership.read();
let rows = dense_rows(self.table, membership.len(), membership.slots(), storage)?;
Ok(DenseSelectionRows::new(self.table, rows))
}
}
macro_rules! lookup_dense_rows {
($outcome:ident) => {
impl<'table, Row: WarpTableRow, Brand, Key> $outcome<'table, Row, Brand, Key> {
pub fn rows_in_dense_order<'workspace>(
self,
storage: &'workspace mut [usize],
) -> Result<
DenseSelectionRows<'table, 'workspace, Row, Brand>,
DenseSelectionStorageTooSmall,
>
where
Row::Indexes: StableSlotResolver,
{
self.into_selection().rows_in_dense_order(storage)
}
}
};
}
lookup_dense_rows!(SecondaryIndexLookupOutcome);
lookup_dense_rows!(RangeIndexLookupOutcome);
#[cfg(test)]
mod tests {
extern crate std;
use super::*;
use crate::storage::warp_table::{Capacity, WarpTableInsertOutcome};
use core::{num::NonZeroU32, panic::AssertUnwindSafe};
use std::panic::catch_unwind;
#[derive(WarpTableRow)]
struct Item {
#[index(secondary)]
group: u8,
}
#[test]
fn inconsistent_memberships_cannot_become_dense_row_snapshots() -> Result<(), &'static str> {
let mut table =
WarpTable::<Item>::try_new(Capacity::new(NonZeroU32::MIN)).map_err(|_| "table")?;
assert!(matches!(
table.insert(Item { group: 1 }),
WarpTableInsertOutcome::Inserted { .. }
));
let mut storage = [0; 2];
for (count, slots) in [(0, [0]), (2, [0]), (1, [1])] {
assert!(
catch_unwind(AssertUnwindSafe(|| {
let _result = dense_rows(
&table,
count,
slots.into_iter().map(StableSlot::new),
&mut storage,
);
}))
.is_err()
);
}
for reverse in [false, true] {
assert!(
catch_unwind(AssertUnwindSafe(|| {
let mut rows = DenseSelectionRows::new(&table, &[1]);
if reverse {
let _row = rows.next_back();
} else {
let _row = rows.next();
}
}))
.is_err()
);
}
assert!(
catch_unwind(AssertUnwindSafe(
|| DenseSelectionRows::new(&table, &[1]).fold(0, |_, _| 1)
))
.is_err()
);
let mut empty = DenseSelectionRows::new(&table, &[]);
assert_eq!(empty.len(), 0);
assert!(empty.next().is_none());
assert!(empty.next_back().is_none());
table.columns.group.clear();
assert!(
catch_unwind(AssertUnwindSafe(
|| DenseSelectionRows::new(&table, &[0]).fold(0, |_, _| 1)
))
.is_err()
);
let error = DenseSelectionStorageTooSmall {
required: 7,
available: 3,
};
let copied = error;
#[allow(clippy::clone_on_copy, reason = "exercise the public Clone contract")]
let cloned = error.clone();
assert_eq!(copied, cloned);
assert_eq!(
alloc::format!("{error:?}"),
"DenseSelectionStorageTooSmall { required: 7, available: 3 }"
);
Ok(())
}
}