1use crate::{
2 AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTile, Bounds, DevicePixels, PlatformAtlas,
3 Point, Size, platform::AtlasTextureList,
4};
5use anyhow::Result;
6use collections::FxHashMap;
7use etagere::{BucketedAtlasAllocator, size2};
8use parking_lot::Mutex;
9use std::{borrow::Cow, ops, sync::Arc};
10
11fn device_size_to_etagere(size: Size<DevicePixels>) -> etagere::Size {
12 size2(size.width.0, size.height.0)
13}
14
15fn etagere_point_to_device(point: etagere::Point) -> Point<DevicePixels> {
16 Point {
17 x: DevicePixels(point.x),
18 y: DevicePixels(point.y),
19 }
20}
21
22pub(crate) struct WgpuAtlas(Mutex<WgpuAtlasState>);
23
24struct PendingUpload {
25 id: AtlasTextureId,
26 bounds: Bounds<DevicePixels>,
27 data: Vec<u8>,
28}
29
30struct WgpuAtlasState {
31 device: Arc<wgpu::Device>,
32 queue: Arc<wgpu::Queue>,
33 storage: WgpuAtlasStorage,
34 tiles_by_key: FxHashMap<AtlasKey, AtlasTile>,
35 pending_uploads: Vec<PendingUpload>,
36}
37
38pub struct WgpuTextureInfo {
39 pub view: wgpu::TextureView,
40}
41
42impl WgpuAtlas {
43 pub(crate) fn new(device: Arc<wgpu::Device>, queue: Arc<wgpu::Queue>) -> Self {
44 WgpuAtlas(Mutex::new(WgpuAtlasState {
45 device,
46 queue,
47 storage: WgpuAtlasStorage::default(),
48 tiles_by_key: Default::default(),
49 pending_uploads: Vec::new(),
50 }))
51 }
52
53 pub fn before_frame(&self) {
54 let mut lock = self.0.lock();
55 lock.flush_uploads();
56 }
57
58 pub fn get_texture_info(&self, id: AtlasTextureId) -> WgpuTextureInfo {
59 let lock = self.0.lock();
60 let texture = &lock.storage[id];
61 WgpuTextureInfo {
62 view: texture.view.clone(),
63 }
64 }
65}
66
67impl PlatformAtlas for WgpuAtlas {
68 fn get_or_insert_with<'a>(
69 &self,
70 key: &AtlasKey,
71 build: &mut dyn FnMut() -> Result<Option<(Size<DevicePixels>, Cow<'a, [u8]>)>>,
72 ) -> Result<Option<AtlasTile>> {
73 let mut lock = self.0.lock();
74 if let Some(tile) = lock.tiles_by_key.get(key) {
75 Ok(Some(tile.clone()))
76 } else {
77 profiling::scope!("new tile");
78 let Some((size, bytes)) = build()? else {
79 return Ok(None);
80 };
81 let tile = lock.allocate(size, key.texture_kind());
82 lock.upload_texture(tile.texture_id, 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
91 let Some(id) = lock.tiles_by_key.remove(key).map(|tile| tile.texture_id) else {
92 return;
93 };
94
95 let Some(texture_slot) = lock.storage[id.kind].textures.get_mut(id.index as usize) else {
96 return;
97 };
98
99 if let Some(mut texture) = texture_slot.take() {
100 texture.decrement_ref_count();
101 if texture.is_unreferenced() {
102 lock.storage[id.kind]
103 .free_list
104 .push(texture.id.index as usize);
105 } else {
106 *texture_slot = Some(texture);
107 }
108 }
109 }
110}
111
112impl WgpuAtlasState {
113 fn allocate(&mut self, size: Size<DevicePixels>, texture_kind: AtlasTextureKind) -> AtlasTile {
114 {
115 let textures = &mut self.storage[texture_kind];
116
117 if let Some(tile) = textures
118 .iter_mut()
119 .rev()
120 .find_map(|texture| texture.allocate(size))
121 {
122 return tile;
123 }
124 }
125
126 let texture = self.push_texture(size, texture_kind);
127 texture
128 .allocate(size)
129 .expect("Failed to allocate from newly created texture")
130 }
131
132 fn push_texture(
133 &mut self,
134 min_size: Size<DevicePixels>,
135 kind: AtlasTextureKind,
136 ) -> &mut WgpuAtlasTexture {
137 const DEFAULT_ATLAS_SIZE: Size<DevicePixels> = Size {
138 width: DevicePixels(1024),
139 height: DevicePixels(1024),
140 };
141
142 let size = min_size.max(&DEFAULT_ATLAS_SIZE);
143 let format = match kind {
144 AtlasTextureKind::Monochrome => wgpu::TextureFormat::R8Unorm,
145 AtlasTextureKind::Subpixel => wgpu::TextureFormat::Bgra8Unorm,
146 AtlasTextureKind::Polychrome => wgpu::TextureFormat::Bgra8Unorm,
147 };
148
149 let texture = self.device.create_texture(&wgpu::TextureDescriptor {
150 label: Some("atlas"),
151 size: wgpu::Extent3d {
152 width: size.width.0 as u32,
153 height: size.height.0 as u32,
154 depth_or_array_layers: 1,
155 },
156 mip_level_count: 1,
157 sample_count: 1,
158 dimension: wgpu::TextureDimension::D2,
159 format,
160 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
161 view_formats: &[],
162 });
163
164 let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
165
166 let texture_list = &mut self.storage[kind];
167 let index = texture_list.free_list.pop();
168
169 let atlas_texture = WgpuAtlasTexture {
170 id: AtlasTextureId {
171 index: index.unwrap_or(texture_list.textures.len()) as u32,
172 kind,
173 },
174 allocator: BucketedAtlasAllocator::new(device_size_to_etagere(size)),
175 format,
176 texture,
177 view,
178 live_atlas_keys: 0,
179 };
180
181 if let Some(ix) = index {
182 texture_list.textures[ix] = Some(atlas_texture);
183 texture_list
184 .textures
185 .get_mut(ix)
186 .and_then(|t| t.as_mut())
187 .expect("texture must exist")
188 } else {
189 texture_list.textures.push(Some(atlas_texture));
190 texture_list
191 .textures
192 .last_mut()
193 .and_then(|t| t.as_mut())
194 .expect("texture must exist")
195 }
196 }
197
198 fn upload_texture(&mut self, id: AtlasTextureId, bounds: Bounds<DevicePixels>, bytes: &[u8]) {
199 self.pending_uploads.push(PendingUpload {
200 id,
201 bounds,
202 data: bytes.to_vec(),
203 });
204 }
205
206 fn flush_uploads(&mut self) {
207 for upload in self.pending_uploads.drain(..) {
208 let texture = &self.storage[upload.id];
209 let bytes_per_pixel = texture.bytes_per_pixel();
210
211 self.queue.write_texture(
212 wgpu::TexelCopyTextureInfo {
213 texture: &texture.texture,
214 mip_level: 0,
215 origin: wgpu::Origin3d {
216 x: upload.bounds.origin.x.0 as u32,
217 y: upload.bounds.origin.y.0 as u32,
218 z: 0,
219 },
220 aspect: wgpu::TextureAspect::All,
221 },
222 &upload.data,
223 wgpu::TexelCopyBufferLayout {
224 offset: 0,
225 bytes_per_row: Some(upload.bounds.size.width.0 as u32 * bytes_per_pixel as u32),
226 rows_per_image: None,
227 },
228 wgpu::Extent3d {
229 width: upload.bounds.size.width.0 as u32,
230 height: upload.bounds.size.height.0 as u32,
231 depth_or_array_layers: 1,
232 },
233 );
234 }
235 }
236}
237
238#[derive(Default)]
239struct WgpuAtlasStorage {
240 monochrome_textures: AtlasTextureList<WgpuAtlasTexture>,
241 subpixel_textures: AtlasTextureList<WgpuAtlasTexture>,
242 polychrome_textures: AtlasTextureList<WgpuAtlasTexture>,
243}
244
245impl ops::Index<AtlasTextureKind> for WgpuAtlasStorage {
246 type Output = AtlasTextureList<WgpuAtlasTexture>;
247 fn index(&self, kind: AtlasTextureKind) -> &Self::Output {
248 match kind {
249 AtlasTextureKind::Monochrome => &self.monochrome_textures,
250 AtlasTextureKind::Subpixel => &self.subpixel_textures,
251 AtlasTextureKind::Polychrome => &self.polychrome_textures,
252 }
253 }
254}
255
256impl ops::IndexMut<AtlasTextureKind> for WgpuAtlasStorage {
257 fn index_mut(&mut self, kind: AtlasTextureKind) -> &mut Self::Output {
258 match kind {
259 AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
260 AtlasTextureKind::Subpixel => &mut self.subpixel_textures,
261 AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
262 }
263 }
264}
265
266impl ops::Index<AtlasTextureId> for WgpuAtlasStorage {
267 type Output = WgpuAtlasTexture;
268 fn index(&self, id: AtlasTextureId) -> &Self::Output {
269 let textures = match id.kind {
270 AtlasTextureKind::Monochrome => &self.monochrome_textures,
271 AtlasTextureKind::Subpixel => &self.subpixel_textures,
272 AtlasTextureKind::Polychrome => &self.polychrome_textures,
273 };
274 textures[id.index as usize]
275 .as_ref()
276 .expect("texture must exist")
277 }
278}
279
280struct WgpuAtlasTexture {
281 id: AtlasTextureId,
282 allocator: BucketedAtlasAllocator,
283 texture: wgpu::Texture,
284 view: wgpu::TextureView,
285 format: wgpu::TextureFormat,
286 live_atlas_keys: u32,
287}
288
289impl WgpuAtlasTexture {
290 fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
291 let allocation = self.allocator.allocate(device_size_to_etagere(size))?;
292 let tile = AtlasTile {
293 texture_id: self.id,
294 tile_id: allocation.id.into(),
295 padding: 0,
296 bounds: Bounds {
297 origin: etagere_point_to_device(allocation.rectangle.min),
298 size,
299 },
300 };
301 self.live_atlas_keys += 1;
302 Some(tile)
303 }
304
305 fn bytes_per_pixel(&self) -> u8 {
306 match self.format {
307 wgpu::TextureFormat::R8Unorm => 1,
308 wgpu::TextureFormat::Bgra8Unorm => 4,
309 _ => 4,
310 }
311 }
312
313 fn decrement_ref_count(&mut self) {
314 self.live_atlas_keys -= 1;
315 }
316
317 fn is_unreferenced(&self) -> bool {
318 self.live_atlas_keys == 0
319 }
320}