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 // Shared memory can be used only if CPU and GPU share the same memory space.
19 // https://developer.apple.com/documentation/metal/setting-resource-storage-modes
20 unified_memory: device.has_unified_memory(),
21 device: AssertSend(device),
22 monochrome_textures: Default::default(),
23 polychrome_textures: Default::default(),
24 tiles_by_key: Default::default(),
25 }))
26 }
27
28 pub(crate) fn metal_texture(&self, id: AtlasTextureId) -> metal::Texture {
29 self.0.lock().texture(id).metal_texture.clone()
30 }
31}
32
33struct MetalAtlasState {
34 device: AssertSend<Device>,
35 unified_memory: bool,
36 monochrome_textures: AtlasTextureList<MetalAtlasTexture>,
37 polychrome_textures: AtlasTextureList<MetalAtlasTexture>,
38 tiles_by_key: FxHashMap<AtlasKey, AtlasTile>,
39}
40
41impl PlatformAtlas for MetalAtlas {
42 fn get_or_insert_with<'a>(
43 &self,
44 key: &AtlasKey,
45 build: &mut dyn FnMut() -> Result<Option<(Size<DevicePixels>, Cow<'a, [u8]>)>>,
46 ) -> Result<Option<AtlasTile>> {
47 let mut lock = self.0.lock();
48 if let Some(tile) = lock.tiles_by_key.get(key) {
49 Ok(Some(tile.clone()))
50 } else {
51 let Some((size, bytes)) = build()? else {
52 return Ok(None);
53 };
54 let tile = lock
55 .allocate(size, key.texture_kind())
56 .context("failed to allocate")?;
57 let texture = lock.texture(tile.texture_id);
58 texture.upload(tile.bounds, &bytes);
59 lock.tiles_by_key.insert(key.clone(), tile.clone());
60 Ok(Some(tile))
61 }
62 }
63
64 fn remove(&self, key: &AtlasKey) {
65 let mut lock = self.0.lock();
66 let Some(id) = lock.tiles_by_key.get(key).map(|v| v.texture_id) else {
67 return;
68 };
69
70 let textures = match id.kind {
71 AtlasTextureKind::Monochrome => &mut lock.monochrome_textures,
72 AtlasTextureKind::Polychrome => &mut lock.polychrome_textures,
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.decrement_ref_count();
85
86 if texture.is_unreferenced() {
87 textures.free_list.push(id.index as usize);
88 lock.tiles_by_key.remove(key);
89 } else {
90 *texture_slot = Some(texture);
91 }
92 }
93 }
94}
95
96impl MetalAtlasState {
97 fn allocate(
98 &mut self,
99 size: Size<DevicePixels>,
100 texture_kind: AtlasTextureKind,
101 ) -> Option<AtlasTile> {
102 {
103 let textures = match texture_kind {
104 AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
105 AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
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 }
151 texture_descriptor.set_pixel_format(pixel_format);
152 texture_descriptor.set_usage(usage);
153 texture_descriptor.set_storage_mode(if self.unified_memory {
154 metal::MTLStorageMode::Shared
155 } else {
156 metal::MTLStorageMode::Managed
157 });
158 let metal_texture = self.device.new_texture(&texture_descriptor);
159
160 let texture_list = match kind {
161 AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
162 AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
163 };
164
165 let index = texture_list.free_list.pop();
166
167 let atlas_texture = MetalAtlasTexture {
168 id: AtlasTextureId {
169 index: index.unwrap_or(texture_list.textures.len()) as u32,
170 kind,
171 },
172 allocator: etagere::BucketedAtlasAllocator::new(size.into()),
173 metal_texture: AssertSend(metal_texture),
174 live_atlas_keys: 0,
175 };
176
177 if let Some(ix) = index {
178 texture_list.textures[ix] = Some(atlas_texture);
179 texture_list.textures.get_mut(ix)
180 } else {
181 texture_list.textures.push(Some(atlas_texture));
182 texture_list.textures.last_mut()
183 }
184 .unwrap()
185 .as_mut()
186 .unwrap()
187 }
188
189 fn texture(&self, id: AtlasTextureId) -> &MetalAtlasTexture {
190 let textures = match id.kind {
191 crate::AtlasTextureKind::Monochrome => &self.monochrome_textures,
192 crate::AtlasTextureKind::Polychrome => &self.polychrome_textures,
193 };
194 textures[id.index as usize].as_ref().unwrap()
195 }
196}
197
198struct MetalAtlasTexture {
199 id: AtlasTextureId,
200 allocator: BucketedAtlasAllocator,
201 metal_texture: AssertSend<metal::Texture>,
202 live_atlas_keys: u32,
203}
204
205impl MetalAtlasTexture {
206 fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
207 let allocation = self.allocator.allocate(size.into())?;
208 let tile = AtlasTile {
209 texture_id: self.id,
210 tile_id: allocation.id.into(),
211 bounds: Bounds {
212 origin: allocation.rectangle.min.into(),
213 size,
214 },
215 padding: 0,
216 };
217 self.live_atlas_keys += 1;
218 Some(tile)
219 }
220
221 fn upload(&self, bounds: Bounds<DevicePixels>, bytes: &[u8]) {
222 let region = metal::MTLRegion::new_2d(
223 bounds.origin.x.into(),
224 bounds.origin.y.into(),
225 bounds.size.width.into(),
226 bounds.size.height.into(),
227 );
228 self.metal_texture.replace_region(
229 region,
230 0,
231 bytes.as_ptr() as *const _,
232 bounds.size.width.to_bytes(self.bytes_per_pixel()) as u64,
233 );
234 }
235
236 fn bytes_per_pixel(&self) -> u8 {
237 use metal::MTLPixelFormat::*;
238 match self.metal_texture.pixel_format() {
239 A8Unorm | R8Unorm => 1,
240 RGBA8Unorm | BGRA8Unorm => 4,
241 _ => unimplemented!(),
242 }
243 }
244
245 fn decrement_ref_count(&mut self) {
246 self.live_atlas_keys -= 1;
247 }
248
249 fn is_unreferenced(&mut self) -> bool {
250 self.live_atlas_keys == 0
251 }
252}
253
254impl From<Size<DevicePixels>> for etagere::Size {
255 fn from(size: Size<DevicePixels>) -> Self {
256 etagere::Size::new(size.width.into(), size.height.into())
257 }
258}
259
260impl From<etagere::Point> for Point<DevicePixels> {
261 fn from(value: etagere::Point) -> Self {
262 Point {
263 x: DevicePixels::from(value.x),
264 y: DevicePixels::from(value.y),
265 }
266 }
267}
268
269impl From<etagere::Size> for Size<DevicePixels> {
270 fn from(size: etagere::Size) -> Self {
271 Size {
272 width: DevicePixels::from(size.width),
273 height: DevicePixels::from(size.height),
274 }
275 }
276}
277
278impl From<etagere::Rectangle> for Bounds<DevicePixels> {
279 fn from(rectangle: etagere::Rectangle) -> Self {
280 Bounds {
281 origin: rectangle.min.into(),
282 size: rectangle.size().into(),
283 }
284 }
285}
286
287#[derive(Deref, DerefMut)]
288struct AssertSend<T>(T);
289
290unsafe impl<T> Send for AssertSend<T> {}