use std::sync::Arc;
use std::thread;
use ordvec::{RankQuant, RankQuantFastscan};
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;
use crate::{make_corpus, D, N};
#[test]
fn fastscan_b2_top10_matches_avx512_kernel() {
const FD: usize = 128;
const FN: usize = 100;
let mut rng = ChaCha8Rng::seed_from_u64(31337);
let docs: Vec<f32> = (0..FN * FD).map(|_| rng.random_range(-1.0..1.0)).collect();
let queries: Vec<f32> = (0..3 * FD).map(|_| rng.random_range(-1.0..1.0)).collect();
let mut idx = RankQuant::new(FD, 2);
idx.add(&docs);
let ref_res = idx.search_asymmetric(&queries, 10);
let mut fs_idx = RankQuantFastscan::new(FD);
fs_idx.add(&docs);
let fs_res = fs_idx.search(&queries, 10);
for q in 0..3 {
let r_set: std::collections::HashSet<i64> = ref_res.indices[q * 10..(q + 1) * 10]
.iter()
.copied()
.collect();
let f_set: std::collections::HashSet<i64> = fs_res.indices[q * 10..(q + 1) * 10]
.iter()
.copied()
.collect();
let inter = r_set.intersection(&f_set).count();
assert!(
inter >= 9,
"query {q}: FastScan top-10 differs by >1 from AVX-512 kernel: \
ref={:?} fastscan={:?}",
ref_res.indices[q * 10..(q + 1) * 10].to_vec(),
fs_res.indices[q * 10..(q + 1) * 10].to_vec(),
);
}
}
#[test]
fn fastscan_handles_k_zero() {
let corpus = make_corpus(250);
let mut rng = ChaCha8Rng::seed_from_u64(251);
let queries: Vec<f32> = (0..(2 * D)).map(|_| rng.random_range(-1.0..1.0)).collect();
let mut fs = RankQuantFastscan::new(D);
fs.add(&corpus);
let r = fs.search(&queries, 0);
assert_eq!(r.k, 0, "result.k must equal caller's k");
assert!(r.scores.is_empty(), "scores must be empty for k=0");
assert!(r.indices.is_empty(), "indices must be empty for k=0");
}
#[test]
fn fastscan_search_on_empty_corpus_returns_sentinel() {
let fs = RankQuantFastscan::new(D);
let q: Vec<f32> = vec![0.5; D];
let r = fs.search(&q, 10);
assert_eq!(r.nq, 1);
assert_eq!(r.k, 0, "empty corpus clamps k to n_vectors (0)");
assert!(
r.indices.iter().all(|&i| i == -1),
"FastScan empty corpus: indices must all be sentinel -1"
);
}
#[test]
fn fastscan_search_on_empty_query_returns_empty_results() {
let corpus = make_corpus(260);
let mut fs = RankQuantFastscan::new(D);
fs.add(&corpus);
let r = fs.search(&[], 10);
assert_eq!(r.nq, 0);
assert!(r.scores.is_empty());
assert!(r.indices.is_empty());
}
#[test]
fn fastscan_handles_k_greater_than_n_vectors() {
const N_SMALL: usize = 5;
let mut rng = ChaCha8Rng::seed_from_u64(261);
let corpus: Vec<f32> = (0..(N_SMALL * D))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let mut fs = RankQuantFastscan::new(D);
fs.add(&corpus);
let query: Vec<f32> = corpus[0..D].to_vec();
let k = 20usize;
let r = fs.search(&query, k);
assert_eq!(r.k, N_SMALL, "k clamps to n_vectors when k > N");
for slot in 0..N_SMALL {
assert!(
r.indices[slot] >= 0 && (r.indices[slot] as usize) < N_SMALL,
"FastScan k>N: slot {slot} index {} invalid",
r.indices[slot]
);
}
}
#[test]
fn fastscan_search_is_thread_safe() {
let corpus = make_corpus(262);
let mut rng = ChaCha8Rng::seed_from_u64(263);
let queries: Vec<f32> = (0..(4 * D)).map(|_| rng.random_range(-1.0..1.0)).collect();
let mut fs = RankQuantFastscan::new(D);
fs.add(&corpus);
let fs = Arc::new(fs);
let ref_indices = fs.search(&queries, 10).indices;
let handles: Vec<_> = (0..8)
.map(|_| {
let fs = Arc::clone(&fs);
let queries = queries.clone();
let expected = ref_indices.clone();
thread::spawn(move || {
for _ in 0..8 {
assert_eq!(fs.search(&queries, 10).indices, expected);
}
})
})
.collect();
for h in handles {
h.join().unwrap();
}
}
#[test]
fn fastscan_dim_boundary_matrix() {
for &dim in &[64usize, 1024] {
const N_SMALL: usize = 16;
let mut rng = ChaCha8Rng::seed_from_u64(270 + dim as u64);
let corpus: Vec<f32> = (0..(N_SMALL * dim))
.map(|_| rng.random_range(-1.0..1.0))
.collect();
let mut fs = RankQuantFastscan::new(dim);
fs.add(&corpus);
let q: Vec<f32> = corpus[0..dim].to_vec();
let r = fs.search(&q, 5);
assert_eq!(r.k, 5);
for &idx in &r.indices {
assert!(
idx >= 0 && (idx as usize) < N_SMALL,
"FastScan dim={dim}: invalid index {idx}"
);
}
assert_eq!(
r.indices[0], 0,
"FastScan dim={dim}: top-1 of self-query should be 0"
);
}
}
#[test]
#[should_panic(expected = "incremental add()")]
fn fastscan_second_add_panics_per_v1_contract() {
let corpus = make_corpus(280);
let mut fs = RankQuantFastscan::new(D);
fs.add(&corpus);
fs.add(&corpus); }
#[test]
fn fastscan_construct_then_metadata_roundtrips() {
let corpus = make_corpus(281);
let mut fs = RankQuantFastscan::new(D);
assert!(fs.is_empty());
assert_eq!(fs.len(), 0);
assert_eq!(fs.dim(), D);
assert_eq!(fs.bytes_per_vec(), D / 2);
assert_eq!(fs.byte_size(), 0);
fs.add(&corpus);
assert!(!fs.is_empty());
assert_eq!(fs.len(), N);
assert_eq!(fs.dim(), D);
assert_eq!(fs.bytes_per_vec(), D / 2);
assert_eq!(fs.byte_size(), N * (D / 2));
}
#[test]
#[should_panic(expected = "divisible by 4")]
fn fastscan_new_rejects_dim_2_not_multiple_of_4() {
let _ = RankQuantFastscan::new(2);
}
#[test]
#[should_panic(expected = "divisible by 4")]
fn fastscan_new_rejects_dim_6_not_multiple_of_4() {
let _ = RankQuantFastscan::new(6);
}
#[test]
#[should_panic(expected = "fit in u16")]
fn fastscan_new_rejects_dim_above_u16_max() {
let _ = RankQuantFastscan::new(65_536);
}