#![allow(clippy::unwrap_used)]
use iqdb_filter::{FilterEvaluator, FilterStrategy, MetadataIndex, StrategySelector};
use iqdb_types::{Filter, Metadata, Value};
struct SimIndex {
metadata: Vec<Option<Metadata>>,
scores: Vec<f64>,
index: MetadataIndex<usize>,
}
impl SimIndex {
fn new(indexed_fields: &[&str], rows: Vec<(Option<Metadata>, f64)>) -> Self {
let metadata: Vec<Option<Metadata>> = rows.iter().map(|(m, _)| m.clone()).collect();
let scores: Vec<f64> = rows.iter().map(|(_, s)| *s).collect();
let index = MetadataIndex::build(
indexed_fields,
metadata.iter().enumerate().map(|(i, m)| (i, m.as_ref())),
);
Self {
metadata,
scores,
index,
}
}
fn metadata_of(&self, key: usize) -> Option<&Metadata> {
self.metadata[key].as_ref()
}
fn rank(&self, mut keys: Vec<usize>, query: f64, k: usize) -> Vec<usize> {
keys.sort_by(|&a, &b| {
let da = (self.scores[a] - query).abs();
let db = (self.scores[b] - query).abs();
da.partial_cmp(&db).unwrap().then(a.cmp(&b))
});
keys.truncate(k);
keys
}
fn search_scan(&self, evaluator: &FilterEvaluator, query: f64, k: usize) -> Vec<usize> {
let kept: Vec<usize> = evaluator
.prefilter((0..self.metadata.len()).map(|i| (i, self.metadata_of(i))))
.collect();
self.rank(kept, query, k)
}
fn search_indexed(&self, evaluator: &FilterEvaluator, query: f64, k: usize) -> Vec<usize> {
let kept: Vec<usize> = match self.index.candidates(evaluator) {
Some(candidates) => candidates
.into_iter()
.filter(|&key| evaluator.evaluate(self.metadata_of(key)))
.collect(),
None => evaluator
.prefilter((0..self.metadata.len()).map(|i| (i, self.metadata_of(i))))
.collect(),
};
self.rank(kept, query, k)
}
}
fn meta(pairs: &[(&str, Value)]) -> Metadata {
pairs
.iter()
.map(|(k, v)| ((*k).to_string(), v.clone()))
.collect()
}
fn corpus() -> SimIndex {
let rows = vec![
(
Some(meta(&[
("lang", Value::String("rust".into())),
("year", Value::Int(2026)),
("published", Value::Bool(true)),
])),
0.9,
),
(
Some(meta(&[
("lang", Value::String("go".into())),
("year", Value::Int(2024)),
("published", Value::Bool(true)),
])),
0.5,
),
(
Some(meta(&[
("lang", Value::String("rust".into())),
("year", Value::Int(2019)),
("published", Value::Bool(false)),
])),
0.7,
),
(
Some(meta(&[
("lang", Value::String("rust".into())),
("year", Value::Int(2026)),
("published", Value::Bool(true)),
])),
0.2,
),
(None, 0.4),
];
SimIndex::new(&["lang", "year", "published"], rows)
}
fn assert_paths_agree(filter: Filter, query: f64, k: usize) {
let index = corpus();
let evaluator = FilterEvaluator::new(filter).unwrap();
let scanned = index.search_scan(&evaluator, query, k);
let indexed = index.search_indexed(&evaluator, query, k);
assert_eq!(
scanned, indexed,
"index-accelerated search diverged from full scan"
);
}
#[test]
fn indexed_equality_matches_scan() {
assert_paths_agree(Filter::eq("lang", Value::String("rust".into())), 0.5, 10);
}
#[test]
fn compound_and_matches_scan() {
assert_paths_agree(
Filter::and(vec![
Filter::eq("lang", Value::String("rust".into())),
Filter::eq("published", Value::Bool(true)),
]),
0.5,
10,
);
}
#[test]
fn unbounded_predicate_falls_back_but_agrees() {
assert_paths_agree(Filter::gt("year", Value::Int(2020)), 0.5, 10);
assert_paths_agree(
Filter::not(Filter::eq("lang", Value::String("rust".into()))),
0.5,
10,
);
}
#[test]
fn mixed_resolvable_and_unresolvable_agrees() {
assert_paths_agree(
Filter::and(vec![
Filter::eq("lang", Value::String("rust".into())),
Filter::gt("year", Value::Int(2020)),
]),
0.5,
10,
);
}
#[test]
fn truncates_to_k() {
let index = corpus();
let evaluator = FilterEvaluator::new(Filter::eq("lang", Value::String("rust".into()))).unwrap();
let top1 = index.search_indexed(&evaluator, 0.0, 1);
assert_eq!(top1.len(), 1);
assert_eq!(top1, [3]);
}
#[test]
fn postfilter_refills_top_k_lazily() {
let index = corpus();
let evaluator = FilterEvaluator::new(Filter::eq("published", Value::Bool(true))).unwrap();
let ranked = index.rank(
(0..index.metadata.len()).collect(),
0.5,
index.metadata.len(),
);
let published: Vec<usize> = evaluator
.postfilter(ranked.into_iter().map(|i| (i, index.metadata_of(i))))
.take(2)
.collect();
assert_eq!(published.len(), 2);
for key in published {
assert!(evaluator.evaluate(index.metadata_of(key)));
}
}
#[test]
fn strategy_selection_drives_the_public_surface() {
let index = corpus();
let narrow = FilterEvaluator::new(Filter::eq("lang", Value::String("go".into()))).unwrap();
let selector = StrategySelector::new();
assert_eq!(selector.choose(&narrow), FilterStrategy::PreFilter);
assert_eq!(
selector.choose_with_index(&narrow, &index.index),
FilterStrategy::PreFilter
);
}