use std::collections::{ BTreeMap, HashMap };
use std::hash::Hash;
struct Slot<V>
{
value: V,
bytes: usize,
used: u64,
}
pub( crate ) struct Lru<K, V>
{
slots: HashMap<K, Slot<V>>,
order: BTreeMap<u64, K>,
tick: u64,
bytes: usize,
budget: usize,
}
impl<K: Copy + Eq + Hash, V> Lru<K, V>
{
pub( crate ) fn new( budget: usize ) -> Self
{
Self { slots: HashMap::new(), order: BTreeMap::new(), tick: 0, bytes: 0, budget }
}
pub( crate ) fn get( &mut self, key: &K ) -> Option<&V>
{
let slot = self.slots.get_mut( key )?;
self.order.remove( &slot.used );
self.tick += 1;
slot.used = self.tick;
self.order.insert( self.tick, *key );
Some( &slot.value )
}
pub( crate ) fn insert( &mut self, key: K, value: V, bytes: usize ) -> Vec<V>
{
let mut evicted: Vec<V> = self.remove( &key ).into_iter().collect();
self.tick += 1;
self.order.insert( self.tick, key );
self.slots.insert( key, Slot { value, bytes, used: self.tick } );
self.bytes += bytes;
while self.bytes > self.budget && self.slots.len() > 1
{
let Some( ( _, oldest ) ) = self.order.pop_first() else { break };
if let Some( slot ) = self.slots.remove( &oldest )
{
self.bytes -= slot.bytes;
evicted.push( slot.value );
}
}
evicted
}
fn remove( &mut self, key: &K ) -> Option<V>
{
let slot = self.slots.remove( key )?;
self.order.remove( &slot.used );
self.bytes -= slot.bytes;
Some( slot.value )
}
}
#[ cfg( test ) ]
mod tests
{
use super::Lru;
#[ test ]
fn evicts_the_least_recently_used()
{
let mut lru = Lru::new( 30 );
assert!( lru.insert( 1, 'a', 10 ).is_empty() );
assert!( lru.insert( 2, 'b', 10 ).is_empty() );
assert!( lru.insert( 3, 'c', 10 ).is_empty() );
assert_eq!( lru.get( &1 ), Some( &'a' ) );
assert_eq!( lru.insert( 4, 'd', 10 ), vec![ 'b' ] );
assert_eq!( lru.get( &2 ), None );
assert_eq!( lru.get( &1 ), Some( &'a' ) );
assert_eq!( lru.get( &3 ), Some( &'c' ) );
}
#[ test ]
fn keeps_a_value_larger_than_the_budget()
{
let mut lru = Lru::new( 10 );
lru.insert( 1, 'a', 5 );
assert_eq!( lru.insert( 2, 'b', 50 ), vec![ 'a' ] );
assert_eq!( lru.get( &2 ), Some( &'b' ) );
}
#[ test ]
fn replacing_a_key_returns_the_old_value()
{
let mut lru = Lru::new( 100 );
lru.insert( 1, 'a', 10 );
assert_eq!( lru.insert( 1, 'z', 10 ), vec![ 'a' ] );
assert_eq!( lru.get( &1 ), Some( &'z' ) );
assert_eq!( lru.bytes, 10 );
}
}