use std::collections::{HashMap, VecDeque};
use std::hash::Hash;
pub const DEFAULT_CAPACITY: usize = 64;
#[derive(Debug, Clone)]
pub struct DirBatch<K: Eq + Hash + Clone, E> {
map: HashMap<K, Vec<E>>,
order: VecDeque<K>,
capacity: usize,
}
impl<K: Eq + Hash + Clone, E> DirBatch<K, E> {
pub fn new(capacity: usize) -> Self {
Self {
map: HashMap::new(),
order: VecDeque::new(),
capacity: capacity.max(1),
}
}
pub fn stage(&mut self, key: K, entry: E) -> Option<(K, Vec<E>)> {
if let Some(batch) = self.map.get_mut(&key) {
batch.push(entry);
return None;
}
let victim = if self.map.len() >= self.capacity {
self.order
.pop_front()
.and_then(|old| self.map.remove(&old).map(|entries| (old, entries)))
} else {
None
};
self.order.push_back(key.clone());
self.map.insert(key, vec![entry]);
victim
}
pub fn peek(&self, key: &K) -> Option<&Vec<E>> {
self.map.get(key)
}
pub fn take(&mut self, key: &K) -> Option<Vec<E>> {
let entries = self.map.remove(key)?;
if let Some(pos) = self.order.iter().position(|k| k == key) {
self.order.remove(pos);
}
Some(entries)
}
pub fn drain_all(&mut self) -> Vec<(K, Vec<E>)> {
self.order.clear();
self.map.drain().collect()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn stages_without_eviction_under_capacity() {
let mut b: DirBatch<u32, &str> = DirBatch::new(4);
assert!(b.stage(1, "a").is_none());
assert!(b.stage(1, "b").is_none());
assert!(b.stage(2, "c").is_none());
let mut full: DirBatch<u32, u32> = DirBatch::new(1);
assert!(full.stage(7, 0).is_none());
assert!(full.stage(7, 1).is_none());
assert_eq!(full.take(&7), Some(vec![0, 1]));
}
#[test]
fn evicts_oldest_when_a_new_dir_overflows() {
let mut b: DirBatch<u32, u32> = DirBatch::new(2);
assert!(b.stage(1, 10).is_none());
assert!(b.stage(2, 20).is_none());
let victim = b.stage(3, 30);
assert_eq!(victim, Some((1, vec![10])));
assert_eq!(b.take(&2), Some(vec![20]));
assert_eq!(b.take(&3), Some(vec![30]));
assert!(b.is_empty());
}
#[test]
fn drain_all_returns_everything() {
let mut b: DirBatch<u32, u32> = DirBatch::new(8);
b.stage(1, 10);
b.stage(1, 11);
b.stage(2, 20);
let mut all = b.drain_all();
all.sort();
assert_eq!(all, vec![(1, vec![10, 11]), (2, vec![20])]);
assert!(b.is_empty());
}
}