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