1use anyhow::{Context as _, Result};
2use collections::FxHashMap;
3use etagere::{BucketedAtlasAllocator, size2};
4use gpui::{
5 AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTextureList, AtlasTile, Bounds, DevicePixels,
6 PlatformAtlas, Point, Size,
7};
8use parking_lot::Mutex;
9use std::{borrow::Cow, ops, sync::Arc};
10
11use crate::WgpuContext;
12
13fn device_size_to_etagere(size: Size<DevicePixels>) -> etagere::Size {
14 size2(size.width.0, size.height.0)
15}
16
17fn etagere_point_to_device(point: etagere::Point) -> Point<DevicePixels> {
18 Point {
19 x: DevicePixels(point.x),
20 y: DevicePixels(point.y),
21 }
22}
23
24pub struct WgpuAtlas(Mutex<WgpuAtlasState>);
25
26struct PendingUpload {
27 id: AtlasTextureId,
28 bounds: Bounds<DevicePixels>,
29 data: Vec<u8>,
30}
31
32struct WgpuAtlasState {
33 device: Arc<wgpu::Device>,
34 queue: Arc<wgpu::Queue>,
35 max_texture_size: u32,
36 color_texture_format: wgpu::TextureFormat,
37 storage: WgpuAtlasStorage,
38 tiles_by_key: FxHashMap<AtlasKey, AtlasTile>,
39 pending_uploads: Vec<PendingUpload>,
40}
41
42pub struct WgpuTextureInfo {
43 pub view: wgpu::TextureView,
44}
45
46impl WgpuAtlas {
47 pub fn new(
48 device: Arc<wgpu::Device>,
49 queue: Arc<wgpu::Queue>,
50 color_texture_format: wgpu::TextureFormat,
51 ) -> Self {
52 let max_texture_size = device.limits().max_texture_dimension_2d;
53 WgpuAtlas(Mutex::new(WgpuAtlasState {
54 device,
55 queue,
56 max_texture_size,
57 color_texture_format,
58 storage: WgpuAtlasStorage::default(),
59 tiles_by_key: Default::default(),
60 pending_uploads: Vec::new(),
61 }))
62 }
63
64 pub fn from_context(context: &WgpuContext) -> Self {
65 Self::new(
66 context.device.clone(),
67 context.queue.clone(),
68 context.color_texture_format(),
69 )
70 }
71
72 pub fn before_frame(&self) {
73 let mut lock = self.0.lock();
74 lock.flush_uploads();
75 }
76
77 pub fn get_texture_info(&self, id: AtlasTextureId) -> WgpuTextureInfo {
78 let lock = self.0.lock();
79 let texture = &lock.storage[id];
80 WgpuTextureInfo {
81 view: texture.view.clone(),
82 }
83 }
84
85 /// Handles device lost by clearing all textures and cached tiles.
86 /// The atlas will lazily recreate textures as needed on subsequent frames.
87 pub fn handle_device_lost(&self, context: &WgpuContext) {
88 let mut lock = self.0.lock();
89 lock.device = context.device.clone();
90 lock.queue = context.queue.clone();
91 lock.color_texture_format = context.color_texture_format();
92 lock.storage = WgpuAtlasStorage::default();
93 lock.tiles_by_key.clear();
94 lock.pending_uploads.clear();
95 }
96}
97
98impl PlatformAtlas for WgpuAtlas {
99 fn get_or_insert_with<'a>(
100 &self,
101 key: &AtlasKey,
102 build: &mut dyn FnMut() -> Result<Option<(Size<DevicePixels>, Cow<'a, [u8]>)>>,
103 ) -> Result<Option<AtlasTile>> {
104 let mut lock = self.0.lock();
105 if let Some(tile) = lock.tiles_by_key.get(key) {
106 Ok(Some(tile.clone()))
107 } else {
108 profiling::scope!("new tile");
109 let Some((size, bytes)) = build()? else {
110 return Ok(None);
111 };
112 let tile = lock
113 .allocate(size, key.texture_kind())
114 .context("failed to allocate")?;
115 lock.upload_texture(tile.texture_id, tile.bounds, bytes);
116 lock.tiles_by_key.insert(key.clone(), tile.clone());
117 Ok(Some(tile))
118 }
119 }
120
121 fn remove(&self, key: &AtlasKey) {
122 let mut lock = self.0.lock();
123
124 let Some(id) = lock.tiles_by_key.remove(key).map(|tile| tile.texture_id) else {
125 return;
126 };
127
128 let Some(texture_slot) = lock.storage[id.kind].textures.get_mut(id.index as usize) else {
129 return;
130 };
131
132 if let Some(mut texture) = texture_slot.take() {
133 texture.decrement_ref_count();
134 if texture.is_unreferenced() {
135 lock.pending_uploads
136 .retain(|upload| upload.id != texture.id);
137 lock.storage[id.kind]
138 .free_list
139 .push(texture.id.index as usize);
140 } else {
141 *texture_slot = Some(texture);
142 }
143 }
144 }
145}
146
147impl WgpuAtlasState {
148 fn allocate(
149 &mut self,
150 size: Size<DevicePixels>,
151 texture_kind: AtlasTextureKind,
152 ) -> Option<AtlasTile> {
153 {
154 let textures = &mut self.storage[texture_kind];
155
156 if let Some(tile) = textures
157 .iter_mut()
158 .rev()
159 .find_map(|texture| texture.allocate(size))
160 {
161 return Some(tile);
162 }
163 }
164
165 let texture = self.push_texture(size, texture_kind);
166 texture.allocate(size)
167 }
168
169 fn push_texture(
170 &mut self,
171 min_size: Size<DevicePixels>,
172 kind: AtlasTextureKind,
173 ) -> &mut WgpuAtlasTexture {
174 const DEFAULT_ATLAS_SIZE: Size<DevicePixels> = Size {
175 width: DevicePixels(1024),
176 height: DevicePixels(1024),
177 };
178 let max_texture_size = self.max_texture_size as i32;
179 let max_atlas_size = Size {
180 width: DevicePixels(max_texture_size),
181 height: DevicePixels(max_texture_size),
182 };
183
184 let size = min_size.min(&max_atlas_size).max(&DEFAULT_ATLAS_SIZE);
185 let format = match kind {
186 AtlasTextureKind::Monochrome => wgpu::TextureFormat::R8Unorm,
187 AtlasTextureKind::Subpixel | AtlasTextureKind::Polychrome => self.color_texture_format,
188 };
189
190 let texture = self.device.create_texture(&wgpu::TextureDescriptor {
191 label: Some("atlas"),
192 size: wgpu::Extent3d {
193 width: size.width.0 as u32,
194 height: size.height.0 as u32,
195 depth_or_array_layers: 1,
196 },
197 mip_level_count: 1,
198 sample_count: 1,
199 dimension: wgpu::TextureDimension::D2,
200 format,
201 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
202 view_formats: &[],
203 });
204
205 let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
206
207 let texture_list = &mut self.storage[kind];
208 let index = texture_list.free_list.pop();
209
210 let atlas_texture = WgpuAtlasTexture {
211 id: AtlasTextureId {
212 index: index.unwrap_or(texture_list.textures.len()) as u32,
213 kind,
214 },
215 allocator: BucketedAtlasAllocator::new(device_size_to_etagere(size)),
216 format,
217 texture,
218 view,
219 live_atlas_keys: 0,
220 };
221
222 if let Some(ix) = index {
223 texture_list.textures[ix] = Some(atlas_texture);
224 texture_list
225 .textures
226 .get_mut(ix)
227 .and_then(|t| t.as_mut())
228 .expect("texture must exist")
229 } else {
230 texture_list.textures.push(Some(atlas_texture));
231 texture_list
232 .textures
233 .last_mut()
234 .and_then(|t| t.as_mut())
235 .expect("texture must exist")
236 }
237 }
238
239 fn upload_texture(
240 &mut self,
241 id: AtlasTextureId,
242 bounds: Bounds<DevicePixels>,
243 bytes: Cow<'_, [u8]>,
244 ) {
245 if let Some(texture) = self.storage.get(id) {
246 let data = swizzle_upload_data(bytes, texture.format);
247 self.pending_uploads
248 .push(PendingUpload { id, bounds, data });
249 }
250 }
251
252 fn flush_uploads(&mut self) {
253 for upload in self.pending_uploads.drain(..) {
254 let Some(texture) = self.storage.get(upload.id) else {
255 continue;
256 };
257 let bytes_per_pixel = texture.bytes_per_pixel();
258
259 self.queue.write_texture(
260 wgpu::TexelCopyTextureInfo {
261 texture: &texture.texture,
262 mip_level: 0,
263 origin: wgpu::Origin3d {
264 x: upload.bounds.origin.x.0 as u32,
265 y: upload.bounds.origin.y.0 as u32,
266 z: 0,
267 },
268 aspect: wgpu::TextureAspect::All,
269 },
270 &upload.data,
271 wgpu::TexelCopyBufferLayout {
272 offset: 0,
273 bytes_per_row: Some(upload.bounds.size.width.0 as u32 * bytes_per_pixel as u32),
274 rows_per_image: None,
275 },
276 wgpu::Extent3d {
277 width: upload.bounds.size.width.0 as u32,
278 height: upload.bounds.size.height.0 as u32,
279 depth_or_array_layers: 1,
280 },
281 );
282 }
283 }
284}
285
286#[derive(Default)]
287struct WgpuAtlasStorage {
288 monochrome_textures: AtlasTextureList<WgpuAtlasTexture>,
289 subpixel_textures: AtlasTextureList<WgpuAtlasTexture>,
290 polychrome_textures: AtlasTextureList<WgpuAtlasTexture>,
291}
292
293impl ops::Index<AtlasTextureKind> for WgpuAtlasStorage {
294 type Output = AtlasTextureList<WgpuAtlasTexture>;
295 fn index(&self, kind: AtlasTextureKind) -> &Self::Output {
296 match kind {
297 AtlasTextureKind::Monochrome => &self.monochrome_textures,
298 AtlasTextureKind::Subpixel => &self.subpixel_textures,
299 AtlasTextureKind::Polychrome => &self.polychrome_textures,
300 }
301 }
302}
303
304impl ops::IndexMut<AtlasTextureKind> for WgpuAtlasStorage {
305 fn index_mut(&mut self, kind: AtlasTextureKind) -> &mut Self::Output {
306 match kind {
307 AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
308 AtlasTextureKind::Subpixel => &mut self.subpixel_textures,
309 AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
310 }
311 }
312}
313
314impl WgpuAtlasStorage {
315 fn get(&self, id: AtlasTextureId) -> Option<&WgpuAtlasTexture> {
316 self[id.kind]
317 .textures
318 .get(id.index as usize)
319 .and_then(|t| t.as_ref())
320 }
321}
322
323impl ops::Index<AtlasTextureId> for WgpuAtlasStorage {
324 type Output = WgpuAtlasTexture;
325 fn index(&self, id: AtlasTextureId) -> &Self::Output {
326 let textures = match id.kind {
327 AtlasTextureKind::Monochrome => &self.monochrome_textures,
328 AtlasTextureKind::Subpixel => &self.subpixel_textures,
329 AtlasTextureKind::Polychrome => &self.polychrome_textures,
330 };
331 textures[id.index as usize]
332 .as_ref()
333 .expect("texture must exist")
334 }
335}
336
337struct WgpuAtlasTexture {
338 id: AtlasTextureId,
339 allocator: BucketedAtlasAllocator,
340 texture: wgpu::Texture,
341 view: wgpu::TextureView,
342 format: wgpu::TextureFormat,
343 live_atlas_keys: u32,
344}
345
346impl WgpuAtlasTexture {
347 fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
348 let allocation = self.allocator.allocate(device_size_to_etagere(size))?;
349 let tile = AtlasTile {
350 texture_id: self.id,
351 tile_id: allocation.id.into(),
352 padding: 0,
353 bounds: Bounds {
354 origin: etagere_point_to_device(allocation.rectangle.min),
355 size,
356 },
357 };
358 self.live_atlas_keys += 1;
359 Some(tile)
360 }
361
362 fn bytes_per_pixel(&self) -> u8 {
363 match self.format {
364 wgpu::TextureFormat::R8Unorm => 1,
365 wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Rgba8Unorm => 4,
366 _ => 4,
367 }
368 }
369
370 fn decrement_ref_count(&mut self) {
371 self.live_atlas_keys -= 1;
372 }
373
374 fn is_unreferenced(&self) -> bool {
375 self.live_atlas_keys == 0
376 }
377}
378
379fn swizzle_upload_data(bytes: Cow<'_, [u8]>, format: wgpu::TextureFormat) -> Vec<u8> {
380 match format {
381 wgpu::TextureFormat::Rgba8Unorm => {
382 debug_assert_eq!(
383 bytes.len() % 4,
384 0,
385 "upload data must be a multiple of 4 bytes"
386 );
387 let mut data = bytes.into_owned();
388 for pixel in data.chunks_exact_mut(4) {
389 pixel.swap(0, 2);
390 }
391 data
392 }
393 _ => bytes.into_owned(),
394 }
395}
396
397#[cfg(all(test, not(target_family = "wasm")))]
398mod tests {
399 use super::*;
400 use gpui::{ImageId, RenderImageParams};
401 use pollster::block_on;
402 use std::sync::Arc;
403
404 fn test_device_and_queue() -> anyhow::Result<(Arc<wgpu::Device>, Arc<wgpu::Queue>)> {
405 block_on(async {
406 let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
407 backends: wgpu::Backends::all(),
408 flags: wgpu::InstanceFlags::default(),
409 backend_options: wgpu::BackendOptions::default(),
410 memory_budget_thresholds: wgpu::MemoryBudgetThresholds::default(),
411 display: None,
412 });
413 let adapter = instance
414 .request_adapter(&wgpu::RequestAdapterOptions {
415 power_preference: wgpu::PowerPreference::LowPower,
416 compatible_surface: None,
417 force_fallback_adapter: false,
418 })
419 .await
420 .map_err(|error| anyhow::anyhow!("failed to request adapter: {error}"))?;
421 let (device, queue) = adapter
422 .request_device(&wgpu::DeviceDescriptor {
423 label: Some("wgpu_atlas_test_device"),
424 required_features: wgpu::Features::empty(),
425 required_limits: wgpu::Limits::downlevel_defaults()
426 .using_resolution(adapter.limits())
427 .using_alignment(adapter.limits()),
428 memory_hints: wgpu::MemoryHints::MemoryUsage,
429 trace: wgpu::Trace::Off,
430 experimental_features: wgpu::ExperimentalFeatures::disabled(),
431 })
432 .await
433 .map_err(|error| anyhow::anyhow!("failed to request device: {error}"))?;
434 Ok((Arc::new(device), Arc::new(queue)))
435 })
436 }
437
438 #[test]
439 fn before_frame_skips_uploads_for_removed_texture() -> anyhow::Result<()> {
440 let (device, queue) = test_device_and_queue()?;
441
442 let atlas = WgpuAtlas::new(device, queue, wgpu::TextureFormat::Bgra8Unorm);
443 let key = AtlasKey::Image(RenderImageParams {
444 image_id: ImageId(1),
445 frame_index: 0,
446 });
447 let size = Size {
448 width: DevicePixels(1),
449 height: DevicePixels(1),
450 };
451 let mut build = || Ok(Some((size, Cow::Owned(vec![0, 0, 0, 255]))));
452
453 // Regression test: before the fix, this panicked in flush_uploads
454 atlas
455 .get_or_insert_with(&key, &mut build)?
456 .expect("tile should be created");
457 atlas.remove(&key);
458 atlas.before_frame();
459 Ok(())
460 }
461
462 #[test]
463 fn swizzle_upload_data_preserves_bgra_uploads() {
464 let input = vec![0x10, 0x20, 0x30, 0x40];
465 assert_eq!(
466 swizzle_upload_data(Cow::Borrowed(&input), wgpu::TextureFormat::Bgra8Unorm),
467 input
468 );
469 }
470
471 #[test]
472 fn swizzle_upload_data_converts_bgra_to_rgba() {
473 let input = vec![0x10, 0x20, 0x30, 0x40, 0xAA, 0xBB, 0xCC, 0xDD];
474 assert_eq!(
475 swizzle_upload_data(Cow::Borrowed(&input), wgpu::TextureFormat::Rgba8Unorm),
476 vec![0x30, 0x20, 0x10, 0x40, 0xCC, 0xBB, 0xAA, 0xDD]
477 );
478 }
479}