Introduce `canvas::Cache` grouping

Caches with the same `Group` will share their text
atlas!

10 files changed, 279 insertions, 101 deletions

diff --git a/graphics/src/cache.rs b/graphics/src/cache.rs
new file mode 100644
index 00000000..7106c392
--- /dev/null
+++ b/graphics/src/cache.rs
@@ -0,0 +1,158 @@
+//! Cache computations and efficiently reuse them.
+use std::cell::RefCell;
+use std::fmt;
+use std::sync::atomic::{self, AtomicU64};
+
+/// A simple cache that stores generated values to avoid recomputation.
+///
+/// Keeps track of the last generated value after clearing.
+pub struct Cache<T> {
+    group: Group,
+    state: RefCell<State<T>>,
+}
+
+impl<T> Cache<T> {
+    /// Creates a new empty [`Cache`].
+    pub fn new() -> Self {
+        Cache {
+            group: Group::unique(),
+            state: RefCell::new(State::Empty { previous: None }),
+        }
+    }
+
+    /// Creates a new empty [`Cache`] with the given [`Group`].
+    ///
+    /// Caches within the same group may reuse internal rendering storage.
+    ///
+    /// You should generally group caches that are likely to change
+    /// together.
+    pub fn with_group(group: Group) -> Self {
+        Cache {
+            group,
+            state: RefCell::new(State::Empty { previous: None }),
+        }
+    }
+
+    /// Returns the [`Group`] of the [`Cache`].
+    pub fn group(&self) -> Group {
+        self.group
+    }
+
+    /// Puts the given value in the [`Cache`].
+    ///
+    /// Notice that, given this is a cache, a mutable reference is not
+    /// necessary to call this method. You can safely update the cache in
+    /// rendering code.
+    pub fn put(&self, value: T) {
+        *self.state.borrow_mut() = State::Filled { current: value };
+    }
+
+    /// Returns a reference cell to the internal [`State`] of the [`Cache`].
+    pub fn state(&self) -> &RefCell<State<T>> {
+        &self.state
+    }
+
+    /// Clears the [`Cache`].
+    pub fn clear(&self)
+    where
+        T: Clone,
+    {
+        use std::ops::Deref;
+
+        let previous = match self.state.borrow().deref() {
+            State::Empty { previous } => previous.clone(),
+            State::Filled { current } => Some(current.clone()),
+        };
+
+        *self.state.borrow_mut() = State::Empty { previous };
+    }
+}
+
+/// A cache group.
+///
+/// Caches that share the same group generally change together.
+///
+/// A cache group can be used to implement certain performance
+/// optimizations during rendering, like batching or sharing atlases.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct Group(u64);
+
+impl Group {
+    /// Generates a new unique cache [`Group`].
+    pub fn unique() -> Self {
+        static NEXT: AtomicU64 = AtomicU64::new(0);
+
+        Self(NEXT.fetch_add(1, atomic::Ordering::Relaxed))
+    }
+}
+
+impl<T> fmt::Debug for Cache<T>
+where
+    T: fmt::Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        use std::ops::Deref;
+
+        let state = self.state.borrow();
+
+        match state.deref() {
+            State::Empty { previous } => {
+                write!(f, "Cache::Empty {{ previous: {previous:?} }}")
+            }
+            State::Filled { current } => {
+                write!(f, "Cache::Filled {{ current: {current:?} }}")
+            }
+        }
+    }
+}
+
+impl<T> Default for Cache<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+/// The state of a [`Cache`].
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum State<T> {
+    /// The [`Cache`] is empty.
+    Empty {
+        /// The previous value of the [`Cache`].
+        previous: Option<T>,
+    },
+    /// The [`Cache`] is filled.
+    Filled {
+        /// The current value of the [`Cache`]
+        current: T,
+    },
+}
+
+/// A piece of data that can be cached.
+pub trait Cached: Sized {
+    /// The type of cache produced.
+    type Cache: Clone;
+
+    /// Loads the [`Cache`] into a proper instance.
+    ///
+    /// [`Cache`]: Self::Cache
+    fn load(cache: &Self::Cache) -> Self;
+
+    /// Caches this value, producing its corresponding [`Cache`].
+    ///
+    /// [`Cache`]: Self::Cache
+    fn cache(self, group: Group, previous: Option<Self::Cache>) -> Self::Cache;
+}
+
+#[cfg(debug_assertions)]
+impl Cached for () {
+    type Cache = ();
+
+    fn load(_cache: &Self::Cache) -> Self {}
+
+    fn cache(
+        self,
+        _group: Group,
+        _previous: Option<Self::Cache>,
+    ) -> Self::Cache {
+    }
+}
diff --git a/graphics/src/cached.rs b/graphics/src/cached.rs
deleted file mode 100644
index c0e4e029..00000000
--- a/graphics/src/cached.rs
+++ /dev/null
@@ -1,24 +0,0 @@
-/// A piece of data that can be cached.
-pub trait Cached: Sized {
-    /// The type of cache produced.
-    type Cache: Clone;
-
-    /// Loads the [`Cache`] into a proper instance.
-    ///
-    /// [`Cache`]: Self::Cache
-    fn load(cache: &Self::Cache) -> Self;
-
-    /// Caches this value, producing its corresponding [`Cache`].
-    ///
-    /// [`Cache`]: Self::Cache
-    fn cache(self, previous: Option<Self::Cache>) -> Self::Cache;
-}
-
-#[cfg(debug_assertions)]
-impl Cached for () {
-    type Cache = ();
-
-    fn load(_cache: &Self::Cache) -> Self {}
-
-    fn cache(self, _previous: Option<Self::Cache>) -> Self::Cache {}
-}
diff --git a/graphics/src/geometry.rs b/graphics/src/geometry.rs
index dd07097f..ab4a7a36 100644
--- a/graphics/src/geometry.rs
+++ b/graphics/src/geometry.rs
@@ -18,8 +18,8 @@ pub use text::Text;
 
 pub use crate::gradient::{self, Gradient};
 
+use crate::cache::Cached;
 use crate::core::{self, Size};
-use crate::Cached;
 
 /// A renderer capable of drawing some [`Self::Geometry`].
 pub trait Renderer: core::Renderer {
diff --git a/graphics/src/geometry/cache.rs b/graphics/src/geometry/cache.rs
index 665e996b..d70cee0b 100644
--- a/graphics/src/geometry/cache.rs
+++ b/graphics/src/geometry/cache.rs
@@ -1,8 +1,8 @@
+use crate::cache::{self, Cached};
 use crate::core::Size;
 use crate::geometry::{self, Frame};
-use crate::Cached;
 
-use std::cell::RefCell;
+pub use cache::Group;
 
 /// A simple cache that stores generated geometry to avoid recomputation.
 ///
@@ -12,7 +12,13 @@ pub struct Cache<Renderer>
 where
     Renderer: geometry::Renderer,
 {
-    state: RefCell<State<Renderer::Geometry>>,
+    raw: crate::Cache<Data<<Renderer::Geometry as Cached>::Cache>>,
+}
+
+#[derive(Debug, Clone)]
+struct Data<T> {
+    bounds: Size,
+    geometry: T,
 }
 
 impl<Renderer> Cache<Renderer>
@@ -22,20 +28,25 @@ where
     /// Creates a new empty [`Cache`].
     pub fn new() -> Self {
         Cache {
-            state: RefCell::new(State::Empty { previous: None }),
+            raw: cache::Cache::new(),
+        }
+    }
+
+    /// Creates a new empty [`Cache`] with the given [`Group`].
+    ///
+    /// Caches within the same group may reuse internal rendering storage.
+    ///
+    /// You should generally group caches that are likely to change
+    /// together.
+    pub fn with_group(group: Group) -> Self {
+        Cache {
+            raw: crate::Cache::with_group(group),
         }
     }
 
     /// Clears the [`Cache`], forcing a redraw the next time it is used.
     pub fn clear(&self) {
-        use std::ops::Deref;
-
-        let previous = match self.state.borrow().deref() {
-            State::Empty { previous } => previous.clone(),
-            State::Filled { geometry, .. } => Some(geometry.clone()),
-        };
-
-        *self.state.borrow_mut() = State::Empty { previous };
+        self.raw.clear();
     }
 
     /// Draws geometry using the provided closure and stores it in the
@@ -56,27 +67,30 @@ where
     ) -> Renderer::Geometry {
         use std::ops::Deref;
 
-        let previous = match self.state.borrow().deref() {
-            State::Empty { previous } => previous.clone(),
-            State::Filled {
-                bounds: cached_bounds,
-                geometry,
-            } => {
-                if *cached_bounds == bounds {
-                    return Cached::load(geometry);
+        let state = self.raw.state();
+
+        let previous = match state.borrow().deref() {
+            cache::State::Empty { previous } => {
+                previous.as_ref().map(|data| data.geometry.clone())
+            }
+            cache::State::Filled { current } => {
+                if current.bounds == bounds {
+                    return Cached::load(&current.geometry);
                 }
 
-                Some(geometry.clone())
+                Some(current.geometry.clone())
             }
         };
 
         let mut frame = Frame::new(renderer, bounds);
         draw_fn(&mut frame);
 
-        let geometry = frame.into_geometry().cache(previous);
+        let geometry = frame.into_geometry().cache(self.raw.group(), previous);
         let result = Cached::load(&geometry);
 
-        *self.state.borrow_mut() = State::Filled { bounds, geometry };
+        *state.borrow_mut() = cache::State::Filled {
+            current: Data { bounds, geometry },
+        };
 
         result
     }
@@ -85,16 +99,10 @@ where
 impl<Renderer> std::fmt::Debug for Cache<Renderer>
 where
     Renderer: geometry::Renderer,
+    <Renderer::Geometry as Cached>::Cache: std::fmt::Debug,
 {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let state = self.state.borrow();
-
-        match *state {
-            State::Empty { .. } => write!(f, "Cache::Empty"),
-            State::Filled { bounds, .. } => {
-                write!(f, "Cache::Filled {{ bounds: {bounds:?} }}")
-            }
-        }
+        write!(f, "{:?}", &self.raw)
     }
 }
 
@@ -106,16 +114,3 @@ where
         Self::new()
     }
 }
-
-enum State<Geometry>
-where
-    Geometry: Cached,
-{
-    Empty {
-        previous: Option<Geometry::Cache>,
-    },
-    Filled {
-        bounds: Size,
-        geometry: Geometry::Cache,
-    },
-}
diff --git a/graphics/src/lib.rs b/graphics/src/lib.rs
index 865ebd97..b5ef55e7 100644
--- a/graphics/src/lib.rs
+++ b/graphics/src/lib.rs
@@ -9,10 +9,10 @@
 )]
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
 mod antialiasing;
-mod cached;
 mod settings;
 mod viewport;
 
+pub mod cache;
 pub mod color;
 pub mod compositor;
 pub mod damage;
@@ -27,7 +27,7 @@ pub mod text;
 pub mod geometry;
 
 pub use antialiasing::Antialiasing;
-pub use cached::Cached;
+pub use cache::Cache;
 pub use compositor::Compositor;
 pub use error::Error;
 pub use gradient::Gradient;
diff --git a/renderer/src/fallback.rs b/renderer/src/fallback.rs
index c932de00..5f69b420 100644
--- a/renderer/src/fallback.rs
+++ b/renderer/src/fallback.rs
@@ -428,8 +428,8 @@ where
 mod geometry {
     use super::Renderer;
     use crate::core::{Point, Radians, Rectangle, Size, Vector};
+    use crate::graphics::cache::{self, Cached};
     use crate::graphics::geometry::{self, Fill, Path, Stroke, Text};
-    use crate::graphics::Cached;
 
     impl<A, B> geometry::Renderer for Renderer<A, B>
     where
@@ -483,21 +483,25 @@ mod geometry {
             }
         }
 
-        fn cache(self, previous: Option<Self::Cache>) -> Self::Cache {
+        fn cache(
+            self,
+            group: cache::Group,
+            previous: Option<Self::Cache>,
+        ) -> Self::Cache {
             match (self, previous) {
                 (
                     Self::Primary(geometry),
                     Some(Geometry::Primary(previous)),
-                ) => Geometry::Primary(geometry.cache(Some(previous))),
+                ) => Geometry::Primary(geometry.cache(group, Some(previous))),
                 (Self::Primary(geometry), None) => {
-                    Geometry::Primary(geometry.cache(None))
+                    Geometry::Primary(geometry.cache(group, None))
                 }
                 (
                     Self::Secondary(geometry),
                     Some(Geometry::Secondary(previous)),
-                ) => Geometry::Secondary(geometry.cache(Some(previous))),
+                ) => Geometry::Secondary(geometry.cache(group, Some(previous))),
                 (Self::Secondary(geometry), None) => {
-                    Geometry::Secondary(geometry.cache(None))
+                    Geometry::Secondary(geometry.cache(group, None))
                 }
                 _ => unreachable!(),
             }
diff --git a/tiny_skia/src/geometry.rs b/tiny_skia/src/geometry.rs
index 117daf41..02b6e1b9 100644
--- a/tiny_skia/src/geometry.rs
+++ b/tiny_skia/src/geometry.rs
@@ -1,9 +1,10 @@
 use crate::core::text::LineHeight;
 use crate::core::{Pixels, Point, Radians, Rectangle, Size, Vector};
+use crate::graphics::cache::{self, Cached};
 use crate::graphics::geometry::fill::{self, Fill};
 use crate::graphics::geometry::stroke::{self, Stroke};
 use crate::graphics::geometry::{self, Path, Style};
-use crate::graphics::{Cached, Gradient, Text};
+use crate::graphics::{Gradient, Text};
 use crate::Primitive;
 
 use std::rc::Rc;
@@ -32,7 +33,7 @@ impl Cached for Geometry {
         Self::Cache(cache.clone())
     }
 
-    fn cache(self, _previous: Option<Cache>) -> Cache {
+    fn cache(self, _group: cache::Group, _previous: Option<Cache>) -> Cache {
         match self {
             Self::Live {
                 primitives,
diff --git a/wgpu/src/geometry.rs b/wgpu/src/geometry.rs
index 60967082..f6213e1d 100644
--- a/wgpu/src/geometry.rs
+++ b/wgpu/src/geometry.rs
@@ -3,6 +3,7 @@ use crate::core::text::LineHeight;
 use crate::core::{
     Pixels, Point, Radians, Rectangle, Size, Transformation, Vector,
 };
+use crate::graphics::cache::{self, Cached};
 use crate::graphics::color;
 use crate::graphics::geometry::fill::{self, Fill};
 use crate::graphics::geometry::{
@@ -10,7 +11,7 @@ use crate::graphics::geometry::{
 };
 use crate::graphics::gradient::{self, Gradient};
 use crate::graphics::mesh::{self, Mesh};
-use crate::graphics::{self, Cached, Text};
+use crate::graphics::{self, Text};
 use crate::text;
 use crate::triangle;
 
@@ -38,7 +39,11 @@ impl Cached for Geometry {
         Geometry::Cached(cache.clone())
     }
 
-    fn cache(self, previous: Option<Self::Cache>) -> Self::Cache {
+    fn cache(
+        self,
+        group: cache::Group,
+        previous: Option<Self::Cache>,
+    ) -> Self::Cache {
         match self {
             Self::Live { meshes, text } => {
                 if let Some(mut previous) = previous {
@@ -51,14 +56,14 @@ impl Cached for Geometry {
                     if let Some(cache) = &mut previous.text {
                         cache.update(text);
                     } else {
-                        previous.text = text::Cache::new(text);
+                        previous.text = text::Cache::new(group, text);
                     }
 
                     previous
                 } else {
                     Cache {
                         meshes: triangle::Cache::new(meshes),
-                        text: text::Cache::new(text),
+                        text: text::Cache::new(group, text),
                     }
                 }
             }
diff --git a/wgpu/src/text.rs b/wgpu/src/text.rs
index 38712660..5806863c 100644
--- a/wgpu/src/text.rs
+++ b/wgpu/src/text.rs
@@ -1,7 +1,8 @@
 use crate::core::alignment;
 use crate::core::{Rectangle, Size, Transformation};
+use crate::graphics::cache;
 use crate::graphics::color;
-use crate::graphics::text::cache::{self, Cache as BufferCache};
+use crate::graphics::text::cache::{self as text_cache, Cache as BufferCache};
 use crate::graphics::text::{font_system, to_color, Editor, Paragraph};
 
 use rustc_hash::FxHashMap;
@@ -35,6 +36,7 @@ pub enum Item {
 #[derive(Debug, Clone)]
 pub struct Cache {
     id: Id,
+    group: cache::Group,
     text: Rc<[Text]>,
     version: usize,
 }
@@ -43,7 +45,7 @@ pub struct Cache {
 pub struct Id(u64);
 
 impl Cache {
-    pub fn new(text: Vec<Text>) -> Option<Self> {
+    pub fn new(group: cache::Group, text: Vec<Text>) -> Option<Self> {
         static NEXT_ID: AtomicU64 = AtomicU64::new(0);
 
         if text.is_empty() {
@@ -52,6 +54,7 @@ impl Cache {
 
         Some(Self {
             id: Id(NEXT_ID.fetch_add(1, atomic::Ordering::Relaxed)),
+            group,
             text: Rc::from(text),
             version: 0,
         })
@@ -65,29 +68,39 @@ impl Cache {
 
 struct Upload {
     renderer: glyphon::TextRenderer,
-    atlas: glyphon::TextAtlas,
     buffer_cache: BufferCache,
     transformation: Transformation,
     version: usize,
     text: rc::Weak<[Text]>,
+    _atlas: rc::Weak<()>,
 }
 
 #[derive(Default)]
 pub struct Storage {
+    groups: FxHashMap<cache::Group, Group>,
     uploads: FxHashMap<Id, Upload>,
 }
 
+struct Group {
+    atlas: glyphon::TextAtlas,
+    previous_uploads: usize,
+    handle: Rc<()>,
+}
+
 impl Storage {
     pub fn new() -> Self {
         Self::default()
     }
 
-    fn get(&self, cache: &Cache) -> Option<&Upload> {
+    fn get(&self, cache: &Cache) -> Option<(&glyphon::TextAtlas, &Upload)> {
         if cache.text.is_empty() {
             return None;
         }
 
-        self.uploads.get(&cache.id)
+        self.groups
+            .get(&cache.group)
+            .map(|group| &group.atlas)
+            .zip(self.uploads.get(&cache.id))
     }
 
     fn prepare(
@@ -101,6 +114,20 @@ impl Storage {
         bounds: Rectangle,
         target_size: Size<u32>,
     ) {
+        let group_count = self.groups.len();
+
+        let group = self.groups.entry(cache.group).or_insert_with(|| {
+            log::info!("New text atlas created (total: {})", group_count + 1);
+
+            Group {
+                atlas: glyphon::TextAtlas::with_color_mode(
+                    device, queue, format, COLOR_MODE,
+                ),
+                previous_uploads: 0,
+                handle: Rc::new(()),
+            }
+        });
+
         match self.uploads.entry(cache.id) {
             hash_map::Entry::Occupied(entry) => {
                 let upload = entry.into_mut();
@@ -114,7 +141,7 @@ impl Storage {
                         queue,
                         encoder,
                         &mut upload.renderer,
-                        &mut upload.atlas,
+                        &mut group.atlas,
                         &mut upload.buffer_cache,
                         &cache.text,
                         bounds,
@@ -127,16 +154,11 @@ impl Storage {
                     upload.transformation = new_transformation;
 
                     upload.buffer_cache.trim();
-                    upload.atlas.trim();
                 }
             }
             hash_map::Entry::Vacant(entry) => {
-                let mut atlas = glyphon::TextAtlas::with_color_mode(
-                    device, queue, format, COLOR_MODE,
-                );
-
                 let mut renderer = glyphon::TextRenderer::new(
-                    &mut atlas,
+                    &mut group.atlas,
                     device,
                     wgpu::MultisampleState::default(),
                     None,
@@ -149,7 +171,7 @@ impl Storage {
                     queue,
                     encoder,
                     &mut renderer,
-                    &mut atlas,
+                    &mut group.atlas,
                     &mut buffer_cache,
                     &cache.text,
                     bounds,
@@ -159,11 +181,11 @@ impl Storage {
 
                 let _ = entry.insert(Upload {
                     renderer,
-                    atlas,
                     buffer_cache,
                     transformation: new_transformation,
                     version: 0,
                     text: Rc::downgrade(&cache.text),
+                    _atlas: Rc::downgrade(&group.handle),
                 });
 
                 log::info!(
@@ -178,6 +200,22 @@ impl Storage {
     pub fn trim(&mut self) {
         self.uploads
             .retain(|_id, upload| upload.text.strong_count() > 0);
+
+        self.groups.retain(|_id, group| {
+            let uploads_alive = Rc::weak_count(&group.handle);
+
+            if uploads_alive == 0 {
+                return false;
+            }
+
+            if uploads_alive < group.previous_uploads {
+                group.atlas.trim();
+            }
+
+            group.previous_uploads = uploads_alive;
+
+            true
+        });
     }
 }
 
@@ -306,10 +344,10 @@ impl Pipeline {
                     layer_count += 1;
                 }
                 Item::Cached { cache, .. } => {
-                    if let Some(upload) = storage.get(cache) {
+                    if let Some((atlas, upload)) = storage.get(cache) {
                         upload
                             .renderer
-                            .render(&upload.atlas, render_pass)
+                            .render(atlas, render_pass)
                             .expect("Render cached text");
                     }
                 }
@@ -345,7 +383,7 @@ fn prepare(
     enum Allocation {
         Paragraph(Paragraph),
         Editor(Editor),
-        Cache(cache::KeyHash),
+        Cache(text_cache::KeyHash),
         Raw(Arc<glyphon::Buffer>),
     }
 
@@ -369,7 +407,7 @@ fn prepare(
             } => {
                 let (key, _) = buffer_cache.allocate(
                     font_system,
-                    cache::Key {
+                    text_cache::Key {
                         content,
                         size: f32::from(*size),
                         line_height: f32::from(*line_height),
diff --git a/widget/src/canvas.rs b/widget/src/canvas.rs
index 42f92de0..be09f163 100644
--- a/widget/src/canvas.rs
+++ b/widget/src/canvas.rs
@@ -6,6 +6,7 @@ mod program;
 pub use event::Event;
 pub use program::Program;
 
+pub use crate::graphics::cache::Group;
 pub use crate::graphics::geometry::{
     fill, gradient, path, stroke, Fill, Gradient, LineCap, LineDash, LineJoin,
     Path, Stroke, Style, Text,