1use anyhow::{Context as _, Result};
2use collections::FxHashMap;
3use derive_more::{Deref, DerefMut};
4use etagere::BucketedAtlasAllocator;
5use gpui::{
6 AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTextureList, AtlasTile, Bounds, DevicePixels,
7 PlatformAtlas, Point, Size,
8};
9use metal::Device;
10use parking_lot::Mutex;
11use std::borrow::Cow;
12
13pub(crate) struct MetalAtlas(Mutex<MetalAtlasState>);
14
15impl MetalAtlas {
16 pub(crate) fn new(device: Device, is_apple_gpu: bool) -> Self {
17 MetalAtlas(Mutex::new(MetalAtlasState {
18 device: AssertSend(device),
19 is_apple_gpu,
20 monochrome_textures: Default::default(),
21 polychrome_textures: Default::default(),
22 tiles_by_key: Default::default(),
23 }))
24 }
25
26 pub(crate) fn metal_texture(&self, id: AtlasTextureId) -> metal::Texture {
27 self.0.lock().texture(id).metal_texture.clone()
28 }
29}
30
31struct MetalAtlasState {
32 device: AssertSend<Device>,
33 is_apple_gpu: bool,
34 monochrome_textures: AtlasTextureList<MetalAtlasTexture>,
35 polychrome_textures: AtlasTextureList<MetalAtlasTexture>,
36 tiles_by_key: FxHashMap<AtlasKey, AtlasTile>,
37}
38
39impl PlatformAtlas for MetalAtlas {
40 fn get_or_insert_with<'a>(
41 &self,
42 key: &AtlasKey,
43 build: &mut dyn FnMut() -> Result<Option<(Size<DevicePixels>, Cow<'a, [u8]>)>>,
44 ) -> Result<Option<AtlasTile>> {
45 let mut lock = self.0.lock();
46 if let Some(tile) = lock.tiles_by_key.get(key) {
47 Ok(Some(tile.clone()))
48 } else {
49 let Some((size, bytes)) = build()? else {
50 return Ok(None);
51 };
52 let tile = lock
53 .allocate(size, key.texture_kind())
54 .context("failed to allocate")?;
55 let texture = lock.texture(tile.texture_id);
56 texture.upload(tile.bounds, &bytes);
57 lock.tiles_by_key.insert(key.clone(), tile.clone());
58 Ok(Some(tile))
59 }
60 }
61
62 fn remove(&self, key: &AtlasKey) {
63 let mut lock = self.0.lock();
64 let Some(id) = lock.tiles_by_key.remove(key).map(|v| v.texture_id) else {
65 return;
66 };
67
68 let textures = match id.kind {
69 AtlasTextureKind::Monochrome => &mut lock.monochrome_textures,
70 AtlasTextureKind::Polychrome => &mut lock.polychrome_textures,
71 AtlasTextureKind::Subpixel => unreachable!(),
72 };
73
74 let Some(texture_slot) = textures
75 .textures
76 .iter_mut()
77 .find(|texture| texture.as_ref().is_some_and(|v| v.id == id))
78 else {
79 return;
80 };
81
82 if let Some(mut texture) = texture_slot.take() {
83 texture.decrement_ref_count();
84 if texture.is_unreferenced() {
85 textures.free_list.push(id.index as usize);
86 } else {
87 *texture_slot = Some(texture);
88 }
89 }
90 }
91}
92
93impl MetalAtlasState {
94 fn allocate(
95 &mut self,
96 size: Size<DevicePixels>,
97 texture_kind: AtlasTextureKind,
98 ) -> Option<AtlasTile> {
99 {
100 let textures = match texture_kind {
101 AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
102 AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
103 AtlasTextureKind::Subpixel => unreachable!(),
104 };
105
106 if let Some(tile) = textures
107 .iter_mut()
108 .rev()
109 .find_map(|texture| texture.allocate(size))
110 {
111 return Some(tile);
112 }
113 }
114
115 let texture = self.push_texture(size, texture_kind);
116 texture.allocate(size)
117 }
118
119 fn push_texture(
120 &mut self,
121 min_size: Size<DevicePixels>,
122 kind: AtlasTextureKind,
123 ) -> &mut MetalAtlasTexture {
124 const DEFAULT_ATLAS_SIZE: Size<DevicePixels> = Size {
125 width: DevicePixels(1024),
126 height: DevicePixels(1024),
127 };
128 // Max texture size on all modern Apple GPUs. Anything bigger than that crashes in validateWithDevice.
129 const MAX_ATLAS_SIZE: Size<DevicePixels> = Size {
130 width: DevicePixels(16384),
131 height: DevicePixels(16384),
132 };
133 let size = min_size.min(&MAX_ATLAS_SIZE).max(&DEFAULT_ATLAS_SIZE);
134 let texture_descriptor = metal::TextureDescriptor::new();
135 texture_descriptor.set_width(size.width.into());
136 texture_descriptor.set_height(size.height.into());
137 let pixel_format;
138 let usage;
139 match kind {
140 AtlasTextureKind::Monochrome => {
141 pixel_format = metal::MTLPixelFormat::A8Unorm;
142 usage = metal::MTLTextureUsage::ShaderRead;
143 }
144 AtlasTextureKind::Polychrome => {
145 pixel_format = metal::MTLPixelFormat::BGRA8Unorm;
146 usage = metal::MTLTextureUsage::ShaderRead;
147 }
148 AtlasTextureKind::Subpixel => unreachable!(),
149 }
150 texture_descriptor.set_pixel_format(pixel_format);
151 texture_descriptor.set_usage(usage);
152 // Shared memory mode can be used only on Apple GPU families
153 // https://developer.apple.com/documentation/metal/mtlresourceoptions/storagemodeshared
154 texture_descriptor.set_storage_mode(if self.is_apple_gpu {
155 metal::MTLStorageMode::Shared
156 } else {
157 metal::MTLStorageMode::Managed
158 });
159 let metal_texture = self.device.new_texture(&texture_descriptor);
160
161 let texture_list = match kind {
162 AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
163 AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
164 AtlasTextureKind::Subpixel => unreachable!(),
165 };
166
167 let index = texture_list.free_list.pop();
168
169 let atlas_texture = MetalAtlasTexture {
170 id: AtlasTextureId {
171 index: index.unwrap_or(texture_list.textures.len()) as u32,
172 kind,
173 },
174 allocator: etagere::BucketedAtlasAllocator::new(size_to_etagere(size)),
175 metal_texture: AssertSend(metal_texture),
176 live_atlas_keys: 0,
177 };
178
179 if let Some(ix) = index {
180 texture_list.textures[ix] = Some(atlas_texture);
181 texture_list.textures.get_mut(ix)
182 } else {
183 texture_list.textures.push(Some(atlas_texture));
184 texture_list.textures.last_mut()
185 }
186 .unwrap()
187 .as_mut()
188 .unwrap()
189 }
190
191 fn texture(&self, id: AtlasTextureId) -> &MetalAtlasTexture {
192 let textures = match id.kind {
193 AtlasTextureKind::Monochrome => &self.monochrome_textures,
194 AtlasTextureKind::Polychrome => &self.polychrome_textures,
195 AtlasTextureKind::Subpixel => unreachable!(),
196 };
197 textures[id.index as usize].as_ref().unwrap()
198 }
199}
200
201struct MetalAtlasTexture {
202 id: AtlasTextureId,
203 allocator: BucketedAtlasAllocator,
204 metal_texture: AssertSend<metal::Texture>,
205 live_atlas_keys: u32,
206}
207
208impl MetalAtlasTexture {
209 fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
210 let allocation = self.allocator.allocate(size_to_etagere(size))?;
211 let tile = AtlasTile {
212 texture_id: self.id,
213 tile_id: allocation.id.into(),
214 bounds: Bounds {
215 origin: point_from_etagere(allocation.rectangle.min),
216 size,
217 },
218 padding: 0,
219 };
220 self.live_atlas_keys += 1;
221 Some(tile)
222 }
223
224 fn upload(&self, bounds: Bounds<DevicePixels>, bytes: &[u8]) {
225 let region = metal::MTLRegion::new_2d(
226 bounds.origin.x.into(),
227 bounds.origin.y.into(),
228 bounds.size.width.into(),
229 bounds.size.height.into(),
230 );
231 self.metal_texture.replace_region(
232 region,
233 0,
234 bytes.as_ptr() as *const _,
235 bounds.size.width.to_bytes(self.bytes_per_pixel()) as u64,
236 );
237 }
238
239 fn bytes_per_pixel(&self) -> u8 {
240 use metal::MTLPixelFormat::*;
241 match self.metal_texture.pixel_format() {
242 A8Unorm | R8Unorm => 1,
243 RGBA8Unorm | BGRA8Unorm => 4,
244 _ => unimplemented!(),
245 }
246 }
247
248 fn decrement_ref_count(&mut self) {
249 self.live_atlas_keys -= 1;
250 }
251
252 fn is_unreferenced(&mut self) -> bool {
253 self.live_atlas_keys == 0
254 }
255}
256
257fn size_to_etagere(size: Size<DevicePixels>) -> etagere::Size {
258 etagere::Size::new(size.width.into(), size.height.into())
259}
260
261fn point_from_etagere(value: etagere::Point) -> Point<DevicePixels> {
262 Point {
263 x: DevicePixels::from(value.x),
264 y: DevicePixels::from(value.y),
265 }
266}
267
268#[derive(Deref, DerefMut)]
269struct AssertSend<T>(T);
270
271unsafe impl<T> Send for AssertSend<T> {}
272
273#[cfg(test)]
274mod tests {
275 use super::*;
276 use gpui::PlatformAtlas;
277 use std::borrow::Cow;
278
279 fn create_atlas() -> Option<MetalAtlas> {
280 let device = metal::Device::system_default()?;
281 Some(MetalAtlas::new(device, true))
282 }
283
284 fn make_image_key(image_id: usize, frame_index: usize) -> AtlasKey {
285 AtlasKey::Image(gpui::RenderImageParams {
286 image_id: gpui::ImageId(image_id),
287 frame_index,
288 })
289 }
290
291 fn insert_tile(atlas: &MetalAtlas, key: &AtlasKey, size: Size<DevicePixels>) -> AtlasTile {
292 atlas
293 .get_or_insert_with(key, &mut || {
294 let byte_count = (size.width.0 as usize) * (size.height.0 as usize) * 4;
295 Ok(Some((size, Cow::Owned(vec![0u8; byte_count]))))
296 })
297 .expect("allocation should succeed")
298 .expect("callback returns Some")
299 }
300
301 #[test]
302 fn test_remove_clears_stale_keys_from_tiles_by_key() {
303 let Some(atlas) = create_atlas() else {
304 return;
305 };
306
307 let small = Size {
308 width: DevicePixels(64),
309 height: DevicePixels(64),
310 };
311
312 let key_a = make_image_key(1, 0);
313 let key_b = make_image_key(2, 0);
314 let key_c = make_image_key(3, 0);
315
316 let tile_a = insert_tile(&atlas, &key_a, small);
317 let tile_b = insert_tile(&atlas, &key_b, small);
318 let tile_c = insert_tile(&atlas, &key_c, small);
319
320 assert_eq!(tile_a.texture_id, tile_b.texture_id);
321 assert_eq!(tile_b.texture_id, tile_c.texture_id);
322
323 // Remove A: texture still has B and C, so it stays.
324 // The key for A must be removed from tiles_by_key.
325 atlas.remove(&key_a);
326
327 // Remove B: texture still has C.
328 atlas.remove(&key_b);
329
330 // Remove C: texture becomes unreferenced and is deleted.
331 atlas.remove(&key_c);
332
333 // Re-inserting A must allocate a fresh tile on a new texture,
334 // NOT return a stale tile referencing the deleted texture.
335 let tile_a2 = insert_tile(&atlas, &key_a, small);
336
337 // The texture must actually exist — this would panic before the fix.
338 let _texture = atlas.metal_texture(tile_a2.texture_id);
339 }
340
341 #[test]
342 fn test_remove_nonexistent_key_is_noop() {
343 let Some(atlas) = create_atlas() else {
344 return;
345 };
346 let key = make_image_key(999, 0);
347 atlas.remove(&key);
348 }
349}