tenant.openagents/omega

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