use crate::lru::{swap_value, CacheError, RawLRU};
use crate::{Cache, DefaultEvictCallback, DefaultHashBuilder, KeyRef, PutResult};
use core::borrow::Borrow;
use core::hash::{BuildHasher, Hash};
pub struct SegmentedCacheBuilder<FH = DefaultHashBuilder, RH = DefaultHashBuilder> {
probationary_size: usize,
protected_size: usize,
probationary_hasher: Option<RH>,
protected_hasher: Option<FH>,
}
impl Default for SegmentedCacheBuilder {
fn default() -> Self {
Self {
probationary_size: 0,
protected_size: 0,
probationary_hasher: Some(DefaultHashBuilder::default()),
protected_hasher: Some(DefaultHashBuilder::default()),
}
}
}
impl SegmentedCacheBuilder {
pub fn new(probationary_size: usize, protected_size: usize) -> Self {
Self {
probationary_size,
protected_size,
probationary_hasher: Some(DefaultHashBuilder::default()),
protected_hasher: Some(DefaultHashBuilder::default()),
}
}
}
impl<FH: BuildHasher, RH: BuildHasher> SegmentedCacheBuilder<FH, RH> {
pub fn set_probationary_size(self, size: usize) -> Self {
SegmentedCacheBuilder {
probationary_size: size,
protected_size: self.protected_size,
probationary_hasher: self.probationary_hasher,
protected_hasher: self.protected_hasher,
}
}
pub fn set_protected_size(self, size: usize) -> Self {
SegmentedCacheBuilder {
probationary_size: self.probationary_size,
protected_size: size,
probationary_hasher: self.probationary_hasher,
protected_hasher: self.protected_hasher,
}
}
pub fn set_probationary_hasher<NRH: BuildHasher>(
self,
hasher: NRH,
) -> SegmentedCacheBuilder<FH, NRH> {
SegmentedCacheBuilder {
probationary_size: self.probationary_size,
protected_size: self.protected_size,
probationary_hasher: Some(hasher),
protected_hasher: self.protected_hasher,
}
}
pub fn set_protected_hasher<NFH: BuildHasher>(
self,
hasher: NFH,
) -> SegmentedCacheBuilder<NFH, RH> {
SegmentedCacheBuilder {
probationary_size: self.probationary_size,
protected_size: self.protected_size,
probationary_hasher: self.probationary_hasher,
protected_hasher: Some(hasher),
}
}
pub fn finalize<K: Hash + Eq, V>(self) -> Result<SegmentedCache<K, V, FH, RH>, CacheError> {
if self.protected_size == 0 {
return Err(CacheError::InvalidSize(0));
}
if self.probationary_size == 0 {
return Err(CacheError::InvalidSize(0));
}
Ok(SegmentedCache {
probationary_size: self.probationary_size,
probationary: RawLRU::with_hasher(
self.probationary_size,
self.probationary_hasher.unwrap(),
)
.unwrap(),
protected_size: self.protected_size,
protected: RawLRU::with_hasher(self.protected_size, self.protected_hasher.unwrap())
.unwrap(),
})
}
}
pub struct SegmentedCache<K, V, FH = DefaultHashBuilder, RH = DefaultHashBuilder> {
probationary_size: usize,
probationary: RawLRU<K, V, DefaultEvictCallback, RH>,
protected_size: usize,
protected: RawLRU<K, V, DefaultEvictCallback, FH>,
}
impl<K: Hash + Eq + Clone, V: Clone, FH: BuildHasher + Clone, RH: BuildHasher + Clone> Clone
for SegmentedCache<K, V, FH, RH>
{
fn clone(&self) -> Self {
Self {
probationary_size: self.probationary_size,
probationary: self.probationary.clone(),
protected_size: self.protected_size,
protected: self.protected.clone(),
}
}
}
impl<K: Hash + Eq, V> SegmentedCache<K, V> {
pub fn new(probationary_size: usize, protected_size: usize) -> Result<Self, CacheError> {
SegmentedCacheBuilder::new(probationary_size, protected_size).finalize()
}
pub fn builder(probationary_size: usize, protected_size: usize) -> SegmentedCacheBuilder {
SegmentedCacheBuilder::new(probationary_size, protected_size)
}
}
impl<K: Hash + Eq, V, FH: BuildHasher, RH: BuildHasher> SegmentedCache<K, V, FH, RH> {
pub fn from_builder(builder: SegmentedCacheBuilder<FH, RH>) -> Result<Self, CacheError> {
builder.finalize()
}
pub fn put_protected(&mut self, k: K, v: V) -> PutResult<K, V> {
self.protected.put(k, v)
}
pub fn peek_lru_from_probationary(&mut self) -> Option<(&K, &V)> {
self.probationary.peek_lru()
}
pub fn peek_lru_mut_from_probationary(&mut self) -> Option<(&K, &mut V)> {
self.probationary.peek_lru_mut()
}
pub fn peek_mru_from_probationary(&mut self) -> Option<(&K, &V)> {
self.probationary.peek_mru()
}
pub fn peek_mru_mut_from_probationary(&mut self) -> Option<(&K, &mut V)> {
self.probationary.peek_mru_mut()
}
pub fn peek_lru_from_protected(&self) -> Option<(&K, &V)> {
self.protected.peek_lru()
}
pub fn peek_lru_mut_from_protected(&mut self) -> Option<(&K, &mut V)> {
self.protected.peek_lru_mut()
}
pub fn peek_mru_from_protected(&self) -> Option<(&K, &V)> {
self.protected.peek_mru()
}
pub fn peek_mru_mut_from_protected(&mut self) -> Option<(&K, &mut V)> {
self.protected.peek_mru_mut()
}
pub fn remove_lru_from_probationary(&mut self) -> Option<(K, V)> {
self.probationary.remove_lru()
}
pub fn remove_lru_from_protected(&mut self) -> Option<(K, V)> {
self.protected.remove_lru()
}
pub fn protected_len(&self) -> usize {
self.protected.len()
}
pub fn probationary_len(&self) -> usize {
self.probationary.len()
}
pub fn probationary_cap(&self) -> usize {
self.probationary_size
}
pub fn protected_cap(&self) -> usize {
self.protected_size
}
fn move_to_protected<T, Q>(&mut self, k: &Q, v: T) -> Option<T>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
if let Some(ent) = self.probationary.remove_and_return_ent(k) {
match self.protected.put_or_evict_nonnull(ent) {
None => Some(v),
Some(old_ent) => {
self.probationary.put_nonnull(old_ent);
Some(v)
}
}
} else {
None
}
}
}
impl<K: Hash + Eq, V, FH: BuildHasher, RH: BuildHasher> Cache<K, V>
for SegmentedCache<K, V, FH, RH>
{
fn put(&mut self, k: K, mut v: V) -> PutResult<K, V> {
let key_ref = KeyRef { k: &k };
if let Some(ent_ptr) = self.protected.map.get_mut(&key_ref).map(|bks| bks.as_ptr()) {
self.protected.update(&mut v, ent_ptr);
return PutResult::Update(v);
}
if self.probationary.contains(&k) {
return self
.probationary
.remove_and_return_ent(&k)
.and_then(|mut ent| {
unsafe {
let ent_ptr = ent.as_mut();
swap_value(&mut v, ent_ptr);
}
self.protected
.put_or_evict_nonnull(ent)
.map(|evicted_ent| self.probationary.put_nonnull(evicted_ent))
})
.unwrap_or(PutResult::<K, V>::Update(v));
}
self.probationary.put(k, v)
}
fn get<Q>(&mut self, k: &Q) -> Option<&V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
match self
.protected
.get_(k)
{
Some(v) => {
Some(unsafe { core::mem::transmute(v) })
}
None => {
match self.probationary.peek_(k) {
Some(v) => {
self.move_to_protected(k, v)
.map(|v| unsafe { core::mem::transmute(v) })
}
None => None,
}
}
}
}
fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
match self
.protected
.get_mut_(k)
{
Some(v) => Some(unsafe { core::mem::transmute(v) }),
None => {
match self.probationary.peek_mut_(k) {
Some(v) => self
.move_to_protected(k, v)
.map(|v| unsafe { core::mem::transmute(v) }),
None => None,
}
}
}
}
fn peek<Q>(&self, k: &Q) -> Option<&V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.protected.peek(k).or_else(|| self.probationary.peek(k))
}
fn peek_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
match self.protected.peek_mut(k) {
Some(v) => Some(v),
None => self.probationary.peek_mut(k),
}
}
fn contains<Q>(&self, k: &Q) -> bool
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.protected.contains(k) || self.probationary.contains(k)
}
fn remove<Q>(&mut self, k: &Q) -> Option<V>
where
K: Borrow<Q>,
Q: Hash + Eq + ?Sized,
{
self.probationary
.remove(k)
.or_else(|| self.protected.remove(k))
}
fn purge(&mut self) {
self.probationary.purge();
self.protected.purge();
}
fn len(&self) -> usize {
self.protected.len() + self.probationary.len()
}
fn cap(&self) -> usize {
self.protected_size + self.probationary_size
}
fn is_empty(&self) -> bool {
self.protected.is_empty() && self.probationary.is_empty()
}
}