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 remove(&self, key: &AtlasKey) {
 61        let mut lock = self.0.lock();
 62        let Some(id) = lock.tiles_by_key.get(key).map(|v| v.texture_id) else {
 63            return;
 64        };
 65
 66        let textures = match id.kind {
 67            AtlasTextureKind::Monochrome => &mut lock.monochrome_textures,
 68            AtlasTextureKind::Polychrome => &mut lock.polychrome_textures,
 69            AtlasTextureKind::Subpixel => unreachable!(),
 70        };
 71
 72        let Some(texture_slot) = textures
 73            .textures
 74            .iter_mut()
 75            .find(|texture| texture.as_ref().is_some_and(|v| v.id == id))
 76        else {
 77            return;
 78        };
 79
 80        if let Some(mut texture) = texture_slot.take() {
 81            texture.decrement_ref_count();
 82
 83            if texture.is_unreferenced() {
 84                textures.free_list.push(id.index as usize);
 85                lock.tiles_by_key.remove(key);
 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        let metal_texture = self.device.new_texture(&texture_descriptor);
153
154        let texture_list = match kind {
155            AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
156            AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
157            AtlasTextureKind::Subpixel => unreachable!(),
158        };
159
160        let index = texture_list.free_list.pop();
161
162        let atlas_texture = MetalAtlasTexture {
163            id: AtlasTextureId {
164                index: index.unwrap_or(texture_list.textures.len()) as u32,
165                kind,
166            },
167            allocator: etagere::BucketedAtlasAllocator::new(size.into()),
168            metal_texture: AssertSend(metal_texture),
169            live_atlas_keys: 0,
170        };
171
172        if let Some(ix) = index {
173            texture_list.textures[ix] = Some(atlas_texture);
174            texture_list.textures.get_mut(ix)
175        } else {
176            texture_list.textures.push(Some(atlas_texture));
177            texture_list.textures.last_mut()
178        }
179        .unwrap()
180        .as_mut()
181        .unwrap()
182    }
183
184    fn texture(&self, id: AtlasTextureId) -> &MetalAtlasTexture {
185        let textures = match id.kind {
186            crate::AtlasTextureKind::Monochrome => &self.monochrome_textures,
187            crate::AtlasTextureKind::Polychrome => &self.polychrome_textures,
188            crate::AtlasTextureKind::Subpixel => unreachable!(),
189        };
190        textures[id.index as usize].as_ref().unwrap()
191    }
192}
193
194struct MetalAtlasTexture {
195    id: AtlasTextureId,
196    allocator: BucketedAtlasAllocator,
197    metal_texture: AssertSend<metal::Texture>,
198    live_atlas_keys: u32,
199}
200
201impl MetalAtlasTexture {
202    fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
203        let allocation = self.allocator.allocate(size.into())?;
204        let tile = AtlasTile {
205            texture_id: self.id,
206            tile_id: allocation.id.into(),
207            bounds: Bounds {
208                origin: allocation.rectangle.min.into(),
209                size,
210            },
211            padding: 0,
212        };
213        self.live_atlas_keys += 1;
214        Some(tile)
215    }
216
217    fn upload(&self, bounds: Bounds<DevicePixels>, bytes: &[u8]) {
218        let region = metal::MTLRegion::new_2d(
219            bounds.origin.x.into(),
220            bounds.origin.y.into(),
221            bounds.size.width.into(),
222            bounds.size.height.into(),
223        );
224        self.metal_texture.replace_region(
225            region,
226            0,
227            bytes.as_ptr() as *const _,
228            bounds.size.width.to_bytes(self.bytes_per_pixel()) as u64,
229        );
230    }
231
232    fn bytes_per_pixel(&self) -> u8 {
233        use metal::MTLPixelFormat::*;
234        match self.metal_texture.pixel_format() {
235            A8Unorm | R8Unorm => 1,
236            RGBA8Unorm | BGRA8Unorm => 4,
237            _ => unimplemented!(),
238        }
239    }
240
241    fn decrement_ref_count(&mut self) {
242        self.live_atlas_keys -= 1;
243    }
244
245    fn is_unreferenced(&mut self) -> bool {
246        self.live_atlas_keys == 0
247    }
248}
249
250impl From<Size<DevicePixels>> for etagere::Size {
251    fn from(size: Size<DevicePixels>) -> Self {
252        etagere::Size::new(size.width.into(), size.height.into())
253    }
254}
255
256impl From<etagere::Point> for Point<DevicePixels> {
257    fn from(value: etagere::Point) -> Self {
258        Point {
259            x: DevicePixels::from(value.x),
260            y: DevicePixels::from(value.y),
261        }
262    }
263}
264
265impl From<etagere::Size> for Size<DevicePixels> {
266    fn from(size: etagere::Size) -> Self {
267        Size {
268            width: DevicePixels::from(size.width),
269            height: DevicePixels::from(size.height),
270        }
271    }
272}
273
274impl From<etagere::Rectangle> for Bounds<DevicePixels> {
275    fn from(rectangle: etagere::Rectangle) -> Self {
276        Bounds {
277            origin: rectangle.min.into(),
278            size: rectangle.size().into(),
279        }
280    }
281}
282
283#[derive(Deref, DerefMut)]
284struct AssertSend<T>(T);
285
286unsafe impl<T> Send for AssertSend<T> {}