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