signed/crates/workspace/src/image_cache.rs

use std::collections::{HashMap, VecDeque};
use std::mem::take;
use std::sync::Arc;

use futures::FutureExt;
use gpui::{
    App, AppContext, Asset, AssetLogger, Entity, Global, ImageAssetLoader, ImageCache,
    ImageCacheError, ImageCacheItem, ImageSource, RenderImage, Resource, Window, hash,
};

/// Upper bound on the number of images the shared cache retains. Loading a
/// new image evicts the least recently used entry once this is reached.
const MAX_IMAGES: usize = 128;

/// Global handle to the shared image cache, installed by [`init`].
struct SharedImageCache(Entity<LruImageCache>);

impl Global for SharedImageCache {}

/// Create the shared image cache and install it as a global. Call once at
/// startup (see `desktop/src/main.rs`), before any window opens.
pub fn init(cx: &mut App) -> Entity<LruImageCache> {
    let cache = LruImageCache::new(MAX_IMAGES, cx);
    cx.set_global(SharedImageCache(cache.clone()));
    cache
}

/// The shared image cache. Panics if [`init`] hasn't been called.
pub fn global(cx: &App) -> Entity<LruImageCache> {
    cx.global::<SharedImageCache>().0.clone()
}

/// Drop every cached image, freeing the decoded image data, the GPU textures
/// and the raw fetched bytes. Images currently on screen are re-fetched on
/// the next frame.
pub fn clear(cx: &mut App, window: &mut Window) {
    let cache = global(cx);
    cache.update(cx, |cache, cx| cache.clear(window, cx));
}

/// Clear the shared cache when `view` is released — its last strong
/// reference is gone, e.g. the user closed the panel it was rendered in.
pub fn clear_on_release<T: 'static>(view: &Entity<T>, window: &Window, cx: &mut App) {
    let cache = global(cx).downgrade();
    cx.observe_release_in(view, window, move |_view, window, cx| {
        if let Some(cache) = cache.upgrade() {
            cache.update(cx, |cache, cx| cache.clear(window, cx));
        }
    })
    .detach();
}

/// A bounded LRU image cache.
///
/// Entries keep the [`Resource`] they were loaded from so that eviction,
/// clearing and release can also remove the asset from the asset system,
/// freeing the raw fetched bytes that would otherwise stay in memory.
pub struct LruImageCache {
    max_items: usize,
    /// Most recently used hashes first.
    usage: VecDeque<u64>,
    cache: HashMap<u64, (ImageCacheItem, Resource)>,
}

impl LruImageCache {
    /// Create a cache that holds at most `max_items` images. Cached images
    /// are dropped from every window when the cache is released.
    pub fn new(max_items: usize, cx: &mut App) -> Entity<Self> {
        let max_items = max_items.max(1);
        cx.new(|cx| {
            cx.on_release(|this: &mut Self, cx| {
                for (_, entry) in take(&mut this.cache) {
                    unload(entry, None, cx);
                }
            })
            .detach();
            Self {
                max_items,
                usage: VecDeque::with_capacity(max_items),
                cache: HashMap::with_capacity(max_items),
            }
        })
    }

    /// Drop every cached image and remove every cached asset.
    pub fn clear(&mut self, window: &mut Window, cx: &mut App) {
        self.usage.clear();
        for (_, entry) in take(&mut self.cache) {
            unload(entry, Some(window), cx);
        }
    }
}

/// Drop a cache entry's decoded data and remove its resource from the asset
/// system, so both the decoded image and the raw fetched bytes are freed.
///
/// The atlas texture is freed **on the next frame, before it paints**, never
/// in the middle of one: the release that empties the cache can run at the
/// end of a frame's draw — after the scene was built, before it is presented
/// — and eviction runs while a frame is painting (`load` is called from
/// paint). The next frame also has to repaint every view instead of replaying
/// their recorded paint commands: `cached()` views reuse recorded commands
/// across frames, and those commands reference the atlas tiles being freed,
/// so a replay would hand the renderer a scene full of freed texture ids.
/// `window.refresh()` disables that reuse for exactly one frame.
///
/// `window` restricts the atlas removal to the current window; `None` removes
/// it from all windows (the cache entity is being released at shutdown, when
/// no scene is in flight).
fn unload(
    (mut item, resource): (ImageCacheItem, Resource),
    window: Option<&mut Window>,
    cx: &mut App,
) {
    if let Some(Ok(image)) = item.get() {
        match window {
            Some(window) => {
                window.on_next_frame(move |window, cx| {
                    window.refresh();
                    cx.drop_image(image, Some(window));
                });
            }
            None => cx.drop_image(image, None),
        }
    }
    ImageSource::Resource(resource).remove_asset(cx);
}

impl ImageCache for LruImageCache {
    fn load(
        &mut self,
        resource: &Resource,
        window: &mut Window,
        cx: &mut App,
    ) -> Option<Result<Arc<RenderImage>, ImageCacheError>> {
        debug_assert_eq!(self.usage.len(), self.cache.len());
        debug_assert!(self.cache.len() <= self.max_items);

        let hash = hash(resource);

        if let Some((item, _)) = self.cache.get_mut(&hash) {
            let current_ix = self
                .usage
                .iter()
                .position(|used| *used == hash)
                .expect("cache and usage list must stay in sync");
            self.usage.remove(current_ix);
            self.usage.push_front(hash);
            return item.get();
        }

        let fut = AssetLogger::<ImageAssetLoader>::load(resource.clone(), cx);
        let task = cx.background_executor().spawn(fut).shared();

        if self.usage.len() >= self.max_items {
            let oldest = self
                .usage
                .pop_back()
                .expect("usage list and cache must stay in sync");
            let entry = self
                .cache
                .remove(&oldest)
                .expect("usage list and cache must stay in sync");
            unload(entry, Some(window), cx);
        }

        self.cache.insert(
            hash,
            (ImageCacheItem::Loading(task.clone()), resource.clone()),
        );
        self.usage.push_front(hash);

        let entity = window.current_view();
        window
            .spawn(cx, {
                async move |cx| {
                    if let Err(error) = task.await {
                        log::error!("failed to load image into cache: {:?}", error);
                    }
                    cx.on_next_frame(move |_, cx| {
                        cx.notify(entity);
                    });
                }
            })
            .detach();

        None
    }
}