use crate::translator::Translator;
use commonware_runtime::{
Metrics,
telemetry::metrics::{Counter, Gauge, MetricsExt as _},
};
use commonware_utils::{Widen, handles::Table};
use core::cmp::Ordering;
use std::collections::VecDeque;
pub struct Index<T: Translator> {
translator: T,
table: Table<u64>,
keys: VecDeque<T::Key>,
start: u64,
items: Gauge,
pruned: Counter,
}
impl<T: Translator> Index<T> {
pub fn new(ctx: impl Metrics, translator: T, start: u64, capacity: usize) -> Self {
Self {
translator,
table: Table::with_capacity(capacity),
keys: VecDeque::with_capacity(capacity),
start,
items: ctx.gauge("items", "Number of positions in the index"),
pruned: ctx.counter("pruned", "Number of positions pruned from the index"),
}
}
fn offset(start: u64, position: u64) -> usize {
usize::try_from(position - start).expect("retained positions must fit in usize")
}
fn key_at(keys: &VecDeque<T::Key>, start: u64, position: u64) -> T::Key {
keys[Self::offset(start, position)]
}
#[inline]
pub fn insert(&mut self, key: &[u8], position: u64) {
assert!(
position >= self.start,
"position {position} must be at or above the index start {}",
self.start
);
let translated = self.translator.transform(key);
let offset = Self::offset(self.start, position);
match offset.cmp(&self.keys.len()) {
Ordering::Equal => self.keys.push_back(translated),
Ordering::Less => self.keys[offset] = translated,
Ordering::Greater => self.keys.resize(offset + 1, translated),
}
self.table
.insert(self.translator.hash_one(translated), position, |position| {
self.translator
.hash_one(Self::key_at(&self.keys, self.start, position))
});
self.items.inc();
}
#[inline]
pub fn get(&self, key: &[u8]) -> impl Iterator<Item = u64> + '_ {
let translated = self.translator.transform(key);
self.table
.get(self.translator.hash_one(translated))
.filter(move |&position| Self::key_at(&self.keys, self.start, position) == translated)
}
pub fn prune(&mut self, start: u64) {
if start <= self.start {
return;
}
let retired = Self::offset(self.start, start).min(self.keys.len());
let mut removed = 0u64;
for (offset, key) in self.keys.drain(..retired).enumerate() {
let hash = self.translator.hash_one(key);
if self.table.remove(hash, self.start + Widen::widen(offset)) {
removed += 1;
}
}
self.start = start;
self.items.dec_by(removed as i64);
self.pruned.inc_by(removed);
}
pub const fn len(&self) -> usize {
self.table.len()
}
pub const fn is_empty(&self) -> bool {
self.table.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::translator::{EightCap, TwoCap};
use commonware_macros::test_traced;
use commonware_runtime::{Runner as _, deterministic};
use commonware_utils::GOLDEN_RATIO;
use core::hash::BuildHasher as _;
fn positions<T: Translator>(index: &Index<T>, key: &[u8]) -> Vec<u64> {
let mut positions: Vec<u64> = index.get(key).collect();
positions.sort_unstable();
positions
}
#[test_traced]
fn test_get_confirms_translated_key() {
deterministic::Runner::default().start(|context| async move {
let mut index = Index::new(context, TwoCap, 0, 0);
index.insert(b"aa1", 1);
index.insert(b"aa2", 3);
index.insert(b"bb1", 4);
assert_eq!(positions(&index, b"aa1"), [1, 3]);
assert_eq!(positions(&index, b"aa9"), [1, 3]);
assert_eq!(positions(&index, b"bb1"), [4]);
assert!(positions(&index, b"cc1").is_empty());
assert_eq!(index.len(), 3);
});
}
#[test_traced]
fn test_get_rejects_candidates_sharing_slot_and_fingerprint() {
deterministic::Runner::default().start(|context| async move {
const INVERSE: u64 = 0xf1de_83e1_9937_733d;
assert_eq!(GOLDEN_RATIO.wrapping_mul(INVERSE), 1);
let k1 = 0x1234_5678_9abc_def0u64;
let k2 = k1.wrapping_add(INVERSE << 16);
let mut index = Index::new(context, EightCap, 0, 0);
let (h1, h2) = (index.translator.hash_one(k1), index.translator.hash_one(k2));
assert_eq!(h2, h1.wrapping_add(1 << 16));
assert_eq!(h1 >> 56, h2 >> 56);
index.insert(&k1.to_be_bytes(), 1);
index.insert(&k2.to_be_bytes(), 3);
assert_eq!(positions(&index, &k1.to_be_bytes()), [1]);
assert_eq!(positions(&index, &k2.to_be_bytes()), [3]);
});
}
#[test_traced]
fn test_prune_retires_exactly_the_pruned_positions() {
deterministic::Runner::default().start(|context| async move {
let mut index = Index::new(context, EightCap, 0, 0);
for (i, position) in [1, 3, 4, 9, 20].into_iter().enumerate() {
index.insert(&[i as u8; 8], position);
}
assert_eq!(index.len(), 5);
assert_eq!(index.keys.len(), 21);
index.prune(2);
assert_eq!(index.len(), 4);
assert_eq!(index.keys.len(), 19);
assert!(positions(&index, &[0u8; 8]).is_empty());
assert_eq!(positions(&index, &[1u8; 8]), [3]);
assert_eq!(index.pruned.get(), 1);
index.prune(1);
assert_eq!(index.len(), 4);
index.prune(9);
assert_eq!(index.len(), 2);
assert_eq!(positions(&index, &[3u8; 8]), [9]);
assert_eq!(positions(&index, &[4u8; 8]), [20]);
index.prune(100);
assert!(index.is_empty());
assert!(index.keys.is_empty());
assert_eq!(index.pruned.get(), 5);
index.insert(&[9u8; 8], 101);
assert_eq!(positions(&index, &[9u8; 8]), [101]);
assert_eq!(index.keys.len(), 2);
});
}
#[test_traced]
fn test_start_above_zero() {
deterministic::Runner::default().start(|context| async move {
let start = 40_967;
let mut index = Index::new(context, EightCap, start, 0);
index.insert(b"first___", start);
index.insert(b"second__", start + 5_000);
assert_eq!(positions(&index, b"first___"), [start]);
index.prune(start + 1);
assert!(positions(&index, b"first___").is_empty());
assert_eq!(positions(&index, b"second__"), [start + 5_000]);
assert_eq!(index.len(), 1);
});
}
#[test_traced]
fn test_growth_rehashes_through_the_keys() {
deterministic::Runner::default().start(|context| async move {
let mut index = Index::new(context, TwoCap, 0, 0);
let key = |i: u32| {
let [hi, lo] = ((i % 5_000) as u16).to_be_bytes();
[hi, lo, (i / 5_000) as u8]
};
for i in 0..20_000u32 {
index.insert(&key(i), 2 * i as u64 + 1);
}
assert_eq!(index.len(), 20_000);
for i in (0..20_000u32).step_by(997) {
assert!(positions(&index, &key(i)).contains(&(2 * i as u64 + 1)));
}
index.prune(20_001);
assert_eq!(index.len(), 10_000);
assert!(!positions(&index, &key(0)).contains(&1));
assert!(positions(&index, &key(10_000)).contains(&20_001));
assert!(positions(&index, &key(19_999)).contains(&39_999));
assert_eq!(index.pruned.get(), 10_000);
});
}
}