metal_atlas.rs

  1use crate::{
  2    AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTile, Bounds, DevicePixels, PlatformAtlas,
  3    Point, Size, platform::AtlasTextureList,
  4};
  5use anyhow::{Context as _, Result};
  6use collections::FxHashMap;
  7use derive_more::{Deref, DerefMut};
  8use etagere::BucketedAtlasAllocator;
  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) -> Self {
 17        MetalAtlas(Mutex::new(MetalAtlasState {
 18            device: AssertSend(device),
 19            monochrome_textures: Default::default(),
 20            polychrome_textures: Default::default(),
 21            tiles_by_key: Default::default(),
 22        }))
 23    }
 24
 25    pub(crate) fn metal_texture(&self, id: AtlasTextureId) -> metal::Texture {
 26        self.0.lock().texture(id).metal_texture.clone()
 27    }
 28}
 29
 30struct MetalAtlasState {
 31    device: AssertSend<Device>,
 32    monochrome_textures: AtlasTextureList<MetalAtlasTexture>,
 33    polychrome_textures: AtlasTextureList<MetalAtlasTexture>,
 34    tiles_by_key: FxHashMap<AtlasKey, AtlasTile>,
 35}
 36
 37impl PlatformAtlas for MetalAtlas {
 38    fn get_or_insert_with<'a>(
 39        &self,
 40        key: &AtlasKey,
 41        build: &mut dyn FnMut() -> Result<Option<(Size<DevicePixels>, Cow<'a, [u8]>)>>,
 42    ) -> Result<Option<AtlasTile>> {
 43        let mut lock = self.0.lock();
 44        if let Some(tile) = lock.tiles_by_key.get(key) {
 45            Ok(Some(tile.clone()))
 46        } else {
 47            let Some((size, bytes)) = build()? else {
 48                return Ok(None);
 49            };
 50            let tile = lock
 51                .allocate(size, key.texture_kind())
 52                .context("failed to allocate")?;
 53            let texture = lock.texture(tile.texture_id);
 54            texture.upload(tile.bounds, &bytes);
 55            lock.tiles_by_key.insert(key.clone(), tile.clone());
 56            Ok(Some(tile))
 57        }
 58    }
 59
 60    fn clear_glyphs(&self) {
 61        let mut lock = self.0.lock();
 62        let mut tile_ids_by_texture = FxHashMap::default();
 63        lock.tiles_by_key.retain(|key, tile| {
 64            if matches!(key, AtlasKey::Glyph(_)) {
 65                tile_ids_by_texture
 66                    .entry(tile.texture_id)
 67                    .or_insert_with(Vec::new)
 68                    .push(tile.tile_id);
 69                false
 70            } else {
 71                true
 72            }
 73        });
 74
 75        for (texture_id, tile_ids) in tile_ids_by_texture {
 76            let texture_slot = &mut lock.monochrome_textures.textures[texture_id.index as usize];
 77            if let Some(texture) = texture_slot {
 78                texture.live_atlas_keys -= tile_ids.len() as u32;
 79                if texture.is_unreferenced() {
 80                    texture_slot.take();
 81                    lock.monochrome_textures
 82                        .free_list
 83                        .push(texture_id.index as usize);
 84                }
 85            }
 86        }
 87    }
 88
 89    fn remove(&self, key: &AtlasKey) {
 90        let mut lock = self.0.lock();
 91        let Some(id) = lock.tiles_by_key.get(key).map(|v| v.texture_id) else {
 92            return;
 93        };
 94
 95        let textures = match id.kind {
 96            AtlasTextureKind::Monochrome => &mut lock.monochrome_textures,
 97            AtlasTextureKind::Polychrome => &mut lock.polychrome_textures,
 98        };
 99
100        let Some(texture_slot) = textures
101            .textures
102            .iter_mut()
103            .find(|texture| texture.as_ref().is_some_and(|v| v.id == id))
104        else {
105            return;
106        };
107
108        if let Some(mut texture) = texture_slot.take() {
109            texture.decrement_ref_count();
110
111            if texture.is_unreferenced() {
112                textures.free_list.push(id.index as usize);
113                lock.tiles_by_key.remove(key);
114            } else {
115                *texture_slot = Some(texture);
116            }
117        }
118    }
119}
120
121impl MetalAtlasState {
122    fn allocate(
123        &mut self,
124        size: Size<DevicePixels>,
125        texture_kind: AtlasTextureKind,
126    ) -> Option<AtlasTile> {
127        {
128            let textures = match texture_kind {
129                AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
130                AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
131            };
132
133            if let Some(tile) = textures
134                .iter_mut()
135                .rev()
136                .find_map(|texture| texture.allocate(size))
137            {
138                return Some(tile);
139            }
140        }
141
142        let texture = self.push_texture(size, texture_kind);
143        texture.allocate(size)
144    }
145
146    fn push_texture(
147        &mut self,
148        min_size: Size<DevicePixels>,
149        kind: AtlasTextureKind,
150    ) -> &mut MetalAtlasTexture {
151        const DEFAULT_ATLAS_SIZE: Size<DevicePixels> = Size {
152            width: DevicePixels(1024),
153            height: DevicePixels(1024),
154        };
155        // Max texture size on all modern Apple GPUs. Anything bigger than that crashes in validateWithDevice.
156        const MAX_ATLAS_SIZE: Size<DevicePixels> = Size {
157            width: DevicePixels(16384),
158            height: DevicePixels(16384),
159        };
160        let size = min_size.min(&MAX_ATLAS_SIZE).max(&DEFAULT_ATLAS_SIZE);
161        let texture_descriptor = metal::TextureDescriptor::new();
162        texture_descriptor.set_width(size.width.into());
163        texture_descriptor.set_height(size.height.into());
164        let pixel_format;
165        let usage;
166        match kind {
167            AtlasTextureKind::Monochrome => {
168                pixel_format = metal::MTLPixelFormat::A8Unorm;
169                usage = metal::MTLTextureUsage::ShaderRead;
170            }
171            AtlasTextureKind::Polychrome => {
172                pixel_format = metal::MTLPixelFormat::BGRA8Unorm;
173                usage = metal::MTLTextureUsage::ShaderRead;
174            }
175        }
176        texture_descriptor.set_pixel_format(pixel_format);
177        texture_descriptor.set_usage(usage);
178        let metal_texture = self.device.new_texture(&texture_descriptor);
179
180        let texture_list = match kind {
181            AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
182            AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
183        };
184
185        let index = texture_list.free_list.pop();
186
187        let atlas_texture = MetalAtlasTexture {
188            id: AtlasTextureId {
189                index: index.unwrap_or(texture_list.textures.len()) as u32,
190                kind,
191            },
192            allocator: etagere::BucketedAtlasAllocator::new(size.into()),
193            metal_texture: AssertSend(metal_texture),
194            live_atlas_keys: 0,
195        };
196
197        if let Some(ix) = index {
198            texture_list.textures[ix] = Some(atlas_texture);
199            texture_list.textures.get_mut(ix).unwrap().as_mut().unwrap()
200        } else {
201            texture_list.textures.push(Some(atlas_texture));
202            texture_list.textures.last_mut().unwrap().as_mut().unwrap()
203        }
204    }
205
206    fn texture(&self, id: AtlasTextureId) -> &MetalAtlasTexture {
207        let textures = match id.kind {
208            crate::AtlasTextureKind::Monochrome => &self.monochrome_textures,
209            crate::AtlasTextureKind::Polychrome => &self.polychrome_textures,
210        };
211        textures[id.index as usize].as_ref().unwrap()
212    }
213}
214
215struct MetalAtlasTexture {
216    id: AtlasTextureId,
217    allocator: BucketedAtlasAllocator,
218    metal_texture: AssertSend<metal::Texture>,
219    live_atlas_keys: u32,
220}
221
222impl MetalAtlasTexture {
223    fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
224        let allocation = self.allocator.allocate(size.into())?;
225        let tile = AtlasTile {
226            texture_id: self.id,
227            tile_id: allocation.id.into(),
228            bounds: Bounds {
229                origin: allocation.rectangle.min.into(),
230                size,
231            },
232            padding: 0,
233        };
234        self.live_atlas_keys += 1;
235        Some(tile)
236    }
237
238    fn upload(&self, bounds: Bounds<DevicePixels>, bytes: &[u8]) {
239        let region = metal::MTLRegion::new_2d(
240            bounds.origin.x.into(),
241            bounds.origin.y.into(),
242            bounds.size.width.into(),
243            bounds.size.height.into(),
244        );
245        self.metal_texture.replace_region(
246            region,
247            0,
248            bytes.as_ptr() as *const _,
249            bounds.size.width.to_bytes(self.bytes_per_pixel()) as u64,
250        );
251    }
252
253    fn bytes_per_pixel(&self) -> u8 {
254        use metal::MTLPixelFormat::*;
255        match self.metal_texture.pixel_format() {
256            A8Unorm | R8Unorm => 1,
257            RGBA8Unorm | BGRA8Unorm => 4,
258            _ => unimplemented!(),
259        }
260    }
261
262    fn decrement_ref_count(&mut self) {
263        self.live_atlas_keys -= 1;
264    }
265
266    fn is_unreferenced(&mut self) -> bool {
267        self.live_atlas_keys == 0
268    }
269}
270
271impl From<Size<DevicePixels>> for etagere::Size {
272    fn from(size: Size<DevicePixels>) -> Self {
273        etagere::Size::new(size.width.into(), size.height.into())
274    }
275}
276
277impl From<etagere::Point> for Point<DevicePixels> {
278    fn from(value: etagere::Point) -> Self {
279        Point {
280            x: DevicePixels::from(value.x),
281            y: DevicePixels::from(value.y),
282        }
283    }
284}
285
286impl From<etagere::Size> for Size<DevicePixels> {
287    fn from(size: etagere::Size) -> Self {
288        Size {
289            width: DevicePixels::from(size.width),
290            height: DevicePixels::from(size.height),
291        }
292    }
293}
294
295impl From<etagere::Rectangle> for Bounds<DevicePixels> {
296    fn from(rectangle: etagere::Rectangle) -> Self {
297        Bounds {
298            origin: rectangle.min.into(),
299            size: rectangle.size().into(),
300        }
301    }
302}
303
304#[derive(Deref, DerefMut)]
305struct AssertSend<T>(T);
306
307unsafe impl<T> Send for AssertSend<T> {}