use std::num::NonZeroUsize;
use chia_protocol::Bytes32;
use dig_block::{L2Block, L2BlockHeader};
use lru::LruCache;
use parking_lot::RwLock;
pub struct ShardedLruCache<V: Clone> {
shards: Vec<RwLock<LruCache<Bytes32, V>>>,
num_shards: usize,
}
impl<V: Clone> ShardedLruCache<V> {
pub fn new(total_capacity: usize, num_shards: usize) -> Self {
let num_shards = num_shards.max(1);
let per_shard = (total_capacity / num_shards).max(1);
let nz = NonZeroUsize::new(per_shard).expect("per-shard capacity is at least 1");
let shards = (0..num_shards)
.map(|_| RwLock::new(LruCache::new(nz)))
.collect();
Self { shards, num_shards }
}
#[inline]
fn shard_index(&self, key: &Bytes32) -> usize {
let b = key.as_ref()[0] as usize;
if self.num_shards.is_power_of_two() {
b & (self.num_shards - 1)
} else {
b % self.num_shards
}
}
pub fn get_clone(&self, key: &Bytes32) -> Option<V> {
let i = self.shard_index(key);
let mut guard = self.shards[i].write();
guard.get(key).cloned()
}
pub fn insert(&self, key: Bytes32, value: V) {
let i = self.shard_index(&key);
let mut guard = self.shards[i].write();
guard.put(key, value);
}
#[inline]
#[must_use]
pub fn contains(&self, key: &Bytes32) -> bool {
let i = self.shard_index(key);
let guard = self.shards[i].read();
guard.peek(key).is_some()
}
pub fn remove(&self, key: &Bytes32) {
let i = self.shard_index(key);
let mut guard = self.shards[i].write();
let _ = guard.pop(key);
}
}
pub type ShardedBlockCache = ShardedLruCache<L2Block>;
pub type ShardedHeaderCache = ShardedLruCache<L2BlockHeader>;