Adjust image cache APIs to enable ElementState based APIs (#29243)

cc: @sunli829 @huacnlee @probably-neb 

I really liked the earlier PR, but had an idea for how to utilize the
element state so that you don't need to construct the cache externally.
I've updated the APIs to introduce an `ImageCacheProvider` trait, and
added an example implementation of it to the image gallery :)

Release Notes:

- N/A

crates/gpui/examples/image_gallery.rs

@@ -1,10 +1,14 @@
+use futures::FutureExt;
 use gpui::{
-    App, AppContext, Application, Bounds, ClickEvent, Context, Entity, HashMapImageCache,
+    App, AppContext, Application, Asset as _, AssetLogger, Bounds, ClickEvent, Context, ElementId,
-    KeyBinding, Menu, MenuItem, SharedString, TitlebarOptions, Window, WindowBounds, WindowOptions,
+    Entity, HashMapImageCache, ImageAssetLoader, ImageCache, ImageCacheProvider, KeyBinding, Menu,
-    actions, div, image_cache, img, prelude::*, px, rgb, size,
+    MenuItem, SharedString, TitlebarOptions, Window, WindowBounds, WindowOptions, actions, div,
+    hash, image_cache, img, prelude::*, px, rgb, size,
 };
 use reqwest_client::ReqwestClient;
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
+
+const IMAGES_IN_GALLERY: usize = 30;
 
 struct ImageGallery {
     image_key: String,
@@ -34,57 +38,209 @@ impl Render for ImageGallery {
         let image_url: SharedString =
             format!("https://picsum.photos/400/200?t={}", self.image_key).into();
 
-        image_cache(&self.image_cache).child(
+        div()
-            div()
+            .flex()
-                .id("main")
+            .flex_col()
-                .font_family(".SystemUIFont")
+            .text_color(gpui::white())
-                .bg(rgb(0xE9E9E9))
+            .child("Manually managed image cache:")
-                .overflow_y_scroll()
+            .child(
-                .p_4()
+                image_cache(self.image_cache.clone()).child(
-                .size_full()
+                div()
-                .flex()
+                    .id("main")
-                .flex_col()
+                    .font_family(".SystemUIFont")
-                .items_center()
+                    .text_color(gpui::black())
-                .gap_2()
+                    .bg(rgb(0xE9E9E9))
-                .child(
+                    .overflow_y_scroll()
-                    div()
+                    .p_4()
-                        .w_full()
+                    .size_full()
-                        .flex()
+                    .flex()
-                        .flex_row()
+                    .flex_col()
-                        .justify_between()
+                    .items_center()
-                        .child(format!(
+                    .gap_2()
-                            "Example to show images and test memory usage (Rendered: {} images).",
+                    .child(
-                            self.total_count
+                        div()
-                        ))
+                            .w_full()
-                        .child(
+                            .flex()
-                            div()
+                            .flex_row()
-                                .id("btn")
+                            .justify_between()
-                                .py_1()
+                            .child(format!(
-                                .px_4()
+                                "Example to show images and test memory usage (Rendered: {} images).",
-                                .bg(gpui::black())
+                                self.total_count
-                                .hover(|this| this.opacity(0.8))
+                            ))
-                                .text_color(gpui::white())
+                            .child(
-                                .text_center()
+                                div()
-                                .w_40()
+                                    .id("btn")
-                                .child("Next Photos")
+                                    .py_1()
-                                .on_click(cx.listener(Self::on_next_image)),
+                                    .px_4()
-                        ),
+                                    .bg(gpui::black())
+                                    .hover(|this| this.opacity(0.8))
+                                    .text_color(gpui::white())
+                                    .text_center()
+                                    .w_40()
+                                    .child("Next Photos")
+                                    .on_click(cx.listener(Self::on_next_image)),
+                            ),
+                    )
+                    .child(
+                        div()
+                            .id("image-gallery")
+                            .flex()
+                            .flex_row()
+                            .flex_wrap()
+                            .gap_x_4()
+                            .gap_y_2()
+                            .justify_around()
+                            .children(
+                                (0..self.items_count)
+                                    .map(|ix| img(format!("{}-{}", image_url, ix)).size_20()),
+                            ),
+                    ),
+            ))
+            .child(
+                "Automatically managed image cache:"
+            )
+            .child(image_cache(simple_lru_cache("lru-cache", IMAGES_IN_GALLERY)).child(
+                div()
+                    .id("main")
+                    .font_family(".SystemUIFont")
+                    .bg(rgb(0xE9E9E9))
+                    .text_color(gpui::black())
+                    .overflow_y_scroll()
+                    .p_4()
+                    .size_full()
+                    .flex()
+                    .flex_col()
+                    .items_center()
+                    .gap_2()
+                    .child(
+                        div()
+                            .id("image-gallery")
+                            .flex()
+                            .flex_row()
+                            .flex_wrap()
+                            .gap_x_4()
+                            .gap_y_2()
+                            .justify_around()
+                            .children(
+                                (0..self.items_count)
+                                    .map(|ix| img(format!("{}-{}", image_url, ix)).size_20()),
+                            ),
+                    )
+            ))
+    }
+}
+
+fn simple_lru_cache(id: impl Into<ElementId>, max_items: usize) -> SimpleLruCacheProvider {
+    SimpleLruCacheProvider {
+        id: id.into(),
+        max_items,
+    }
+}
+
+struct SimpleLruCacheProvider {
+    id: ElementId,
+    max_items: usize,
+}
+
+impl ImageCacheProvider for SimpleLruCacheProvider {
+    fn provide(&mut self, window: &mut Window, cx: &mut App) -> gpui::AnyImageCache {
+        window
+            .with_global_id(self.id.clone(), |global_id, window| {
+                window.with_element_state::<Entity<SimpleLruCache>, _>(
+                    global_id,
+                    |lru_cache, _window| {
+                        let mut lru_cache = lru_cache.unwrap_or_else(|| {
+                            cx.new(|cx| SimpleLruCache::new(self.max_items, cx))
+                        });
+                        if lru_cache.read(cx).max_items != self.max_items {
+                            lru_cache = cx.new(|cx| SimpleLruCache::new(self.max_items, cx));
+                        }
+                        (lru_cache.clone(), lru_cache)
+                    },
                 )
-                .child(
+            })
-                    div()
+            .into()
-                        .id("image-gallery")
+    }
-                        .flex()
+}
-                        .flex_row()
+
-                        .flex_wrap()
+struct SimpleLruCache {
-                        .gap_x_4()
+    max_items: usize,
-                        .gap_y_2()
+    usages: Vec<u64>,
-                        .justify_around()
+    cache: HashMap<u64, gpui::ImageCacheItem>,
-                        .children(
+}
-                            (0..self.items_count)
+
-                                .map(|ix| img(format!("{}-{}", image_url, ix)).size_20()),
+impl SimpleLruCache {
-                        ),
+    fn new(max_items: usize, cx: &mut Context<Self>) -> Self {
-                ),
+        cx.on_release(|simple_cache, cx| {
-        )
+            for (_, mut item) in std::mem::take(&mut simple_cache.cache) {
+                if let Some(Ok(image)) = item.get() {
+                    cx.drop_image(image, None);
+                }
+            }
+        })
+        .detach();
+
+        Self {
+            max_items,
+            usages: Vec::with_capacity(max_items),
+            cache: HashMap::with_capacity(max_items),
+        }
+    }
+}
+
+impl ImageCache for SimpleLruCache {
+    fn load(
+        &mut self,
+        resource: &gpui::Resource,
+        window: &mut Window,
+        cx: &mut App,
+    ) -> Option<Result<Arc<gpui::RenderImage>, gpui::ImageCacheError>> {
+        assert_eq!(self.usages.len(), self.cache.len());
+        assert!(self.cache.len() <= self.max_items);
+
+        let hash = hash(resource);
+
+        if let Some(item) = self.cache.get_mut(&hash) {
+            let current_ix = self
+                .usages
+                .iter()
+                .position(|item| *item == hash)
+                .expect("cache and usages must stay in sync");
+            self.usages.remove(current_ix);
+            self.usages.insert(0, hash);
+
+            return item.get();
+        }
+
+        let fut = AssetLogger::<ImageAssetLoader>::load(resource.clone(), cx);
+        let task = cx.background_executor().spawn(fut).shared();
+        if self.usages.len() == self.max_items {
+            let oldest = self.usages.pop().unwrap();
+            let mut image = self
+                .cache
+                .remove(&oldest)
+                .expect("cache and usages must be in sync");
+            if let Some(Ok(image)) = image.get() {
+                cx.drop_image(image, Some(window));
+            }
+        }
+        self.cache
+            .insert(hash, gpui::ImageCacheItem::Loading(task.clone()));
+        self.usages.insert(0, hash);
+
+        let entity = window.current_view();
+        window
+            .spawn(cx, {
+                async move |cx| {
+                    _ = task.await;
+                    cx.on_next_frame(move |_, cx| {
+                        cx.notify(entity);
+                    });
+                }
+            })
+            .detach();
+
+        None
     }
 }
 
@@ -124,7 +280,7 @@ fn main() {
         cx.open_window(window_options, |_, cx| {
             cx.new(|ctx| ImageGallery {
                 image_key: "".into(),
-                items_count: 99,
+                items_count: IMAGES_IN_GALLERY,
                 total_count: 0,
                 image_cache: HashMapImageCache::new(ctx),
             })

crates/gpui/src/elements/image_cache.rs

@@ -10,9 +10,10 @@ use smallvec::SmallVec;
 use std::{collections::HashMap, fmt, sync::Arc};
 
 /// An image cache element, all its child img elements will use the cache specified by this element.
-pub fn image_cache<I: ImageCache>(image_cache: &Entity<I>) -> ImageCacheElement {
+/// Note that this could as simple as passing an `Entity<T: ImageCache>`
+pub fn image_cache(image_cache_provider: impl ImageCacheProvider) -> ImageCacheElement {
     ImageCacheElement {
-        image_cache: image_cache.clone().into(),
+        image_cache_provider: Box::new(image_cache_provider),
         style: StyleRefinement::default(),
         children: SmallVec::default(),
     }
@@ -68,7 +69,7 @@ mod any_image_cache {
 
 /// An image cache element.
 pub struct ImageCacheElement {
-    image_cache: AnyImageCache,
+    image_cache_provider: Box<dyn ImageCacheProvider>,
     style: StyleRefinement,
     children: SmallVec<[AnyElement; 2]>,
 }
@@ -107,7 +108,8 @@ impl Element for ImageCacheElement {
         window: &mut Window,
         cx: &mut App,
     ) -> (LayoutId, Self::RequestLayoutState) {
-        window.with_image_cache(self.image_cache.clone(), |window| {
+        let image_cache = self.image_cache_provider.provide(window, cx);
+        window.with_image_cache(image_cache, |window| {
             let child_layout_ids = self
                 .children
                 .iter_mut()
@@ -142,7 +144,8 @@ impl Element for ImageCacheElement {
         window: &mut Window,
         cx: &mut App,
     ) {
-        window.with_image_cache(self.image_cache.clone(), |window| {
+        let image_cache = self.image_cache_provider.provide(window, cx);
+        window.with_image_cache(image_cache, |window| {
             for child in &mut self.children {
                 child.paint(window, cx);
             }
@@ -150,22 +153,39 @@ impl Element for ImageCacheElement {
     }
 }
 
-type ImageLoadingTask = Shared<Task<Result<Arc<RenderImage>, ImageCacheError>>>;
+/// An image loading task associated with an image cache.
+pub type ImageLoadingTask = Shared<Task<Result<Arc<RenderImage>, ImageCacheError>>>;
 
-enum CacheItem {
+/// An image cache item
+pub enum ImageCacheItem {
+    /// The associated image is currently loading
     Loading(ImageLoadingTask),
+    /// This item has loaded an image.
     Loaded(Result<Arc<RenderImage>, ImageCacheError>),
 }
 
-impl CacheItem {
+impl std::fmt::Debug for ImageCacheItem {
-    fn get(&mut self) -> Option<Result<Arc<RenderImage>, ImageCacheError>> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let status = match self {
+            ImageCacheItem::Loading(_) => &"Loading...".to_string(),
+            ImageCacheItem::Loaded(render_image) => &format!("{:?}", render_image),
+        };
+        f.debug_struct("ImageCacheItem")
+            .field("status", status)
+            .finish()
+    }
+}
+
+impl ImageCacheItem {
+    /// Attempt to get the image from the cache item.
+    pub fn get(&mut self) -> Option<Result<Arc<RenderImage>, ImageCacheError>> {
         match self {
-            CacheItem::Loading(task) => {
+            ImageCacheItem::Loading(task) => {
                 let res = task.now_or_never()?;
-                *self = CacheItem::Loaded(res.clone());
+                *self = ImageCacheItem::Loaded(res.clone());
                 Some(res)
             }
-            CacheItem::Loaded(res) => Some(res.clone()),
+            ImageCacheItem::Loaded(res) => Some(res.clone()),
         }
     }
 }
@@ -183,8 +203,21 @@ pub trait ImageCache: 'static {
     ) -> Option<Result<Arc<RenderImage>, ImageCacheError>>;
 }
 
+/// An object that can create an ImageCache during the render phase.
+/// See the ImageCache trait for more information.
+pub trait ImageCacheProvider: 'static {
+    /// Called during the request_layout phase to create an ImageCache.
+    fn provide(&mut self, _window: &mut Window, _cx: &mut App) -> AnyImageCache;
+}
+
+impl<T: ImageCache> ImageCacheProvider for Entity<T> {
+    fn provide(&mut self, _window: &mut Window, _cx: &mut App) -> AnyImageCache {
+        self.clone().into()
+    }
+}
+
 /// An implementation of ImageCache, that uses an LRU caching strategy to unload images when the cache is full
-pub struct HashMapImageCache(HashMap<u64, CacheItem>);
+pub struct HashMapImageCache(HashMap<u64, ImageCacheItem>);
 
 impl fmt::Debug for HashMapImageCache {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -227,7 +260,7 @@ impl HashMapImageCache {
 
         let fut = AssetLogger::<ImageAssetLoader>::load(source.clone(), cx);
         let task = cx.background_executor().spawn(fut).shared();
-        self.0.insert(hash, CacheItem::Loading(task.clone()));
+        self.0.insert(hash, ImageCacheItem::Loading(task.clone()));
 
         let entity = window.current_view();
         window

crates/gpui/src/window.rs

@@ -1467,6 +1467,20 @@ impl Window {
         self.rem_size = rem_size.into();
     }
 
+    /// Acquire a globally unique identifier for the given ElementId.
+    /// Only valid for the duration of the provided closure.
+    pub fn with_global_id<R>(
+        &mut self,
+        element_id: ElementId,
+        f: impl FnOnce(&GlobalElementId, &mut Self) -> R,
+    ) -> R {
+        self.element_id_stack.push(element_id);
+        let global_id = GlobalElementId(self.element_id_stack.clone());
+        let result = f(&global_id, self);
+        self.element_id_stack.pop();
+        result
+    }
+
     /// Executes the provided function with the specified rem size.
     ///
     /// This method must only be called as part of element drawing.