Add a free list to the atlas textures for quick re-use of empty slots

Notably, we've reverted the texture index reuse, as it would impose cost upon callers who do not
remove images from the atlas.

Co-authored-by: Mikayla Maki <mikayla@zed.dev>

crates/gpui/src/platform.rs

@@ -43,6 +43,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},
@@ -548,6 +549,41 @@ 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(..)
+    }
+
+    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

@@ -2689,6 +2689,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