metal_atlas.rs

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