use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec::Vec;
use proptest::prelude::*;
use super::{
CandidateBlock, PreferredRows, RangeIndexInconsistency, RangeIndexOrderingDirection,
RangeIndexOrderingError, RangeIndexPlanes, preferred_partition, slot_blocks,
};
use crate::storage::query_workspace::{
QueryChunk, QueryWorkspace, QueryWorkspaceError, WordOperation,
};
use crate::storage::warp_table::secondary_index::{SlotSet, SlotSetInsertOutcome, StableSlot};
fn assert_blocks(values: BTreeSet<u32>) -> Result<(), &'static str> {
let mut slots = SlotSet::empty();
let mut expected = BTreeMap::<u32, u64>::new();
for value in values {
let SlotSetInsertOutcome::Inserted { slot } = slots
.try_prepare_insert(StableSlot::new(value))
.map_err(|_| "distinct slot insertion failed")?
.commit();
assert_eq!(slot.value(), value);
*expected.entry((value / 64).wrapping_mul(64)).or_default() |= 1_u64 << (value % 64);
}
let actual = slot_blocks(slots.read());
assert_eq!(actual.capacity(), expected.len());
assert_eq!(
actual
.iter()
.map(|block| (block.slot_base, block.boundary))
.collect::<Vec<_>>(),
expected.into_iter().collect::<Vec<_>>()
);
for block in actual {
assert_eq!((block.leaders, block.planes), (0, &[][..]));
}
Ok(())
}
#[test]
fn range_masks_preserve_boundaries_gaps_and_the_final_slot() -> Result<(), &'static str> {
for values in [
BTreeSet::new(),
BTreeSet::from([u32::MAX]),
BTreeSet::from([0, 1, 63, 64, 65, 65_535, 65_536, u32::MAX]),
(0..64).collect(),
(1..65).collect(),
(63..128).collect(),
(0..65_537).collect(),
(65_471..65_600).collect(),
(u32::MAX - 127..=u32::MAX).collect(),
(0..64).chain(65..129).collect(),
(1..65).chain(66..130).collect(),
(0..64).chain(u32::MAX - 63..=u32::MAX).collect(),
(0..64).chain([65, 67, 129]).collect(),
[0].into_iter().chain(64..192).collect(),
(0..192).step_by(2).collect(),
] {
assert_blocks(values)?;
}
Ok(())
}
#[test]
fn workspace_blocks_preserve_partial_words_chunks_and_high_slots() -> Result<(), &'static str> {
for capacity in [
0_usize, 1, 63, 64, 65, 255, 256, 257, 511, 512, 513, 1023, 1024, 1025, 2047, 2048, 2049,
4092, 4095, 4096, 4097, 65535, 65536, 65537, 131073,
] {
let mut workspace = QueryWorkspace::new(
capacity,
core::iter::repeat_with(|| QueryChunk::EMPTY)
.take(capacity.div_ceil(65536).saturating_add(1))
.collect::<Vec<_>>(),
alloc::boxed::Box::new(QueryChunk::EMPTY),
)
.map_err(|_| "workspace")?;
for dense in [false, true] {
for values in [
BTreeSet::new(),
(0..capacity as u32).step_by(11).collect(),
(0..capacity as u32).rev().take(3).collect(),
] {
workspace.clear();
let mut expected = BTreeMap::<u32, u64>::new();
{
let mut membership = workspace.membership().map_err(|_| "membership")?;
for value in values {
membership.insert(value).map_err(|_| "insert")?;
*expected.entry(value & !63).or_default() |= 1_u64 << (value & 63);
}
if dense {
membership
.write_words(
capacity,
WordOperation::Combine(
crate::storage::query_workspace::MembershipOperation::Union,
),
core::iter::empty::<Result<u64, QueryWorkspaceError>>(),
)
.map_err(|_| "dense")?;
}
assert_eq!(
super::workspace_block_count(&membership).map_err(|_| "count")?,
expected.len()
);
let mut actual = Vec::new();
super::visit_workspace_blocks(&membership, |block| {
actual.push((block.slot_base, block.boundary))
})
.map_err(|_| "blocks")?;
assert_eq!(
actual,
expected
.iter()
.map(|(base, mask)| (*base, *mask))
.collect::<Vec<_>>()
);
}
let (chunks, capacity, len) = workspace.read_storage();
let actual = slot_blocks(super::SlotSetRead::Workspace {
chunks,
capacity,
len,
});
assert_eq!(actual.capacity(), expected.len());
assert_eq!(
actual
.iter()
.map(|block| (block.slot_base, block.boundary))
.collect::<Vec<_>>(),
expected.into_iter().collect::<Vec<_>>()
);
}
}
}
Ok(())
}
#[test]
fn workspace_capacity_does_not_hide_candidates_outside_the_ordering_index()
-> Result<(), &'static str> {
let level = super::RangeIndexOrderingLevel {
direction: RangeIndexOrderingDirection::SmallestFirst,
planes: Ok(RangeIndexPlanes {
field: "rank",
maximum_rows: 3,
words_per_block: 2,
blocks: &[7, 1],
}),
};
let mut chunks = [QueryChunk::EMPTY];
let mut workspace =
QueryWorkspace::new(65, &mut chunks, QueryChunk::EMPTY).map_err(|_| "workspace")?;
let expected = Some(RangeIndexOrderingError::Inconsistent {
field: "rank",
reason: RangeIndexInconsistency::SlotOutsideCapacity { maximum_rows: 3 },
});
{
let mut membership = workspace.membership().map_err(|_| "membership")?;
membership.insert(64).map_err(|_| "insert")?;
membership
.write_words(
65,
WordOperation::Combine(crate::storage::query_workspace::MembershipOperation::Union),
core::iter::empty::<Result<u64, QueryWorkspaceError>>(),
)
.map_err(|_| "dense")?;
assert_eq!(
super::rank(
&level,
&[],
super::RangeIndexRankingCandidates::Workspace(&membership),
core::num::NonZeroUsize::MIN
)
.err(),
expected
);
}
let (chunks, capacity, len) = workspace.read_storage();
assert_eq!(
super::rank(
&level,
&[],
super::RangeIndexRankingCandidates::Slots(super::SlotSetRead::Workspace {
chunks: chunks.as_slice(),
capacity,
len
}),
core::num::NonZeroUsize::MIN
)
.err(),
expected
);
Ok(())
}
proptest! {
#[test]
fn candidate_masks_match_independent_membership(
runs in prop::collection::vec((any::<u32>(), 64_u32..257), 0..17),
singles in prop::collection::vec(any::<u32>(), 0..65),
) {
let mut values = BTreeSet::new();
for (first, length) in runs {
values.extend(first..=first.saturating_add(length));
}
assert_blocks(values.clone()).map_err(proptest::test_runner::TestCaseError::fail)?;
values.extend(singles);
assert_blocks(values).map_err(proptest::test_runner::TestCaseError::fail)?;
}
}
#[test]
fn bounded_counting_preserves_exact_promotions() -> Result<(), &'static str> {
let words = [0b1011_u64];
let planes = RangeIndexPlanes {
field: "rank",
maximum_rows: 192,
words_per_block: 2,
blocks: &[],
};
for remaining_limit in [0, 1, 3, 4, 5, usize::MAX] {
let mut blocks = [
CandidateBlock::new(0, u64::MAX),
CandidateBlock::new(64, 0b1100),
CandidateBlock::new(128, 0),
];
for block in &mut blocks {
block.planes = &words;
}
let outcome = preferred_partition(
&planes,
&blocks,
0,
RangeIndexOrderingDirection::LargestFirst,
remaining_limit,
)
.map_err(|_| "valid planes were refused")?;
match outcome {
PreferredRows::Promote { count } => {
assert!(remaining_limit >= 4);
assert_eq!(count, 4);
}
PreferredRows::Narrow => assert!(remaining_limit < 4),
PreferredRows::Unchanged => return Err("mixed partition was unchanged"),
}
}
match preferred_partition(
&planes,
&[],
0,
RangeIndexOrderingDirection::LargestFirst,
0,
)
.map_err(|_| "empty candidates were refused")?
{
PreferredRows::Unchanged => {}
PreferredRows::Promote { .. } | PreferredRows::Narrow => {
return Err("empty candidates changed");
}
}
Ok(())
}
#[test]
fn bounded_counting_still_validates_planes_after_exceeding_the_limit() {
let planes = RangeIndexPlanes {
field: "rank",
maximum_rows: 128,
words_per_block: 2,
blocks: &[],
};
for (word, remaining_limit) in [
(u64::MAX, 0),
(u64::MAX, 64),
(u64::MAX - 1, 0),
(u64::MAX - 1, 63),
] {
let words = [word];
let mut blocks = [CandidateBlock::new(0, u64::MAX), CandidateBlock::new(64, 1)];
if let Some(first) = blocks.first_mut() {
first.planes = &words;
}
let outcome = preferred_partition(
&planes,
&blocks,
0,
RangeIndexOrderingDirection::LargestFirst,
remaining_limit,
);
let outcome = outcome.and_then(|partition| {
super::settle(
&planes,
&mut blocks,
0,
RangeIndexOrderingDirection::LargestFirst,
&mut |_| {},
&partition,
)
});
assert_eq!(
outcome.err(),
Some(RangeIndexOrderingError::Inconsistent {
field: "rank",
reason: RangeIndexInconsistency::BitPositionOutsideWord { bit: 0 },
})
);
}
}
#[test]
fn refinement_preserves_a_late_partition_error() {
for direction in [
RangeIndexOrderingDirection::LargestFirst,
RangeIndexOrderingDirection::SmallestFirst,
] {
let words = [match direction {
RangeIndexOrderingDirection::LargestFirst => u64::MAX - 1,
RangeIndexOrderingDirection::SmallestFirst => 1,
}];
let level = super::RangeIndexOrderingLevel {
direction,
planes: Ok(RangeIndexPlanes {
field: "rank",
maximum_rows: 128,
words_per_block: 2,
blocks: &[],
}),
};
let mut first = CandidateBlock::new(0, u64::MAX);
first.planes = &words;
let mut blocks = [first, CandidateBlock::new(64, 1)];
let mut settled = Vec::new();
assert_eq!(
super::refine_candidates(
&level,
[],
&mut blocks,
core::num::NonZeroUsize::MIN,
|block| settled.push(block),
),
Err(RangeIndexOrderingError::Inconsistent {
field: "rank",
reason: RangeIndexInconsistency::BitPositionOutsideWord { bit: 0 },
})
);
assert!(settled.is_empty());
}
}
fn assert_partition(
inputs: &[(u64, u64, u64)],
remaining: usize,
largest: bool,
) -> Result<(), &'static str> {
let planes = RangeIndexPlanes {
field: "partition",
maximum_rows: inputs.len().saturating_mul(64),
words_per_block: 2,
blocks: &[],
};
let mut blocks = Vec::new();
let mut candidates = BTreeSet::new();
let mut leaders = BTreeSet::new();
let mut preferred = BTreeSet::new();
for (ordinal, (boundary, plane, previous)) in inputs.iter().enumerate() {
let base = (ordinal as u32) << 6;
let boundary = boundary | 1;
let previous = previous & !boundary;
let mut block = CandidateBlock::new(base, boundary);
block.leaders = previous;
block.planes = core::slice::from_ref(plane);
blocks.push(block);
for bit in 0..64 {
let slot = base | bit;
if boundary & (1_u64 << bit) != 0 {
candidates.insert(slot);
if (plane & (1_u64 << bit) != 0) == largest {
preferred.insert(slot);
}
}
if previous & (1_u64 << bit) != 0 {
leaders.insert(slot);
}
}
}
let expected_boundary = if preferred.len() <= remaining {
leaders.extend(&preferred);
candidates
.difference(&preferred)
.copied()
.collect::<BTreeSet<_>>()
} else {
preferred.clone()
};
let unchanged = expected_boundary == candidates;
let outcome = preferred_partition(
&planes,
&blocks,
0,
if largest {
RangeIndexOrderingDirection::LargestFirst
} else {
RangeIndexOrderingDirection::SmallestFirst
},
remaining,
)
.map_err(|_| "partition refused")?;
match outcome {
PreferredRows::Unchanged => assert!(unchanged),
PreferredRows::Promote { count } => {
assert!(!unchanged);
assert_eq!(count, preferred.len());
assert!(count <= remaining);
}
PreferredRows::Narrow => {
assert!(!unchanged);
assert!(preferred.len() > remaining);
}
}
let mut actual_leaders = BTreeSet::new();
let mut actual_boundary = BTreeSet::new();
let active = super::settle(
&planes,
&mut blocks,
0,
if largest {
RangeIndexOrderingDirection::LargestFirst
} else {
RangeIndexOrderingDirection::SmallestFirst
},
&mut |block| {
for bit in 0..64 {
if block.slots & (1_u64 << bit) != 0 {
assert!(actual_leaders.insert(block.slot_base | bit));
}
}
},
&outcome,
)
.map_err(|_| "settlement refused")?;
for block in blocks.iter().take(active) {
assert_ne!(block.boundary, 0);
for bit in 0..64 {
if block.leaders & (1_u64 << bit) != 0 {
assert!(actual_leaders.insert(block.slot_base | bit));
}
if block.boundary & (1_u64 << bit) != 0 {
assert!(actual_boundary.insert(block.slot_base | bit));
}
}
}
assert_eq!(actual_leaders, leaders);
assert_eq!(actual_boundary, expected_boundary);
Ok(())
}
#[test]
fn uniform_bits_preserve_candidates_and_previously_promoted_leaders() -> Result<(), &'static str> {
for inputs in [
Vec::new(),
alloc::vec![(0b11, 0, 0b1100), (0b101, 0, 0b1010)],
alloc::vec![(0b11, u64::MAX, 0b1100), (0b101, u64::MAX, 0b1010)],
alloc::vec![(0b11, 0b11, 0b1100), (0b101, 0, 0b1010)],
alloc::vec![(u64::MAX, u64::MAX, 0); 17],
] {
for remaining in [0, 1, 3, 4, 63, 64, 1000, usize::MAX] {
for largest in [false, true] {
assert_partition(&inputs, remaining, largest)?;
}
}
}
Ok(())
}
proptest! {
#[test]
fn refinement_partition_matches_individual_slot_sets(
inputs in prop::collection::vec((any::<u64>(), any::<u64>(), any::<u64>()), 0..32),
remaining in 0..2049_usize,
largest in any::<bool>(),
) {
assert_partition(&inputs, remaining, largest).map_err(proptest::test_runner::TestCaseError::fail)?;
}
}
fn assert_distinguishing_bits(values: &[(u64, bool)], bits: usize) {
let mask = u64::MAX
.checked_shr(64_usize.saturating_sub(bits) as u32)
.unwrap_or(0);
let storage = values
.chunks(64)
.map(|rows| {
let mut boundary = 0_u64;
let mut planes = alloc::vec![0_u64; bits];
for (slot, (value, selected)) in rows.iter().enumerate() {
let member = 1_u64 << slot;
if *selected {
boundary |= member;
}
for (bit, plane) in planes.iter_mut().enumerate() {
if value & (1_u64 << bit) != 0 {
*plane |= member;
}
}
}
(boundary, planes)
})
.collect::<Vec<_>>();
let blocks = storage
.iter()
.filter(|(boundary, _)| *boundary != 0)
.map(|(boundary, planes)| CandidateBlock {
planes,
..CandidateBlock::new(0, *boundary)
})
.collect::<Vec<_>>();
let selected = values
.iter()
.filter(|(_, selected)| *selected)
.map(|(value, _)| value & mask);
let minimum = selected.clone().min().unwrap_or(0);
let maximum = selected.max().unwrap_or(0);
let expected = 64_u32.saturating_sub((minimum ^ maximum).leading_zeros());
assert_eq!(super::distinguishing_bits(&blocks), expected as usize);
}
#[test]
fn distinguishing_bits_preserve_uniform_prefixes_across_blocks_and_key_widths() {
assert_distinguishing_bits(&[], 64);
for bits in [8, 16, 32, 64] {
for prefix in [0, u64::MAX, 0xaaaa_5555_aaaa_5555] {
assert_distinguishing_bits(&alloc::vec![(prefix, true); 193], bits);
for bit in 0..bits {
for position in [0, 1, 63, 64, 65, 127, 192] {
let mut values = alloc::vec![(prefix, true); 193];
if let Some((value, _)) = values.get_mut(position) {
*value ^= 1_u64 << bit;
}
assert_distinguishing_bits(&values, bits);
}
}
}
}
}
#[test]
fn prefix_discovery_does_not_hide_absent_members_in_later_blocks_or_levels()
-> Result<(), &'static str> {
let mut slots = SlotSet::empty();
for slot in [0, 1, 65] {
let SlotSetInsertOutcome::Inserted { .. } = slots
.try_prepare_insert(StableSlot::new(slot))
.map_err(|_| "slot insertion")?
.commit();
}
for high_plane in [0, 2] {
let words = [
3, 0, 0, 0, 0, 0, 0, 0, high_plane, 1, 0, 0, 0, 0, 0, 0, 0, 0,
];
let uniform_words = [3, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0];
let level = |blocks| super::RangeIndexOrderingLevel {
direction: RangeIndexOrderingDirection::SmallestFirst,
planes: Ok(RangeIndexPlanes {
field: "rank",
maximum_rows: 128,
words_per_block: 9,
blocks,
}),
};
let uniform = level(&uniform_words);
let corrupt = level(&words);
for (first, later) in [
(&corrupt, &[][..]),
(&uniform, core::slice::from_ref(&corrupt)),
] {
let expected = RangeIndexOrderingError::Inconsistent {
field: "rank",
reason: RangeIndexInconsistency::MemberAbsent,
};
assert_eq!(
super::rank(
first,
later,
super::RangeIndexRankingCandidates::Slots(slots.read()),
core::num::NonZeroUsize::MIN,
)
.err(),
Some(expected)
);
let mut storage = [core::mem::MaybeUninit::uninit(); 512];
assert_eq!(
super::rank_with(
first,
later,
super::RangeIndexRankingCandidates::Slots(slots.read()),
core::num::NonZeroUsize::MIN,
crate::storage::warp_table::StorageLease::new(&mut storage),
)
.err(),
Some(super::RankingWorkspaceError::Ordering(expected))
);
}
}
Ok(())
}
proptest! {
#[test]
fn distinguishing_bits_match_selected_key_extrema(
values in prop::collection::vec((any::<u64>(), any::<bool>()), 0..258),
prefix in any::<u64>(),
varying in 0..65_u32,
bits in prop::sample::select(alloc::vec![8_usize, 16, 32, 64]),
) {
let mask = u64::MAX.checked_shr(64_u32.saturating_sub(varying)).unwrap_or(0);
let values = values.into_iter().map(|(value, selected)| {
(if selected { prefix & !mask | value & mask } else { value }, selected)
}).collect::<Vec<_>>();
assert_distinguishing_bits(&values, bits);
}
}