gpui: Add drop_image (#19772)

This PR adds a function, WindowContext::drop_image, to manually remove a
RenderImage from the sprite atlas. In addition, PlatformAtlas::remove
was added to support this behavior. Previously, there was no way to
request a RenderImage to be removed from the sprite atlas, and since
they are not removed automatically the sprite would remain in video
memory once added until the window was closed. This PR allows a
developer to request the image be dropped from memory manually, however
it does not add automatic removal.

Release Notes:

- N/A

---------

Co-authored-by: Piotr Osiewicz <24362066+osiewicz@users.noreply.github.com>
Co-authored-by: Mikayla Maki <mikayla@zed.dev>

crates/gpui/src/platform.rs

@@ -46,6 +46,7 @@ use smallvec::SmallVec;
 use std::borrow::Cow;
 use std::hash::{Hash, Hasher};
 use std::io::Cursor;
+use std::ops;
 use std::time::{Duration, Instant};
 use std::{
     fmt::{self, Debug},
@@ -561,6 +562,42 @@ pub(crate) trait PlatformAtlas: Send + Sync {
         key: &AtlasKey,
         build: &mut dyn FnMut() -> Result<Option<(Size<DevicePixels>, Cow<'a, [u8]>)>>,
     ) -> Result<Option<AtlasTile>>;
+    fn remove(&self, key: &AtlasKey);
+}
+
+struct AtlasTextureList<T> {
+    textures: Vec<Option<T>>,
+    free_list: Vec<usize>,
+}
+
+impl<T> Default for AtlasTextureList<T> {
+    fn default() -> Self {
+        Self {
+            textures: Vec::default(),
+            free_list: Vec::default(),
+        }
+    }
+}
+
+impl<T> ops::Index<usize> for AtlasTextureList<T> {
+    type Output = Option<T>;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.textures[index]
+    }
+}
+
+impl<T> AtlasTextureList<T> {
+    #[allow(unused)]
+    fn drain(&mut self) -> std::vec::Drain<Option<T>> {
+        self.free_list.clear();
+        self.textures.drain(..)
+    }
+
+    #[allow(dead_code)]
+    fn iter_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut T> {
+        self.textures.iter_mut().flatten()
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]

crates/gpui/src/platform/blade/blade_atlas.rs

@@ -1,6 +1,6 @@
 use crate::{
-    AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTile, Bounds, DevicePixels, PlatformAtlas,
-    Point, Size,
+    platform::AtlasTextureList, AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTile, Bounds,
+    DevicePixels, PlatformAtlas, Point, Size,
 };
 use anyhow::Result;
 use blade_graphics as gpu;
@@ -67,7 +67,7 @@ impl BladeAtlas {
     pub(crate) fn clear_textures(&self, texture_kind: AtlasTextureKind) {
         let mut lock = self.0.lock();
         let textures = &mut lock.storage[texture_kind];
-        for texture in textures {
+        for texture in textures.iter_mut() {
             texture.clear();
         }
     }
@@ -130,19 +130,48 @@ impl PlatformAtlas for BladeAtlas {
             Ok(Some(tile))
         }
     }
+
+    fn remove(&self, key: &AtlasKey) {
+        let mut lock = self.0.lock();
+
+        let Some(id) = lock.tiles_by_key.remove(key).map(|tile| tile.texture_id) else {
+            return;
+        };
+
+        let Some(texture_slot) = lock.storage[id.kind].textures.get_mut(id.index as usize) else {
+            return;
+        };
+
+        if let Some(mut texture) = texture_slot.take() {
+            texture.decrement_ref_count();
+            if texture.is_unreferenced() {
+                lock.storage[id.kind]
+                    .free_list
+                    .push(texture.id.index as usize);
+                texture.destroy(&lock.gpu);
+            } else {
+                *texture_slot = Some(texture);
+            }
+        }
+    }
 }
 
 impl BladeAtlasState {
     fn allocate(&mut self, size: Size<DevicePixels>, texture_kind: AtlasTextureKind) -> AtlasTile {
-        let textures = &mut self.storage[texture_kind];
-        textures
-            .iter_mut()
-            .rev()
-            .find_map(|texture| texture.allocate(size))
-            .unwrap_or_else(|| {
-                let texture = self.push_texture(size, texture_kind);
-                texture.allocate(size).unwrap()
-            })
+        {
+            let textures = &mut self.storage[texture_kind];
+
+            if let Some(tile) = textures
+                .iter_mut()
+                .rev()
+                .find_map(|texture| texture.allocate(size))
+            {
+                return tile;
+            }
+        }
+
+        let texture = self.push_texture(size, texture_kind);
+        texture.allocate(size).unwrap()
     }
 
     fn push_texture(
@@ -198,21 +227,30 @@ impl BladeAtlasState {
             },
         );
 
-        let textures = &mut self.storage[kind];
+        let texture_list = &mut self.storage[kind];
+        let index = texture_list.free_list.pop();
+
         let atlas_texture = BladeAtlasTexture {
             id: AtlasTextureId {
-                index: textures.len() as u32,
+                index: index.unwrap_or(texture_list.textures.len()) as u32,
                 kind,
             },
             allocator: etagere::BucketedAtlasAllocator::new(size.into()),
             format,
             raw,
             raw_view,
+            live_atlas_keys: 0,
         };
 
         self.initializations.push(atlas_texture.id);
-        textures.push(atlas_texture);
-        textures.last_mut().unwrap()
+
+        if let Some(ix) = index {
+            texture_list.textures[ix] = Some(atlas_texture);
+            texture_list.textures.get_mut(ix).unwrap().as_mut().unwrap()
+        } else {
+            texture_list.textures.push(Some(atlas_texture));
+            texture_list.textures.last_mut().unwrap().as_mut().unwrap()
+        }
     }
 
     fn upload_texture(&mut self, id: AtlasTextureId, bounds: Bounds<DevicePixels>, bytes: &[u8]) {
@@ -258,13 +296,13 @@ impl BladeAtlasState {
 
 #[derive(Default)]
 struct BladeAtlasStorage {
-    monochrome_textures: Vec<BladeAtlasTexture>,
-    polychrome_textures: Vec<BladeAtlasTexture>,
-    path_textures: Vec<BladeAtlasTexture>,
+    monochrome_textures: AtlasTextureList<BladeAtlasTexture>,
+    polychrome_textures: AtlasTextureList<BladeAtlasTexture>,
+    path_textures: AtlasTextureList<BladeAtlasTexture>,
 }
 
 impl ops::Index<AtlasTextureKind> for BladeAtlasStorage {
-    type Output = Vec<BladeAtlasTexture>;
+    type Output = AtlasTextureList<BladeAtlasTexture>;
     fn index(&self, kind: AtlasTextureKind) -> &Self::Output {
         match kind {
             crate::AtlasTextureKind::Monochrome => &self.monochrome_textures,
@@ -292,19 +330,19 @@ impl ops::Index<AtlasTextureId> for BladeAtlasStorage {
             crate::AtlasTextureKind::Polychrome => &self.polychrome_textures,
             crate::AtlasTextureKind::Path => &self.path_textures,
         };
-        &textures[id.index as usize]
+        textures[id.index as usize].as_ref().unwrap()
     }
 }
 
 impl BladeAtlasStorage {
     fn destroy(&mut self, gpu: &gpu::Context) {
-        for mut texture in self.monochrome_textures.drain(..) {
+        for mut texture in self.monochrome_textures.drain().flatten() {
             texture.destroy(gpu);
         }
-        for mut texture in self.polychrome_textures.drain(..) {
+        for mut texture in self.polychrome_textures.drain().flatten() {
             texture.destroy(gpu);
         }
-        for mut texture in self.path_textures.drain(..) {
+        for mut texture in self.path_textures.drain().flatten() {
             texture.destroy(gpu);
         }
     }
@@ -316,6 +354,7 @@ struct BladeAtlasTexture {
     raw: gpu::Texture,
     raw_view: gpu::TextureView,
     format: gpu::TextureFormat,
+    live_atlas_keys: u32,
 }
 
 impl BladeAtlasTexture {
@@ -334,6 +373,7 @@ impl BladeAtlasTexture {
                 size,
             },
         };
+        self.live_atlas_keys += 1;
         Some(tile)
     }
 
@@ -345,6 +385,14 @@ impl BladeAtlasTexture {
     fn bytes_per_pixel(&self) -> u8 {
         self.format.block_info().size
     }
+
+    fn decrement_ref_count(&mut self) {
+        self.live_atlas_keys -= 1;
+    }
+
+    fn is_unreferenced(&mut self) -> bool {
+        self.live_atlas_keys == 0
+    }
 }
 
 impl From<Size<DevicePixels>> for etagere::Size {

crates/gpui/src/platform/mac/metal_atlas.rs

@@ -1,6 +1,6 @@
 use crate::{
-    AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTile, Bounds, DevicePixels, PlatformAtlas,
-    Point, Size,
+    platform::AtlasTextureList, AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTile, Bounds,
+    DevicePixels, PlatformAtlas, Point, Size,
 };
 use anyhow::{anyhow, Result};
 use collections::FxHashMap;
@@ -42,7 +42,7 @@ impl MetalAtlas {
             AtlasTextureKind::Polychrome => &mut lock.polychrome_textures,
             AtlasTextureKind::Path => &mut lock.path_textures,
         };
-        for texture in textures {
+        for texture in textures.iter_mut() {
             texture.clear();
         }
     }
@@ -50,9 +50,9 @@ impl MetalAtlas {
 
 struct MetalAtlasState {
     device: AssertSend<Device>,
-    monochrome_textures: Vec<MetalAtlasTexture>,
-    polychrome_textures: Vec<MetalAtlasTexture>,
-    path_textures: Vec<MetalAtlasTexture>,
+    monochrome_textures: AtlasTextureList<MetalAtlasTexture>,
+    polychrome_textures: AtlasTextureList<MetalAtlasTexture>,
+    path_textures: AtlasTextureList<MetalAtlasTexture>,
     tiles_by_key: FxHashMap<AtlasKey, AtlasTile>,
 }
 
@@ -78,6 +78,38 @@ impl PlatformAtlas for MetalAtlas {
             Ok(Some(tile))
         }
     }
+
+    fn remove(&self, key: &AtlasKey) {
+        let mut lock = self.0.lock();
+        let Some(id) = lock.tiles_by_key.get(key).map(|v| v.texture_id) else {
+            return;
+        };
+
+        let textures = match id.kind {
+            AtlasTextureKind::Monochrome => &mut lock.monochrome_textures,
+            AtlasTextureKind::Polychrome => &mut lock.polychrome_textures,
+            AtlasTextureKind::Path => &mut lock.polychrome_textures,
+        };
+
+        let Some(texture_slot) = textures
+            .textures
+            .iter_mut()
+            .find(|texture| texture.as_ref().is_some_and(|v| v.id == id))
+        else {
+            return;
+        };
+
+        if let Some(mut texture) = texture_slot.take() {
+            texture.decrement_ref_count();
+
+            if texture.is_unreferenced() {
+                textures.free_list.push(id.index as usize);
+                lock.tiles_by_key.remove(key);
+            } else {
+                *texture_slot = Some(texture);
+            }
+        }
+    }
 }
 
 impl MetalAtlasState {
@@ -86,20 +118,24 @@ impl MetalAtlasState {
         size: Size<DevicePixels>,
         texture_kind: AtlasTextureKind,
     ) -> Option<AtlasTile> {
-        let textures = match texture_kind {
-            AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
-            AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
-            AtlasTextureKind::Path => &mut self.path_textures,
-        };
+        {
+            let textures = match texture_kind {
+                AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
+                AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
+                AtlasTextureKind::Path => &mut self.path_textures,
+            };
 
-        textures
-            .iter_mut()
-            .rev()
-            .find_map(|texture| texture.allocate(size))
-            .or_else(|| {
-                let texture = self.push_texture(size, texture_kind);
-                texture.allocate(size)
-            })
+            if let Some(tile) = textures
+                .iter_mut()
+                .rev()
+                .find_map(|texture| texture.allocate(size))
+            {
+                return Some(tile);
+            }
+        }
+
+        let texture = self.push_texture(size, texture_kind);
+        texture.allocate(size)
     }
 
     fn push_texture(
@@ -140,21 +176,31 @@ impl MetalAtlasState {
         texture_descriptor.set_usage(usage);
         let metal_texture = self.device.new_texture(&texture_descriptor);
 
-        let textures = match kind {
+        let texture_list = match kind {
             AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
             AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
             AtlasTextureKind::Path => &mut self.path_textures,
         };
+
+        let index = texture_list.free_list.pop();
+
         let atlas_texture = MetalAtlasTexture {
             id: AtlasTextureId {
-                index: textures.len() as u32,
+                index: index.unwrap_or(texture_list.textures.len()) as u32,
                 kind,
             },
             allocator: etagere::BucketedAtlasAllocator::new(size.into()),
             metal_texture: AssertSend(metal_texture),
+            live_atlas_keys: 0,
         };
-        textures.push(atlas_texture);
-        textures.last_mut().unwrap()
+
+        if let Some(ix) = index {
+            texture_list.textures[ix] = Some(atlas_texture);
+            texture_list.textures.get_mut(ix).unwrap().as_mut().unwrap()
+        } else {
+            texture_list.textures.push(Some(atlas_texture));
+            texture_list.textures.last_mut().unwrap().as_mut().unwrap()
+        }
     }
 
     fn texture(&self, id: AtlasTextureId) -> &MetalAtlasTexture {
@@ -163,7 +209,7 @@ impl MetalAtlasState {
             crate::AtlasTextureKind::Polychrome => &self.polychrome_textures,
             crate::AtlasTextureKind::Path => &self.path_textures,
         };
-        &textures[id.index as usize]
+        textures[id.index as usize].as_ref().unwrap()
     }
 }
 
@@ -171,6 +217,7 @@ struct MetalAtlasTexture {
     id: AtlasTextureId,
     allocator: BucketedAtlasAllocator,
     metal_texture: AssertSend<metal::Texture>,
+    live_atlas_keys: u32,
 }
 
 impl MetalAtlasTexture {
@@ -189,6 +236,7 @@ impl MetalAtlasTexture {
             },
             padding: 0,
         };
+        self.live_atlas_keys += 1;
         Some(tile)
     }
 
@@ -215,6 +263,14 @@ impl MetalAtlasTexture {
             _ => unimplemented!(),
         }
     }
+
+    fn decrement_ref_count(&mut self) {
+        self.live_atlas_keys -= 1;
+    }
+
+    fn is_unreferenced(&mut self) -> bool {
+        self.live_atlas_keys == 0
+    }
 }
 
 impl From<Size<DevicePixels>> for etagere::Size {

crates/gpui/src/platform/test/window.rs

@@ -339,4 +339,9 @@ impl PlatformAtlas for TestAtlas {
 
         Ok(Some(state.tiles[key].clone()))
     }
+
+    fn remove(&self, key: &AtlasKey) {
+        let mut state = self.0.lock();
+        state.tiles.remove(key);
+    }
 }

crates/gpui/src/window.rs

@@ -2685,6 +2685,20 @@ impl<'a> WindowContext<'a> {
         });
     }
 
+    /// Removes an image from the sprite atlas.
+    pub fn drop_image(&mut self, data: Arc<RenderImage>) -> Result<()> {
+        for frame_index in 0..data.frame_count() {
+            let params = RenderImageParams {
+                image_id: data.id,
+                frame_index,
+            };
+
+            self.window.sprite_atlas.remove(&params.clone().into());
+        }
+
+        Ok(())
+    }
+
     #[must_use]
     /// Add a node to the layout tree for the current frame. Takes the `Style` of the element for which
     /// layout is being requested, along with the layout ids of any children. This method is called during