use super::{FastMap, FastSet, RowHandle};
pub enum KeyState<K> {
Dense(DenseKeys),
Hashed {
counts: FastMap<K, usize>,
left_buckets: FastMap<K, Vec<RowHandle>>,
touched: Vec<K>,
marked: FastSet<K>,
},
}
impl<K: Eq + std::hash::Hash + Clone + 'static> KeyState<K> {
pub fn new() -> Self {
if std::any::TypeId::of::<K>() == std::any::TypeId::of::<usize>() {
Self::Dense(DenseKeys::new())
} else {
Self::Hashed {
counts: FastMap::default(),
left_buckets: FastMap::default(),
touched: Vec::new(),
marked: FastSet::default(),
}
}
}
pub fn add(&mut self, key: &K, n: usize) -> usize {
match self {
Self::Dense(state) => {
let raw = dense_raw(key);
state.add(raw, n)
}
Self::Hashed {
counts,
touched,
marked,
..
} => {
if !marked.contains(key) {
marked.insert(key.clone());
touched.push(key.clone());
}
let entry = counts.entry(key.clone()).or_insert(0);
*entry += n;
*entry
}
}
}
pub fn take(&mut self, key: &K, n: usize) {
match self {
Self::Dense(state) => state.take(dense_raw(key), n),
Self::Hashed {
counts,
touched,
marked,
..
} => {
if !marked.contains(key) {
marked.insert(key.clone());
touched.push(key.clone());
}
if let Some(entry) = counts.get_mut(key) {
*entry -= n;
if *entry == 0 {
counts.remove(key);
}
}
}
}
}
pub fn alive(&self, key: &K) -> bool {
match self {
Self::Dense(state) => state.alive(dense_raw(key)),
Self::Hashed { counts, .. } => counts.contains_key(key),
}
}
pub fn left_insert(&mut self, key: &K, handle: RowHandle) {
match self {
Self::Dense(state) => state.left_insert(dense_raw(key), handle),
Self::Hashed { left_buckets, .. } => {
left_buckets.entry(key.clone()).or_default().push(handle)
}
}
}
pub fn left_remove(&mut self, key: &K, handle: RowHandle) {
match self {
Self::Dense(state) => state.left_remove(dense_raw(key), handle),
Self::Hashed { left_buckets, .. } => {
if let Some(bucket) = left_buckets.get_mut(key) {
if let Some(pos) = bucket.iter().position(|h| *h == handle) {
bucket.swap_remove(pos);
}
if bucket.is_empty() {
left_buckets.remove(key);
}
}
}
}
}
pub fn left_bucket<'a>(&'a self, key: &'a K) -> Option<&'a [RowHandle]> {
match self {
Self::Dense(state) => state.left_bucket(dense_raw(key)),
Self::Hashed { left_buckets, .. } => left_buckets.get(key).map(Vec::as_slice),
}
}
pub fn begin_event(&mut self) {
match self {
Self::Dense(state) => state.begin_event(),
Self::Hashed { .. } => {}
}
}
pub fn drain_touched(&mut self, out: &mut Vec<K>) {
match self {
Self::Dense(state) => {
for raw in state.touched.drain(..) {
state.in_touched[raw] = false;
out.push(key_from_raw::<K>(raw));
}
}
Self::Hashed {
touched, marked, ..
} => {
marked.clear();
out.append(touched);
}
}
}
pub fn clear(&mut self) {
match self {
Self::Dense(state) => state.clear(),
Self::Hashed {
counts,
left_buckets,
..
} => {
counts.clear();
left_buckets.clear();
}
}
}
}
pub struct DenseKeys {
counts: Vec<u32>,
in_touched: Vec<bool>,
touched: Vec<usize>,
left_buckets: Vec<Vec<RowHandle>>,
}
impl DenseKeys {
pub fn new() -> Self {
Self {
counts: Vec::new(),
in_touched: Vec::new(),
touched: Vec::new(),
left_buckets: Vec::new(),
}
}
fn grow(&mut self, raw: usize) {
if self.counts.len() <= raw {
self.counts.resize(raw + 1, 0);
self.in_touched.resize(raw + 1, false);
self.left_buckets.resize_with(raw + 1, Vec::new);
}
}
#[inline]
pub fn add(&mut self, raw: usize, n: usize) -> usize {
if self.counts.len() <= raw {
self.grow(raw);
}
self.counts[raw] += n as u32;
if !self.in_touched[raw] {
self.in_touched[raw] = true;
self.touched.push(raw);
}
self.counts[raw] as usize
}
#[inline]
pub fn take(&mut self, raw: usize, n: usize) {
if self.counts.len() <= raw {
self.grow(raw);
}
debug_assert!(self.counts[raw] >= n as u32, "count underflow");
self.counts[raw] -= n as u32;
if !self.in_touched[raw] {
self.in_touched[raw] = true;
self.touched.push(raw);
}
}
#[inline]
pub fn alive(&self, raw: usize) -> bool {
self.counts.get(raw).copied().unwrap_or(0) > 0
}
#[inline]
pub fn left_insert(&mut self, raw: usize, handle: RowHandle) {
self.grow(raw);
self.left_buckets[raw].push(handle);
}
#[inline]
pub fn left_remove(&mut self, raw: usize, handle: RowHandle) {
if raw >= self.left_buckets.len() {
return;
}
if let Some(pos) = self.left_buckets[raw].iter().position(|h| *h == handle) {
self.left_buckets[raw].swap_remove(pos);
}
}
#[inline]
pub fn left_bucket(&self, raw: usize) -> Option<&[RowHandle]> {
self.left_buckets.get(raw).map(Vec::as_slice)
}
pub fn begin_event(&mut self) {
for raw in self.touched.drain(..) {
self.in_touched[raw] = false;
}
}
pub fn clear(&mut self) {
self.counts.clear();
self.in_touched.clear();
self.touched.clear();
self.left_buckets.clear();
}
}
fn dense_raw<K: 'static>(key: &K) -> usize {
debug_assert_eq!(
std::any::TypeId::of::<K>(),
std::any::TypeId::of::<usize>(),
"dense counts are usize-keyed"
);
unsafe { *(key as *const K as *const usize) }
}
fn key_from_raw<K: 'static>(raw: usize) -> K {
debug_assert_eq!(
std::any::TypeId::of::<K>(),
std::any::TypeId::of::<usize>(),
"dense counts are usize-keyed"
);
unsafe { std::ptr::read(&raw as *const usize as *const K) }
}