use super::partition::Partition;
use crate::index::Cursor as CursorTrait;
use commonware_runtime::telemetry::metrics::{Counter, Gauge};
use std::{
collections::{btree_map, hash_map, BTreeMap, HashMap},
ops::Range,
};
const MUST_CALL_NEXT: &str = "must call Cursor::next()";
const NO_ACTIVE_ITEM: &str = "no active item in Cursor";
enum State {
NeedNext { from: usize },
Active { offset: usize },
Done,
}
enum Backing<'a, K: Ord + Copy, V> {
Soa {
partition: &'a mut Partition<K, V>,
key: K,
run: Range<usize>,
},
Spilled {
spilled: &'a mut HashMap<usize, BTreeMap<K, Vec<V>>>,
partition: usize,
key: K,
},
}
impl<K: Ord + Copy, V> Backing<'_, K, V> {
fn len(&self) -> usize {
match self {
Self::Soa { run, .. } => run.len(),
Self::Spilled {
spilled,
partition,
key,
} => spilled
.get(partition)
.and_then(|inner| inner.get(key))
.map_or(0, Vec::len),
}
}
fn get(&self, off: usize) -> &V {
match self {
Self::Soa { partition, run, .. } => partition.value_at(run.start + off),
Self::Spilled {
spilled,
partition,
key,
} => &spilled
.get(partition)
.and_then(|inner| inner.get(key))
.expect("active cursor must reference a present key")[off],
}
}
fn set(&mut self, off: usize, value: V) {
match self {
Self::Soa { partition, run, .. } => partition.set(run.start + off, value),
Self::Spilled {
spilled,
partition,
key,
} => {
spilled
.get_mut(partition)
.and_then(|inner| inner.get_mut(key))
.expect("active cursor must reference a present key")[off] = value;
}
}
}
fn insert(&mut self, off: usize, value: V) -> bool {
match self {
Self::Soa {
partition,
key,
run,
} => {
#[allow(unstable_name_collisions)]
let created = run.is_empty(); partition.insert_at(run.start + off, *key, value);
run.end += 1;
created
}
Self::Spilled {
spilled,
partition,
key,
} => match spilled.entry(*partition).or_default().entry(*key) {
btree_map::Entry::Occupied(mut run) => {
run.get_mut().insert(off, value);
false
}
btree_map::Entry::Vacant(run) => {
run.insert(vec![value]);
true
}
},
}
}
fn remove(&mut self, off: usize) -> bool {
match self {
Self::Soa { partition, run, .. } => {
partition.remove(run.start + off);
run.end -= 1;
#[allow(unstable_name_collisions)]
run.is_empty()
}
Self::Spilled {
spilled,
partition,
key,
} => {
let hash_map::Entry::Occupied(mut part) = spilled.entry(*partition) else {
unreachable!("active cursor must reference a present partition")
};
let btree_map::Entry::Occupied(mut run) = part.get_mut().entry(*key) else {
unreachable!("active cursor must reference a present key")
};
run.get_mut().remove(off);
if !run.get().is_empty() {
return false;
}
run.remove();
if part.get().is_empty() {
part.remove();
}
true
}
}
}
}
pub struct Cursor<'a, K: Ord + Copy, V> {
backing: Backing<'a, K, V>,
state: State,
keys: &'a Gauge,
items: &'a Gauge,
pruned: &'a Counter,
}
impl<'a, K: Ord + Copy, V> Cursor<'a, K, V> {
pub(super) const fn soa(
partition: &'a mut Partition<K, V>,
key: K,
run: Range<usize>,
keys: &'a Gauge,
items: &'a Gauge,
pruned: &'a Counter,
) -> Self {
Self {
backing: Backing::Soa {
partition,
key,
run,
},
state: State::NeedNext { from: 0 },
keys,
items,
pruned,
}
}
pub(super) const fn spilled(
spilled: &'a mut HashMap<usize, BTreeMap<K, Vec<V>>>,
partition: usize,
key: K,
keys: &'a Gauge,
items: &'a Gauge,
pruned: &'a Counter,
) -> Self {
Self {
backing: Backing::Spilled {
spilled,
partition,
key,
},
state: State::NeedNext { from: 0 },
keys,
items,
pruned,
}
}
}
impl<K: Ord + Copy + Send + Sync, V: Send + Sync> CursorTrait for Cursor<'_, K, V> {
type Value = V;
fn next(&mut self) -> Option<&V> {
let off = match self.state {
State::Done => return None,
State::NeedNext { from } => from,
State::Active { offset } => offset + 1,
};
if off >= self.backing.len() {
self.state = State::Done;
return None;
}
self.state = State::Active { offset: off };
Some(self.backing.get(off))
}
fn update(&mut self, value: V) {
match self.state {
State::NeedNext { .. } => panic!("{MUST_CALL_NEXT}"),
State::Done => panic!("{NO_ACTIVE_ITEM}"),
State::Active { offset } => self.backing.set(offset, value),
}
}
fn insert(&mut self, value: V) {
match self.state {
State::NeedNext { .. } => panic!("{MUST_CALL_NEXT}"),
State::Active { offset } => {
self.backing.insert(offset + 1, value);
self.items.inc();
self.state = State::NeedNext { from: offset + 2 };
}
State::Done => {
let end = self.backing.len();
if self.backing.insert(end, value) {
self.keys.inc();
}
self.items.inc();
}
}
}
fn delete(&mut self) {
let offset = match self.state {
State::NeedNext { .. } => panic!("{MUST_CALL_NEXT}"),
State::Done => panic!("{NO_ACTIVE_ITEM}"),
State::Active { offset } => offset,
};
if self.backing.remove(offset) {
self.keys.dec();
}
self.items.dec();
self.pruned.inc();
self.state = State::NeedNext { from: offset };
}
}