use super::*;
use uqa_core::{
memory::{MemoryBudget, MemoryError},
QueryCancelled, TokenOffsets,
};
mod reference;
#[test]
fn occurrence_validation_borrows_complete_graphs_and_preserves_cancellation() {
let entries = graph_fixture(0, 129);
let (scores, positions) = encode_occurrence_cluster(&entries).unwrap();
let info = allocation_counter::measure(|| {
validate_occurrence_cluster(0, &scores, &positions, || Ok(())).unwrap();
});
assert_eq!(info.count_total, 0);
for end in 0..positions.len() {
assert!(validate_occurrence_cluster(0, &scores, &positions[..end], || Ok(())).is_err());
}
let mut polls = 0;
let result = validate_occurrence_cluster(0, &scores, &positions, || {
polls += 1;
if polls == 20 {
Err(QueryCancelled.into())
} else {
Ok(())
}
});
assert!(matches!(result, Err(StorageBackendError::Cancelled(_))));
assert_eq!(polls, 20);
}
#[test]
fn occurrence_encoding_preserves_wire_bytes_across_varint_widths() {
let entries = vec![
OccurrencePosting {
doc_id: 2,
doc_length: 128,
occurrences: vec![TokenOccurrence {
position: 0,
position_length: 1,
offsets: None,
}],
},
OccurrencePosting {
doc_id: 16_384,
doc_length: u64::MAX,
occurrences: vec![TokenOccurrence {
position: 128,
position_length: 2,
offsets: None,
}],
},
];
let expected_scores = b"UQCS\x02\0\0\0\x02\0\0\0\x01\0\0\0\x02\0\0\x40\x2c\0\0\0\x2f\0\0\0\x2f\0\0\0\x31\0\0\0\x31\0\0\0\x3d\0\0\0\x02\xfe\x7f\x01\x01\x80\x01\xff\xff\xff\xff\xff\xff\xff\xff\xff\x01";
let expected_positions =
b"UQCP\x02\0\0\0\x02\0\0\0\x03\0\0\0\0\0\0\0\x03\0\0\0\x07\0\0\0\0\x01\0\x80\x01\x02\0";
let control = crate::read_control::StorageReadControl::with_limit(4096);
let (scores, positions) =
encode_occurrence_cluster_controlled(entries.iter(), &control).unwrap();
assert_eq!(&*scores, expected_scores);
assert_eq!(&*positions, expected_positions);
assert_eq!(
reference::decode_occurrence_cluster(0, &scores, &positions).unwrap(),
entries
);
drop((scores, positions));
assert_eq!(control.memory().used(), 0);
}
#[test]
fn occurrence_encoding_retains_and_releases_its_complete_budget() {
let entries = graph_fixture(0, 129);
let mut failures = 0;
let mut successes = 0;
for limit in [0, 32, 128, 512, 4096, 16384, 65536] {
let control = crate::read_control::StorageReadControl::with_limit(limit);
match encode_occurrence_cluster_controlled(entries.iter(), &control) {
Ok((scores, positions)) => {
successes += 1;
assert_eq!(
decode_occurrence_cluster(0, &scores, &positions).unwrap(),
entries
);
assert!(control.memory().used() >= scores.capacity() + positions.capacity());
}
Err(StorageBackendError::Memory(_)) => failures += 1,
Err(error) => panic!("unexpected encoding error: {error}"),
}
assert_eq!(control.memory().used(), 0);
}
assert!(failures > 0 && successes > 0);
let control = crate::read_control::StorageReadControl::with_limit(0);
control.cancellation().cancel();
assert!(matches!(
encode_occurrence_cluster_controlled(entries.iter(), &control),
Err(StorageBackendError::Cancelled(_))
));
assert_eq!(control.memory().used(), 0);
}
#[test]
fn successive_cluster_encodings_share_the_allowance_and_cancellation() {
let entries = graph_fixture(0, 129);
let probe = crate::read_control::StorageReadControl::with_limit(1 << 24);
let (scores, positions) = encode_occurrence_cluster_controlled(entries.iter(), &probe).unwrap();
let required_peak = probe.memory().peak();
let retained = probe.memory().used();
drop((scores, positions));
assert_eq!(probe.memory().used(), 0);
let control = crate::read_control::StorageReadControl::with_limit(required_peak);
let first = encode_occurrence_cluster_controlled(entries.iter(), &control).unwrap();
assert_eq!(control.memory().used(), retained);
assert!(matches!(
encode_occurrence_cluster_controlled(entries.iter(), &control),
Err(StorageBackendError::Memory(_))
));
assert_eq!(control.memory().used(), retained);
assert_eq!(
decode_occurrence_cluster(0, &first.0, &first.1).unwrap(),
entries
);
drop(first);
assert_eq!(control.memory().used(), 0);
let second = encode_occurrence_cluster_controlled(entries.iter(), &control).unwrap();
control.cancellation().cancel();
assert!(matches!(
encode_occurrence_cluster_controlled(entries.iter(), &control),
Err(StorageBackendError::Cancelled(_))
));
assert_eq!(control.memory().used(), retained);
drop(second);
assert_eq!(control.memory().used(), 0);
}
fn graph_fixture(cluster: u64, count: usize) -> Vec<OccurrencePosting> {
let base = cluster_base(cluster).unwrap();
(0..count)
.map(|index| OccurrencePosting {
doc_id: base
+ if index + 1 == count {
65_535
} else {
(index * 2) as u64
},
doc_length: 1,
occurrences: (0..=index % 4)
.map(|ordinal| TokenOccurrence {
position: (ordinal / 2) as u32,
position_length: 1 + (ordinal % 3) as u32,
offsets: (ordinal % 2 == 0).then_some(TokenOffsets {
start_utf8: 3,
end_utf8: 7,
start_utf16: 3,
end_utf16: 4,
}),
})
.collect(),
})
.collect()
}
fn compare_decoders(cluster: u64, scores: &[u8], positions: &[u8]) {
let expected = reference::decode_occurrence_cluster(cluster, scores, positions);
let budget = MemoryBudget::new(1 << 24);
let unrelated = budget.reserve(7).unwrap();
let actual = decode_occurrence_cluster_budgeted(cluster, scores, positions, &budget, || Ok(()));
match (&expected, actual) {
(Ok(expected), Ok(actual)) => assert_eq!(actual.as_slice(), expected),
(Err(_), Err(_)) => {}
(expected, actual) => {
panic!("decoder acceptance changed: expected={expected:?}, actual={actual:?}")
}
}
assert_eq!(budget.used(), 7);
for target in [
cluster_base(cluster).unwrap(),
cluster_base(cluster).unwrap() + 1,
cluster_base(cluster).unwrap() + 65_535,
] {
let actual = decode_occurrence_document_budgeted(
cluster,
scores,
positions,
target,
&budget,
|| Ok(()),
);
match (&expected, actual) {
(Ok(expected), Ok(actual)) => {
assert_eq!(actual.as_deref(), expected.iter().find(|posting| posting.doc_id == target));
}
(Err(_), Err(_)) => {},
(expected, actual) => panic!("selected decoder acceptance changed for {target}: expected={expected:?}, actual={actual:?}"),
}
assert_eq!(budget.used(), 7);
}
drop(unrelated);
assert_eq!(budget.used(), 0);
}
#[test]
fn borrowed_decoders_preserve_complete_format_and_corruption_acceptance() {
for cluster in [0, 17, cluster_id(DocId::MAX)] {
for count in [1, 2, 127, 128, 129, 260] {
let entries = graph_fixture(cluster, count);
let (scores, positions) = encode_occurrence_cluster(&entries).unwrap();
assert_eq!(
reference::decode_occurrence_cluster(cluster, &scores, &positions).unwrap(),
entries
);
compare_decoders(cluster, &scores, &positions);
for (which, original) in [(0, &scores), (1, &positions)] {
for at in 0..original.len() {
if count > 2 && at >= 80 && at % 127 != 0 && at + 4 < original.len() {
continue;
}
let mut bytes = original.clone();
bytes[at] ^= 0x80;
let (score, position) = if which == 0 {
(bytes.as_slice(), positions.as_slice())
} else {
(scores.as_slice(), bytes.as_slice())
};
compare_decoders(cluster, score, position);
let (score, position) = if which == 0 {
(&scores[..at], positions.as_slice())
} else {
(scores.as_slice(), &positions[..at])
};
compare_decoders(cluster, score, position);
}
}
}
}
}
#[test]
fn selected_reads_validate_other_documents_with_only_selected_output_memory() {
let entries = graph_fixture(cluster_id(DocId::MAX), 260);
let (scores, positions) = encode_occurrence_cluster(&entries).unwrap();
let bytes = entries[0].occurrences.len() * size_of::<TokenOccurrence>();
let budget = MemoryBudget::new(bytes + 7);
let unrelated = budget.reserve(7).unwrap();
let selected = decode_occurrence_document_budgeted(
cluster_id(DocId::MAX),
&scores,
&positions,
entries[0].doc_id,
&budget,
|| Ok(()),
)
.unwrap()
.unwrap();
assert_eq!(*selected, entries[0]);
assert_eq!(budget.used(), bytes + 7);
assert_eq!(budget.peak(), bytes + 7);
let missing = decode_occurrence_document_budgeted(
cluster_id(DocId::MAX),
&scores,
&positions,
entries[0].doc_id + 1,
&budget,
|| Ok(()),
)
.unwrap();
assert!(missing.is_none());
let mut corrupt_positions = positions.clone();
let offset_count = read_u32(&positions, 12).unwrap() as usize;
let payload_start = HEADER_LEN + offset_count * 4;
let final_start = read_u32(&positions, HEADER_LEN + (entries.len() - 1) * 4).unwrap() as usize;
corrupt_positions[payload_start + final_start + 2] = 2;
assert!(reference::decode_occurrence_cluster(
cluster_id(DocId::MAX),
&scores,
&corrupt_positions
)
.is_err());
assert!(decode_occurrence_document_budgeted(
cluster_id(DocId::MAX),
&scores,
&corrupt_positions,
entries[0].doc_id + 1,
&budget,
|| Ok(())
)
.is_err());
assert_eq!(budget.used(), bytes + 7);
drop(selected);
assert_eq!(budget.used(), 7);
drop(unrelated);
}
#[test]
fn decoder_limits_and_every_callback_failure_release_only_their_output() {
let entries = graph_fixture(0, 4);
let (scores, positions) = encode_occurrence_cluster(&entries).unwrap();
let baseline = MemoryBudget::new(1 << 20);
let mut calls = 0;
let reference = decode_occurrence_cluster_budgeted(0, &scores, &positions, &baseline, || {
calls += 1;
Ok(())
})
.unwrap();
let required = baseline.peak();
assert_eq!(required, reference.reserved_bytes());
for limit in 0..=required {
let budget = MemoryBudget::new(limit + 7);
let unrelated = budget.reserve(7).unwrap();
match decode_occurrence_cluster_budgeted(0, &scores, &positions, &budget, || Ok(())) {
Ok(actual) => {
assert_eq!(*actual, *reference);
assert_eq!(limit, required);
}
Err(StorageBackendError::Memory(MemoryError::Limit { .. })) => {
assert!(limit < required);
}
result => panic!("unexpected allocation result at {limit}: {result:?}"),
}
assert_eq!(budget.used(), 7);
drop(unrelated);
}
let budget = MemoryBudget::new(1 << 20);
let prior = decode_occurrence_document_budgeted(
0,
&scores,
&positions,
entries[0].doc_id,
&budget,
|| Ok(()),
)
.unwrap()
.unwrap();
let retained = budget.used();
for stop in 1..=calls {
let mut current = 0;
assert!(matches!(
decode_occurrence_cluster_budgeted(0, &scores, &positions, &budget, || {
current += 1;
if current == stop {
Err(QueryCancelled.into())
} else {
Ok(())
}
}),
Err(StorageBackendError::Cancelled(_))
));
assert_eq!(budget.used(), retained);
assert_eq!(*prior, entries[0]);
}
for target in [
entries[0].doc_id,
entries[0].doc_id + 1,
entries.last().unwrap().doc_id,
] {
let mut calls = 0;
let output =
decode_occurrence_document_budgeted(0, &scores, &positions, target, &budget, || {
calls += 1;
Ok(())
})
.unwrap();
drop(output);
for stop in 1..=calls {
let mut current = 0;
assert!(matches!(
decode_occurrence_document_budgeted(
0,
&scores,
&positions,
target,
&budget,
|| {
current += 1;
if current == stop {
Err(QueryCancelled.into())
} else {
Ok(())
}
}
),
Err(StorageBackendError::Cancelled(_))
));
assert_eq!(budget.used(), retained);
}
}
drop(prior);
assert_eq!(budget.used(), 0);
}