use crate::index::Cursor as CursorTrait;
use commonware_runtime::telemetry::metrics::{Counter, Gauge};
use std::{
collections::HashMap,
hash::{BuildHasher, Hash},
};
pub type Overflow<K, V, S> = HashMap<K, Vec<V>, S>;
#[cold]
#[inline(never)]
pub(super) fn push_displaced<K: Hash + Eq + Copy, V, S: BuildHasher>(
overflow: &mut Overflow<K, V, S>,
key: K,
old: V,
) {
overflow.entry(key).or_default().push(old);
}
#[derive(Clone, Copy)]
enum Position {
Head,
Overflow(usize),
}
pub trait IndexEntry<V: Send + Sync>: Send + Sync {
type Key;
fn key(&self) -> &Self::Key;
fn get_mut(&mut self) -> &mut V;
fn remove(self);
}
const MUST_CALL_NEXT: &str = "must call Cursor::next()";
const NO_ACTIVE_ITEM: &str = "no active item in Cursor";
#[derive(Clone, Copy)]
enum State {
NeedNext { from: Option<Position> },
Active { pos: Position },
Done,
EntryRemoved,
}
pub struct Cursor<
'a,
K: Hash + Eq + Copy,
V: Send + Sync,
E: IndexEntry<V, Key = K>,
S: BuildHasher,
> {
entry: Option<E>,
overflow: &'a mut Overflow<K, V, S>,
chain: Option<Vec<V>>,
state: State,
keys: &'a Gauge,
items: &'a Gauge,
pruned: &'a Counter,
}
impl<'a, K: Hash + Eq + Copy, V: Send + Sync, E: IndexEntry<V, Key = K>, S: BuildHasher>
Cursor<'a, K, V, E, S>
{
#[inline]
pub(super) const fn new(
entry: E,
overflow: &'a mut Overflow<K, V, S>,
keys: &'a Gauge,
items: &'a Gauge,
pruned: &'a Counter,
) -> Self {
Self {
entry: Some(entry),
overflow,
chain: None,
state: State::NeedNext { from: None },
keys,
items,
pruned,
}
}
#[inline]
fn head_mut(&mut self) -> &mut V {
self.entry.as_mut().unwrap().get_mut()
}
fn chain_mut(&mut self) -> &mut Vec<V> {
let key = self.entry.as_ref().unwrap().key();
let overflow = &mut self.overflow;
self.chain.get_or_insert_with(|| {
if overflow.is_empty() {
Vec::new()
} else {
overflow.remove(key).unwrap_or_default()
}
})
}
}
impl<
K: Hash + Eq + Copy + Send + Sync,
V: Send + Sync,
E: IndexEntry<V, Key = K>,
S: BuildHasher + Send + Sync,
> CursorTrait for Cursor<'_, K, V, E, S>
{
type Value = V;
#[inline]
fn next(&mut self) -> Option<&V> {
let from = match self.state {
State::Done | State::EntryRemoved => return None,
State::NeedNext { from } => from,
State::Active { pos } => Some(pos),
};
let next = match from {
None => Position::Head,
Some(Position::Head) => match self.chain_mut().len() {
0 => {
self.state = State::Done;
return None;
}
len => Position::Overflow(len - 1),
},
Some(Position::Overflow(0)) => {
self.state = State::Done;
return None;
}
Some(Position::Overflow(i)) => Position::Overflow(i - 1),
};
self.state = State::Active { pos: next };
Some(match next {
Position::Head => self.head_mut(),
Position::Overflow(i) => &self.chain.as_ref().unwrap()[i],
})
}
#[inline]
fn update(&mut self, v: V) {
match self.state {
State::NeedNext { .. } => panic!("{MUST_CALL_NEXT}"),
State::Done | State::EntryRemoved => panic!("{NO_ACTIVE_ITEM}"),
State::Active {
pos: Position::Head,
} => *self.head_mut() = v,
State::Active {
pos: Position::Overflow(i),
} => self.chain.as_mut().unwrap()[i] = v,
}
}
fn insert(&mut self, v: V) {
match self.state {
State::NeedNext { .. } => panic!("{MUST_CALL_NEXT}"),
State::Active { pos } => {
self.items.inc();
let chain = self.chain_mut();
let at = match pos {
Position::Head => chain.len(),
Position::Overflow(i) => i,
};
chain.insert(at, v);
self.state = State::NeedNext {
from: Some(Position::Overflow(at)),
};
}
State::EntryRemoved => {
self.items.inc();
*self.head_mut() = v;
self.state = State::Done;
}
State::Done => {
self.items.inc();
self.chain_mut().insert(0, v);
}
}
}
fn delete(&mut self) {
let pos = match self.state {
State::NeedNext { .. } => panic!("{MUST_CALL_NEXT}"),
State::Done | State::EntryRemoved => panic!("{NO_ACTIVE_ITEM}"),
State::Active { pos } => pos,
};
self.pruned.inc();
self.items.dec();
match pos {
Position::Head => {
if let Some(promoted) = self.chain_mut().pop() {
*self.head_mut() = promoted;
self.state = State::NeedNext { from: None };
} else {
self.state = State::EntryRemoved;
}
}
Position::Overflow(i) => {
self.chain.as_mut().unwrap().remove(i);
self.state = State::NeedNext {
from: Some(Position::Overflow(i)),
};
}
}
}
}
impl<K: Hash + Eq + Copy, V: Send + Sync, E: IndexEntry<V, Key = K>, S: BuildHasher> Drop
for Cursor<'_, K, V, E, S>
{
#[inline]
fn drop(&mut self) {
if matches!(self.state, State::EntryRemoved) {
self.keys.dec();
self.entry.take().unwrap().remove();
} else if let Some(chain) = self.chain.take() {
if !chain.is_empty() {
let key = *self.entry.as_ref().unwrap().key();
self.overflow.insert(key, chain);
}
}
}
}
pub struct Values<'a, K: Hash + Eq, V, S: BuildHasher> {
head: Option<&'a V>,
overflow: OverflowValues<'a, K, V, S>,
}
enum OverflowValues<'a, K: Hash + Eq, V, S: BuildHasher> {
Pending {
overflow: &'a Overflow<K, V, S>,
key: K,
},
Iter(std::iter::Rev<std::slice::Iter<'a, V>>),
}
impl<'a, K: Hash + Eq, V, S: BuildHasher> Values<'a, K, V, S> {
pub(super) fn new(head: Option<&'a V>, overflow: &'a Overflow<K, V, S>, key: K) -> Self {
let overflow = match head {
Some(_) => OverflowValues::Pending { overflow, key },
None => OverflowValues::Iter([].iter().rev()),
};
Self { head, overflow }
}
}
impl<'a, K: Hash + Eq, V, S: BuildHasher> Iterator for Values<'a, K, V, S> {
type Item = &'a V;
fn next(&mut self) -> Option<&'a V> {
if let Some(head) = self.head.take() {
return Some(head);
}
loop {
match &mut self.overflow {
OverflowValues::Pending { overflow, key } => {
let values = overflow.get(key).map_or(&[][..], Vec::as_slice);
self.overflow = OverflowValues::Iter(values.iter().rev());
}
OverflowValues::Iter(values) => return values.next(),
}
}
}
}