#![allow(clippy::cast_possible_truncation)]
use parking_lot::RwLock;
use rustc_hash::FxHashMap;
pub(crate) const NUM_SHARDS: usize = 16;
#[derive(Debug, Default)]
struct IndexShard {
entries: FxHashMap<u64, usize>,
}
#[derive(Debug)]
pub struct ShardedIndex {
shards: [RwLock<IndexShard>; NUM_SHARDS],
}
impl Default for ShardedIndex {
fn default() -> Self {
Self::new()
}
}
impl ShardedIndex {
#[must_use]
pub fn new() -> Self {
Self {
shards: std::array::from_fn(|_| RwLock::new(IndexShard::default())),
}
}
#[must_use]
pub fn from_hashmap(map: FxHashMap<u64, usize>) -> Self {
let index = Self::new();
for (id, offset) in map {
index.insert(id, offset);
}
index
}
#[inline]
const fn shard_index(id: u64) -> usize {
(id % NUM_SHARDS as u64) as usize
}
pub fn insert(&self, id: u64, offset: usize) {
let shard_idx = Self::shard_index(id);
let mut shard = self.shards[shard_idx].write();
shard.entries.insert(id, offset);
}
#[must_use]
pub fn get(&self, id: u64) -> Option<usize> {
let shard_idx = Self::shard_index(id);
let shard = self.shards[shard_idx].read();
shard.entries.get(&id).copied()
}
#[must_use]
pub fn contains_key(&self, id: u64) -> bool {
let shard_idx = Self::shard_index(id);
let shard = self.shards[shard_idx].read();
shard.entries.contains_key(&id)
}
pub fn remove(&self, id: u64) -> Option<usize> {
let shard_idx = Self::shard_index(id);
let mut shard = self.shards[shard_idx].write();
shard.entries.remove(&id)
}
#[must_use]
pub fn len(&self) -> usize {
self.shards.iter().map(|s| s.read().entries.len()).sum()
}
#[must_use]
#[allow(dead_code)] pub fn is_empty(&self) -> bool {
self.shards.iter().all(|s| s.read().entries.is_empty())
}
#[allow(dead_code)] pub fn clear(&self) {
for shard in &self.shards {
shard.write().entries.clear();
}
}
pub fn replace_all(&self, new_entries: FxHashMap<u64, usize>) {
let mut guards: Vec<_> = self.shards.iter().map(RwLock::write).collect();
for guard in &mut guards {
guard.entries.clear();
}
for (id, offset) in new_entries {
let shard_idx = Self::shard_index(id);
guards[shard_idx].entries.insert(id, offset);
}
}
#[must_use]
pub fn keys(&self) -> Vec<u64> {
let mut keys = Vec::with_capacity(self.len());
for shard in &self.shards {
let guard = shard.read();
keys.extend(guard.entries.keys().copied());
}
keys
}
#[must_use]
pub fn to_hashmap(&self) -> FxHashMap<u64, usize> {
let mut map = FxHashMap::default();
map.reserve(self.len());
for shard in &self.shards {
let guard = shard.read();
for (&id, &offset) in &guard.entries {
map.insert(id, offset);
}
}
map
}
#[must_use]
#[allow(dead_code)] pub fn max_offset(&self) -> Option<usize> {
let mut max = None;
for shard in &self.shards {
let guard = shard.read();
for &offset in guard.entries.values() {
max = Some(max.map_or(offset, |m: usize| m.max(offset)));
}
}
max
}
#[allow(dead_code)] pub fn reserve(&self, additional: usize) {
let per_shard = additional / NUM_SHARDS + 1;
for shard in &self.shards {
shard.write().entries.reserve(per_shard);
}
}
}