use std::collections::HashMap;
use std::collections::HashSet;
use std::rc::Rc;
struct Entry<T> {
value: Rc<T>,
generation: u64,
}
impl<T> Entry<T> {
fn shared(&self) -> Self {
Self {
value: Rc::clone(&self.value),
generation: self.generation,
}
}
}
const MAX_VIDEO_FRAMES: usize = 512;
const MAX_IMAGES: usize = 128;
const MAX_RETIRED: usize = 128;
#[cfg_attr(not(target_arch = "wasm32"), allow(dead_code))]
pub(crate) const DEFAULT_VIDEO_FRAME_BYTES: u64 = 1 << 30;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct WorkingSetOverBudget {
pub(crate) held: u64,
pub(crate) requested: u64,
pub(crate) budget: u64,
}
impl WorkingSetOverBudget {
pub(crate) const fn needed(&self) -> u64 {
self.held.saturating_add(self.requested)
}
}
pub(crate) struct FrameStore<T> {
frames: HashMap<(String, i64), Entry<T>>,
images: HashMap<String, Entry<T>>,
retired: HashMap<u64, Entry<T>>,
pinned_frames: HashSet<(String, i64)>,
next_generation: u64,
size: fn(&T) -> u64,
video_byte_budget: u64,
}
impl<T> Default for FrameStore<T> {
fn default() -> Self {
Self::sized(|_| 0, u64::MAX)
}
}
impl<T> FrameStore<T> {
pub(crate) fn sized(size: fn(&T) -> u64, video_byte_budget: u64) -> Self {
Self {
frames: HashMap::new(),
images: HashMap::new(),
retired: HashMap::new(),
pinned_frames: HashSet::new(),
next_generation: 1,
size,
video_byte_budget,
}
}
pub(crate) fn set_video_byte_budget(&mut self, bytes: u64) {
self.video_byte_budget = bytes;
self.enforce_video_budget(None);
}
#[cfg(test)]
pub(crate) fn video_bytes(&self) -> u64 {
let mut counted = HashSet::new();
self.frames
.values()
.filter(|entry| counted.insert(entry.generation))
.fold(0_u64, |total, entry| {
total.saturating_add((self.size)(entry.value.as_ref()))
})
}
pub(crate) fn pin(&mut self, keys: impl IntoIterator<Item = (String, i64)>) {
self.pinned_frames.clear();
self.pinned_frames.extend(keys);
}
pub(crate) fn unpin(&mut self) {
self.pinned_frames.clear();
}
fn pinned_generations(
&self,
frame: Option<&(String, i64)>,
image: Option<&str>,
) -> HashSet<u64> {
self.pinned_frames
.iter()
.filter(|key| Some(*key) != frame)
.filter_map(|key| match self.frames.get(key) {
Some(entry) => Some(entry.generation),
None if Some(key.0.as_str()) == image => None,
None => self.images.get(&key.0).map(|entry| entry.generation),
})
.collect()
}
fn image_pinned(&self, source: &str) -> bool {
self.pinned_frames
.iter()
.any(|key| key.0 == source && !self.frames.contains_key(key))
}
fn generation_bytes(&self, generations: &HashSet<u64>) -> u64 {
let mut counted = HashSet::new();
self.frames
.values()
.chain(self.images.values())
.filter(|entry| generations.contains(&entry.generation))
.filter(|entry| counted.insert(entry.generation))
.fold(0_u64, |total, entry| {
total.saturating_add((self.size)(entry.value.as_ref()))
})
}
fn admit_pinned(
&self,
others: &HashSet<u64>,
generation: Option<u64>,
requested: u64,
) -> Result<(), WorkingSetOverBudget> {
if others.is_empty() || generation.is_some_and(|generation| others.contains(&generation)) {
return Ok(());
}
let held = self.generation_bytes(others);
if held.saturating_add(requested) > self.video_byte_budget {
return Err(WorkingSetOverBudget {
held,
requested,
budget: self.video_byte_budget,
});
}
Ok(())
}
fn enforce_video_budget(&mut self, keep: Option<u64>) {
let pinned = self.pinned_generations(None, None);
while self.retained_bytes(self.size) > self.video_byte_budget {
let oldest = self
.generations()
.into_iter()
.filter(|generation| Some(*generation) != keep && !pinned.contains(generation))
.min();
let Some(oldest) = oldest else {
break;
};
self.release_generation(oldest);
}
}
fn entry(&mut self, value: T) -> Entry<T> {
let generation = self.next_generation;
self.next_generation = generation.checked_add(1).unwrap_or(1);
if self.next_generation == 1 {
self.frames.clear();
self.images.clear();
self.retired.clear();
}
Entry {
value: Rc::new(value),
generation,
}
}
pub(crate) fn insert(
&mut self,
key: (String, i64),
value: T,
) -> Result<u64, WorkingSetOverBudget> {
if self.pinned_frames.contains(&key) {
let others = self.pinned_generations(Some(&key), None);
self.admit_pinned(&others, None, (self.size)(&value))?;
}
if !self.frames.contains_key(&key) && self.frames.len() >= MAX_VIDEO_FRAMES {
self.evict_oldest_frame();
}
let entry = self.entry(value);
let generation = entry.generation;
if let Some(old) = self.frames.insert(key, entry) {
self.retire(old);
}
self.enforce_video_budget(Some(generation));
Ok(generation)
}
pub(crate) fn alias(
&mut self,
existing: &(String, i64),
requested: (String, i64),
) -> Result<bool, WorkingSetOverBudget> {
let Some(entry) = self.frames.get(existing).map(Entry::shared) else {
return Ok(false);
};
if existing == &requested {
return Ok(true);
}
let generation = entry.generation;
if self.pinned_frames.contains(&requested) {
let others = self.pinned_generations(Some(&requested), None);
self.admit_pinned(&others, Some(generation), (self.size)(entry.value.as_ref()))?;
}
if !self.frames.contains_key(&requested) && self.frames.len() >= MAX_VIDEO_FRAMES {
self.evict_oldest_frame();
}
if let Some(old) = self.frames.insert(requested, entry)
&& old.generation != generation
{
self.retire(old);
}
Ok(true)
}
pub(crate) fn insert_image(
&mut self,
source: String,
value: T,
) -> Result<u64, WorkingSetOverBudget> {
if self.image_pinned(&source) {
let others = self.pinned_generations(None, Some(&source));
self.admit_pinned(&others, None, (self.size)(&value))?;
}
if !self.images.contains_key(&source) && self.images.len() >= MAX_IMAGES {
self.evict_oldest_image();
}
let entry = self.entry(value);
let generation = entry.generation;
if let Some(old) = self.images.insert(source, entry) {
self.retire(old);
}
self.enforce_video_budget(Some(generation));
Ok(generation)
}
#[cfg_attr(not(target_arch = "wasm32"), allow(dead_code))]
pub(crate) fn holds_image(&self, source: &str, generation: u64) -> bool {
self.images
.get(source)
.is_some_and(|entry| entry.generation == generation)
}
pub(crate) fn get(&self, source: &str, time: i64) -> Option<&T> {
self.lookup(source, time).map(|entry| entry.value.as_ref())
}
#[cfg_attr(not(target_arch = "wasm32"), allow(dead_code))]
pub(crate) fn generation(&self, source: &str, time: i64) -> Option<u64> {
self.lookup(source, time).map(|entry| entry.generation)
}
fn lookup(&self, source: &str, time: i64) -> Option<&Entry<T>> {
self.frames
.get(&(source.to_owned(), time))
.or_else(|| self.images.get(source))
}
pub(crate) fn remove(&mut self, source: &str, time: i64) -> bool {
self.frames.remove(&(source.to_owned(), time)).is_some()
}
pub(crate) fn len(&self) -> usize {
self.generations().len()
}
pub(crate) fn retired_len(&self) -> usize {
self.retired.len()
}
pub(crate) fn retained_bytes(&self, bytes: impl Fn(&T) -> u64) -> u64 {
let mut counted = HashSet::new();
self.frames
.values()
.chain(self.images.values())
.chain(self.retired.values())
.filter(|entry| counted.insert(entry.generation))
.fold(0_u64, |total, entry| {
total.saturating_add(bytes(entry.value.as_ref()))
})
}
pub(crate) fn release_generation(&mut self, generation: u64) -> bool {
let retired = self.retired.remove(&generation).is_some();
let before_frames = self.frames.len();
self.frames
.retain(|_, entry| entry.generation != generation);
let before_images = self.images.len();
self.images
.retain(|_, entry| entry.generation != generation);
retired || before_frames != self.frames.len() || before_images != self.images.len()
}
pub(crate) fn clear(&mut self) {
self.frames.clear();
self.images.clear();
self.retired.clear();
}
pub(crate) fn retain_values(&mut self, keep: impl Fn(&T) -> bool) {
self.frames.retain(|_, entry| keep(&entry.value));
self.images.retain(|_, entry| keep(&entry.value));
self.retired.retain(|_, entry| keep(&entry.value));
}
#[cfg_attr(not(target_arch = "wasm32"), allow(dead_code))]
pub(crate) fn retain_entries(&mut self, keep: impl Fn(&str, &T) -> bool) {
self.frames
.retain(|(source, _), entry| keep(source, &entry.value));
self.images
.retain(|source, entry| keep(source, &entry.value));
}
pub(crate) fn retain_sources(&mut self, keep: impl Fn(&str) -> bool) {
self.frames.retain(|(source, _), _| keep(source));
self.images.retain(|source, _| keep(source));
}
#[cfg_attr(not(target_arch = "wasm32"), allow(dead_code))]
pub(crate) fn release_source(&mut self, source: &str) -> bool {
let before = self.frames.len() + self.images.len();
self.retain_sources(|kept| kept != source);
before != self.frames.len() + self.images.len()
}
pub(crate) fn trim_before(&mut self, source: &str, before: i64) {
self.frames
.retain(|(kept, time), _| kept != source || *time >= before);
}
fn retire(&mut self, entry: Entry<T>) {
while self.retired.len() >= MAX_RETIRED {
let Some(oldest) = self.retired.keys().min().copied() else {
break;
};
self.retired.remove(&oldest);
}
self.retired.insert(entry.generation, entry);
}
fn generations(&self) -> HashSet<u64> {
self.frames
.values()
.chain(self.images.values())
.chain(self.retired.values())
.map(|entry| entry.generation)
.collect()
}
fn evict_oldest_frame(&mut self) {
let pinned = self.pinned_generations(None, None);
let key = self
.frames
.iter()
.filter(|(_, entry)| !pinned.contains(&entry.generation))
.min_by_key(|(_, entry)| entry.generation)
.map(|(key, _)| key.clone());
if let Some(key) = key {
self.frames.remove(&key);
}
}
fn evict_oldest_image(&mut self) {
let pinned = self.pinned_generations(None, None);
let key = self
.images
.iter()
.filter(|(_, entry)| !pinned.contains(&entry.generation))
.min_by_key(|(_, entry)| entry.generation)
.map(|(key, _)| key.clone());
if let Some(key) = key {
self.images.remove(&key);
}
}
}
#[cfg(test)]
#[path = "frame_store_tests.rs"]
mod tests;