use std::collections::BTreeMap;
use std::ops::Range;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use g_math::fixed_point::FixedPoint;
use crate::constants::SEMANTIC_INDEX_MAX_SLICES;
use crate::metric_tree::{EuclideanMetric, MetricVpTree};
pub struct SliceIndex {
epoch: u64,
pub tree: MetricVpTree<Vec<FixedPoint>>,
}
pub struct SemanticIndexCache {
epoch: AtomicU64,
slices: RwLock<BTreeMap<(usize, usize), Arc<SliceIndex>>>,
}
impl SemanticIndexCache {
pub fn new() -> Self {
Self {
epoch: AtomicU64::new(0),
slices: RwLock::new(BTreeMap::new()),
}
}
pub fn bump(&self) {
self.epoch.fetch_add(1, Ordering::SeqCst);
}
pub fn epoch(&self) -> u64 {
self.epoch.load(Ordering::SeqCst)
}
pub fn get_or_build<F>(&self, dim_range: &Range<usize>, snapshot: F) -> Arc<SliceIndex>
where
F: FnOnce() -> Vec<(String, Vec<FixedPoint>)>,
{
let key = (dim_range.start, dim_range.end);
let current = self.epoch.load(Ordering::SeqCst);
{
let slices = self.slices.read().unwrap_or_else(|e| e.into_inner());
if let Some(idx) = slices.get(&key) {
if idx.epoch == current {
return Arc::clone(idx);
}
}
}
let build_epoch = self.epoch.load(Ordering::SeqCst);
let entries = snapshot();
let index = Arc::new(SliceIndex {
epoch: build_epoch,
tree: MetricVpTree::build(entries, &EuclideanMetric),
});
let mut slices = self.slices.write().unwrap_or_else(|e| e.into_inner());
let now = self.epoch.load(Ordering::SeqCst);
slices.retain(|_, idx| idx.epoch == now);
let entry = slices.entry(key).or_insert_with(|| Arc::clone(&index));
let result = Arc::clone(entry);
while slices.len() > SEMANTIC_INDEX_MAX_SLICES {
let evict = slices
.keys()
.find(|&&k| k != key)
.copied()
.expect("cache over capacity implies a second key exists");
slices.remove(&evict);
}
result
}
pub fn cached_slice_count(&self) -> usize {
self.slices.read().unwrap_or_else(|e| e.into_inner()).len()
}
}
impl Default for SemanticIndexCache {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn fp(v: i32) -> FixedPoint {
FixedPoint::from_int(v)
}
fn snapshot_of(n: usize) -> Vec<(String, Vec<FixedPoint>)> {
(0..n).map(|i| (format!("n{}", i), vec![fp(i as i32)])).collect()
}
#[test]
fn cache_hit_and_epoch_invalidation() {
let cache = SemanticIndexCache::new();
let range = 16..17;
let a = cache.get_or_build(&range, || snapshot_of(3));
let b = cache.get_or_build(&range, || panic!("must be served from cache"));
assert!(Arc::ptr_eq(&a, &b));
cache.bump();
let c = cache.get_or_build(&range, || snapshot_of(4));
assert!(!Arc::ptr_eq(&a, &c));
assert_eq!(c.tree.len(), 4);
}
#[test]
fn distinct_slices_get_distinct_trees() {
let cache = SemanticIndexCache::new();
let a = cache.get_or_build(&(16..18), || snapshot_of(2));
let b = cache.get_or_build(&(16..20), || snapshot_of(5));
assert!(!Arc::ptr_eq(&a, &b));
assert_eq!(cache.cached_slice_count(), 2);
}
#[test]
fn eviction_is_bounded_and_deterministic() {
let cache = SemanticIndexCache::new();
for i in 0..(SEMANTIC_INDEX_MAX_SLICES + 4) {
cache.get_or_build(&(i..(i + 1)), || snapshot_of(1));
}
assert!(cache.cached_slice_count() <= SEMANTIC_INDEX_MAX_SLICES);
let last = SEMANTIC_INDEX_MAX_SLICES + 3;
let again = cache.get_or_build(&(last..(last + 1)), || panic!("evicted the hot slice"));
assert_eq!(again.tree.len(), 1);
}
}