Free least recently used

src/lib.rs

@@ -1,6 +1,6 @@
 use std::{
     borrow::Cow,
-    collections::{HashMap, HashSet},
+    collections::HashSet,
     error::Error,
     fmt::{self, Display, Formatter},
     iter,
@@ -15,6 +15,7 @@ use fontdue::{
     layout::{GlyphRasterConfig, Layout},
     Font,
 };
+use recently_used::RecentlyUsedMap;
 use wgpu::{
     BindGroup, BindGroupEntry, BindGroupLayoutEntry, BindingResource, BindingType, BlendState,
     Buffer, BufferBindingType, BufferDescriptor, BufferUsages, ColorTargetState, ColorWrites,
@@ -28,6 +29,8 @@ use wgpu::{
 
 pub use fontdue;
 
+mod recently_used;
+
 #[repr(C)]
 pub struct Color {
     pub r: u8,
@@ -54,7 +57,9 @@ impl Display for PrepareError {
 impl Error for PrepareError {}
 
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub enum RenderError {}
+pub enum RenderError {
+    AtlasFull,
+}
 
 impl Display for RenderError {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
@@ -96,40 +101,22 @@ pub struct Params {
     screen_resolution: Resolution,
 }
 
-fn try_allocate(
+fn try_allocate(atlas: &mut InnerAtlas, width: usize, height: usize) -> Option<Allocation> {
-    atlas: &mut InnerAtlas,
-    layout: &Layout<impl HasColor>,
-    width: usize,
-    height: usize,
-) -> Option<Allocation> {
     let size = size2(width as i32, height as i32);
-    let allocation = atlas.packer.allocate(size);
 
+    loop {
+        let allocation = atlas.packer.allocate(size);
         if allocation.is_some() {
             return allocation;
         }
 
-    // Try to free any allocations not used in the current layout
+        // Try to free least recently used allocation
-    // TODO: use LRU instead
+        let (key, value) = atlas.glyph_cache.entries_least_recently_used().next()?;
-    let used_glyphs = layout
+        atlas
-        .glyphs()
+            .packer
-        .iter()
+            .deallocate(value.atlas_id.expect("cache corrupt"));
-        .map(|gp| gp.key)
+        atlas.glyph_cache.remove(&key);
-        .collect::<HashSet<_>>();
-
-    atlas.glyph_cache.retain(|key, details| {
-        if used_glyphs.contains(&key) {
-            true
-        } else {
-            if let Some(atlas_id) = details.atlas_id {
-                atlas.packer.deallocate(atlas_id)
     }
-            false
-        }
-    });
-
-    // Attempt to reallocate
-    atlas.packer.allocate(size)
 }
 
 struct InnerAtlas {
@@ -138,7 +125,7 @@ struct InnerAtlas {
     packer: BucketedAtlasAllocator,
     width: u32,
     height: u32,
-    glyph_cache: HashMap<GlyphRasterConfig, GlyphDetails>,
+    glyph_cache: RecentlyUsedMap<GlyphRasterConfig, GlyphDetails>,
     params: Params,
     params_buffer: Buffer,
 }
@@ -183,7 +170,7 @@ impl TextAtlas {
             ..Default::default()
         });
 
-        let glyph_cache = HashMap::new();
+        let glyph_cache = RecentlyUsedMap::new();
 
         // Create a render pipeline to use for rendering later
         let shader = device.create_shader_module(&ShaderModuleDescriptor {
@@ -340,6 +327,7 @@ pub struct TextRenderer {
     index_buffer_size: u64,
     vertices_to_render: u32,
     atlas: TextAtlas,
+    glyphs_in_use: HashSet<GlyphRasterConfig>,
 }
 
 impl TextRenderer {
@@ -367,6 +355,7 @@ impl TextRenderer {
             index_buffer_size,
             vertices_to_render: 0,
             atlas: atlas.clone(),
+            glyphs_in_use: HashSet::new(),
         }
     }
 
@@ -402,8 +391,12 @@ impl TextRenderer {
         }
         let mut upload_bounds = None::<UploadBounds>;
 
+        self.glyphs_in_use.clear();
+
         for layout in layouts.iter() {
             for glyph in layout.glyphs() {
+                self.glyphs_in_use.insert(glyph.key);
+
                 let already_on_gpu = self
                     .atlas
                     .inner
@@ -411,6 +404,7 @@ impl TextRenderer {
                     .expect("atlas locked")
                     .glyph_cache
                     .contains_key(&glyph.key);
+
                 if already_on_gpu {
                     continue;
                 }
@@ -422,8 +416,7 @@ impl TextRenderer {
 
                 let (gpu_cache, atlas_id) = if glyph.char_data.rasterize() {
                     // Find a position in the packer
-                    let allocation =
+                    let allocation = match try_allocate(&mut atlas, metrics.width, metrics.height) {
-                        match try_allocate(&mut atlas, layout, metrics.width, metrics.height) {
                         Some(a) => a,
                         None => return Err(PrepareError::AtlasFull),
                     };
@@ -467,6 +460,7 @@ impl TextRenderer {
                     (GpuCache::SkipRasterization, None)
                 };
 
+                if !atlas.glyph_cache.contains_key(&glyph.key) {
                     atlas.glyph_cache.insert(
                         glyph.key,
                         GlyphDetails {
@@ -478,6 +472,7 @@ impl TextRenderer {
                     );
                 }
             }
+        }
 
         if let Some(ub) = upload_bounds {
             let atlas = self.atlas.inner.read().expect("atlas locked");
@@ -599,11 +594,19 @@ impl TextRenderer {
         Ok(())
     }
 
-    pub fn render<'pass>(&'pass mut self, pass: &mut RenderPass<'pass>) -> Result<(), ()> {
+    pub fn render<'pass>(&'pass mut self, pass: &mut RenderPass<'pass>) -> Result<(), RenderError> {
         if self.vertices_to_render == 0 {
             return Ok(());
         }
 
+        // Validate that glyphs haven't been evicted from cache since `prepare`
+        let atlas = self.atlas.inner.read().expect("atlas locked");
+        for glyph in self.glyphs_in_use.iter() {
+            if !atlas.glyph_cache.contains_key(glyph) {
+                return Err(RenderError::AtlasFull);
+            }
+        }
+
         pass.set_pipeline(&self.atlas.pipeline);
         pass.set_bind_group(0, &self.atlas.bind_group, &[]);
         pass.set_vertex_buffer(0, self.vertex_buffer.slice(..));

src/recently_used.rs

@@ -0,0 +1,127 @@
+use std::{
+    borrow::Borrow,
+    collections::{hash_map::Entry, HashMap},
+    hash::Hash,
+};
+
+struct RecentlyUsedItem<V> {
+    entry_idx: usize,
+    value: V,
+}
+
+pub struct RecentlyUsedMap<K: Clone + Copy + Eq + Hash, V> {
+    keys: Vec<K>,
+    map: HashMap<K, RecentlyUsedItem<V>>,
+}
+
+impl<K: Clone + Copy + Eq + Hash, V> RecentlyUsedMap<K, V> {
+    pub fn new() -> Self {
+        Self::with_capacity(0)
+    }
+
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            keys: Vec::with_capacity(capacity),
+            map: HashMap::with_capacity(capacity),
+        }
+    }
+
+    pub fn insert(&mut self, key: K, value: V) {
+        let new_idx = self.keys.len();
+        match self.map.entry(key) {
+            Entry::Occupied(mut occupied) => {
+                let old = occupied.insert(RecentlyUsedItem {
+                    entry_idx: new_idx,
+                    value,
+                });
+                let removed = self.keys.remove(old.entry_idx);
+                self.keys.push(removed);
+            }
+            Entry::Vacant(vacant) => {
+                vacant.insert(RecentlyUsedItem {
+                    entry_idx: new_idx,
+                    value,
+                });
+                self.keys.push(key);
+            }
+        }
+    }
+
+    pub fn remove<Q: ?Sized>(&mut self, key: &Q)
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        if let Some(entry) = self.map.remove(key.borrow()) {
+            self.keys.remove(entry.entry_idx);
+        }
+    }
+
+    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        self.map.get(k).map(|item| &item.value)
+    }
+
+    pub fn contains_key<Q: ?Sized>(&self, k: &Q) -> bool
+    where
+        K: Borrow<Q>,
+        Q: Hash + Eq,
+    {
+        self.map.contains_key(k)
+    }
+
+    pub fn entries_least_recently_used(&self) -> impl Iterator<Item = (K, &V)> + '_ {
+        self.keys.iter().map(|k| (*k, &self.map[k].value))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn insert() {
+        let mut rus = RecentlyUsedMap::new();
+        rus.insert("a", ());
+        rus.insert("b", ());
+        assert_eq!(
+            rus.entries_least_recently_used()
+                .map(|(k, _)| k)
+                .collect::<String>(),
+            "ab"
+        );
+    }
+
+    #[test]
+    fn reinsert() {
+        let mut rus = RecentlyUsedMap::new();
+        rus.insert("a", ());
+        rus.insert("b", ());
+        rus.insert("c", ());
+        rus.insert("a", ());
+        assert_eq!(
+            rus.entries_least_recently_used()
+                .map(|(k, _)| k)
+                .collect::<String>(),
+            "bca"
+        );
+    }
+
+    #[test]
+    fn remove() {
+        let mut rus = RecentlyUsedMap::new();
+        rus.insert("a", ());
+        rus.insert("b", ());
+        rus.remove("a");
+        rus.remove("c");
+        assert_eq!(
+            rus.entries_least_recently_used()
+                .map(|(k, _)| k)
+                .collect::<String>(),
+            "b"
+        );
+    }
+}