use alloc::collections::BTreeMap;
use alloc::rc::Rc;
use alloc::vec::Vec;
pub const DEFAULT_CACHE_SIZE: i64 = -2000;
struct Entry {
data: Rc<Vec<u8>>,
last_used: u64,
}
pub struct PageCache {
map: BTreeMap<u32, Entry>,
capacity: usize,
clock: u64,
}
impl PageCache {
pub fn new(cache_size: i64, page_size: usize) -> PageCache {
PageCache {
map: BTreeMap::new(),
capacity: capacity_for(cache_size, page_size),
clock: 0,
}
}
pub fn set_cache_size(&mut self, cache_size: i64, page_size: usize) {
self.capacity = capacity_for(cache_size, page_size);
self.evict_to_capacity();
}
pub fn capacity(&self) -> usize {
self.capacity
}
pub fn len(&self) -> usize {
self.map.len()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
fn tick(&mut self) -> u64 {
self.clock = self.clock.wrapping_add(1);
self.clock
}
pub fn get(&mut self, number: u32) -> Option<Rc<Vec<u8>>> {
let t = self.tick();
let entry = self.map.get_mut(&number)?;
entry.last_used = t;
Some(Rc::clone(&entry.data))
}
pub fn insert(&mut self, number: u32, data: Rc<Vec<u8>>) {
let t = self.tick();
if !self.map.contains_key(&number) {
while self.map.len() >= self.capacity {
if !self.evict_one() {
break;
}
}
}
self.map.insert(number, Entry { data, last_used: t });
}
pub fn invalidate(&mut self, number: u32) {
self.map.remove(&number);
}
pub fn clear(&mut self) {
self.map.clear();
}
fn evict_one(&mut self) -> bool {
let victim = self
.map
.iter()
.min_by_key(|(_, e)| e.last_used)
.map(|(&n, _)| n);
match victim {
Some(n) => {
self.map.remove(&n);
true
}
None => false,
}
}
fn evict_to_capacity(&mut self) {
while self.map.len() > self.capacity {
if !self.evict_one() {
break;
}
}
}
}
fn capacity_for(cache_size: i64, page_size: usize) -> usize {
let pages = if cache_size >= 0 {
cache_size as u64
} else {
let kib = (cache_size as i128).unsigned_abs() as u64;
let bytes = kib.saturating_mul(1024);
let ps = page_size.max(1) as u64;
bytes / ps
};
(pages as usize).max(1)
}
#[cfg(test)]
mod tests {
use super::*;
use alloc::vec;
fn page(byte: u8, size: usize) -> Rc<Vec<u8>> {
Rc::new(vec![byte; size])
}
#[test]
fn positive_cache_size_is_a_page_count() {
let c = PageCache::new(10, 4096);
assert_eq!(c.capacity(), 10);
}
#[test]
fn negative_cache_size_is_kib_of_memory() {
let c = PageCache::new(-2000, 4096);
assert_eq!(c.capacity(), 500);
let c = PageCache::new(-2000, 1024);
assert_eq!(c.capacity(), 2000);
}
#[test]
fn capacity_is_at_least_one() {
assert_eq!(PageCache::new(0, 4096).capacity(), 1);
assert_eq!(PageCache::new(-1, 65536).capacity(), 1);
}
#[test]
fn evicts_least_recently_used() {
let mut c = PageCache::new(2, 4096);
c.insert(1, page(1, 4096));
c.insert(2, page(2, 4096));
assert!(c.get(1).is_some());
c.insert(3, page(3, 4096));
assert_eq!(c.len(), 2);
assert!(c.get(1).is_some(), "recently used page 1 survived");
assert!(c.get(2).is_none(), "LRU page 2 was evicted");
assert!(c.get(3).is_some(), "newly inserted page 3 present");
}
#[test]
fn never_exceeds_capacity() {
let mut c = PageCache::new(3, 4096);
for n in 1..=100u32 {
c.insert(n, page(n as u8, 4096));
assert!(c.len() <= 3, "resident set stayed within capacity");
}
assert_eq!(c.len(), 3);
}
#[test]
fn refresh_does_not_grow_or_evict() {
let mut c = PageCache::new(2, 4096);
c.insert(1, page(1, 4096));
c.insert(2, page(2, 4096));
c.insert(1, page(9, 4096));
assert_eq!(c.len(), 2);
assert!(c.get(2).is_some());
assert_eq!(c.get(1).unwrap()[0], 9);
}
#[test]
fn shrinking_capacity_evicts_immediately() {
let mut c = PageCache::new(5, 4096);
for n in 1..=5u32 {
c.insert(n, page(n as u8, 4096));
}
assert_eq!(c.len(), 5);
c.set_cache_size(2, 4096);
assert_eq!(c.capacity(), 2);
assert_eq!(c.len(), 2);
}
#[test]
fn invalidate_and_clear() {
let mut c = PageCache::new(4, 4096);
c.insert(1, page(1, 4096));
c.insert(2, page(2, 4096));
c.invalidate(1);
assert!(c.get(1).is_none());
assert!(c.get(2).is_some());
c.clear();
assert!(c.is_empty());
}
}