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(&mut self, id: AtlasTextureId, bounds: Bounds<DevicePixels>, bytes: &[u8]) {
240 let data = self
241 .storage
242 .get(id)
243 .map(|texture| swizzle_upload_data(bytes, texture.format))
244 .unwrap_or_else(|| bytes.to_vec());
245
246 self.pending_uploads
247 .push(PendingUpload { id, bounds, data });
248 }
249
250 fn flush_uploads(&mut self) {
251 for upload in self.pending_uploads.drain(..) {
252 let Some(texture) = self.storage.get(upload.id) else {
253 continue;
254 };
255 let bytes_per_pixel = texture.bytes_per_pixel();
256
257 self.queue.write_texture(
258 wgpu::TexelCopyTextureInfo {
259 texture: &texture.texture,
260 mip_level: 0,
261 origin: wgpu::Origin3d {
262 x: upload.bounds.origin.x.0 as u32,
263 y: upload.bounds.origin.y.0 as u32,
264 z: 0,
265 },
266 aspect: wgpu::TextureAspect::All,
267 },
268 &upload.data,
269 wgpu::TexelCopyBufferLayout {
270 offset: 0,
271 bytes_per_row: Some(upload.bounds.size.width.0 as u32 * bytes_per_pixel as u32),
272 rows_per_image: None,
273 },
274 wgpu::Extent3d {
275 width: upload.bounds.size.width.0 as u32,
276 height: upload.bounds.size.height.0 as u32,
277 depth_or_array_layers: 1,
278 },
279 );
280 }
281 }
282}
283
284#[derive(Default)]
285struct WgpuAtlasStorage {
286 monochrome_textures: AtlasTextureList<WgpuAtlasTexture>,
287 subpixel_textures: AtlasTextureList<WgpuAtlasTexture>,
288 polychrome_textures: AtlasTextureList<WgpuAtlasTexture>,
289}
290
291impl ops::Index<AtlasTextureKind> for WgpuAtlasStorage {
292 type Output = AtlasTextureList<WgpuAtlasTexture>;
293 fn index(&self, kind: AtlasTextureKind) -> &Self::Output {
294 match kind {
295 AtlasTextureKind::Monochrome => &self.monochrome_textures,
296 AtlasTextureKind::Subpixel => &self.subpixel_textures,
297 AtlasTextureKind::Polychrome => &self.polychrome_textures,
298 }
299 }
300}
301
302impl ops::IndexMut<AtlasTextureKind> for WgpuAtlasStorage {
303 fn index_mut(&mut self, kind: AtlasTextureKind) -> &mut Self::Output {
304 match kind {
305 AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
306 AtlasTextureKind::Subpixel => &mut self.subpixel_textures,
307 AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
308 }
309 }
310}
311
312impl WgpuAtlasStorage {
313 fn get(&self, id: AtlasTextureId) -> Option<&WgpuAtlasTexture> {
314 self[id.kind]
315 .textures
316 .get(id.index as usize)
317 .and_then(|t| t.as_ref())
318 }
319}
320
321impl ops::Index<AtlasTextureId> for WgpuAtlasStorage {
322 type Output = WgpuAtlasTexture;
323 fn index(&self, id: AtlasTextureId) -> &Self::Output {
324 let textures = match id.kind {
325 AtlasTextureKind::Monochrome => &self.monochrome_textures,
326 AtlasTextureKind::Subpixel => &self.subpixel_textures,
327 AtlasTextureKind::Polychrome => &self.polychrome_textures,
328 };
329 textures[id.index as usize]
330 .as_ref()
331 .expect("texture must exist")
332 }
333}
334
335struct WgpuAtlasTexture {
336 id: AtlasTextureId,
337 allocator: BucketedAtlasAllocator,
338 texture: wgpu::Texture,
339 view: wgpu::TextureView,
340 format: wgpu::TextureFormat,
341 live_atlas_keys: u32,
342}
343
344impl WgpuAtlasTexture {
345 fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
346 let allocation = self.allocator.allocate(device_size_to_etagere(size))?;
347 let tile = AtlasTile {
348 texture_id: self.id,
349 tile_id: allocation.id.into(),
350 padding: 0,
351 bounds: Bounds {
352 origin: etagere_point_to_device(allocation.rectangle.min),
353 size,
354 },
355 };
356 self.live_atlas_keys += 1;
357 Some(tile)
358 }
359
360 fn bytes_per_pixel(&self) -> u8 {
361 match self.format {
362 wgpu::TextureFormat::R8Unorm => 1,
363 wgpu::TextureFormat::Bgra8Unorm | wgpu::TextureFormat::Rgba8Unorm => 4,
364 _ => 4,
365 }
366 }
367
368 fn decrement_ref_count(&mut self) {
369 self.live_atlas_keys -= 1;
370 }
371
372 fn is_unreferenced(&self) -> bool {
373 self.live_atlas_keys == 0
374 }
375}
376
377fn swizzle_upload_data(bytes: &[u8], format: wgpu::TextureFormat) -> Vec<u8> {
378 match format {
379 wgpu::TextureFormat::Rgba8Unorm => {
380 let mut data = bytes.to_vec();
381 for pixel in data.chunks_exact_mut(4) {
382 pixel.swap(0, 2);
383 }
384 data
385 }
386 _ => bytes.to_vec(),
387 }
388}
389
390#[cfg(all(test, not(target_family = "wasm")))]
391mod tests {
392 use super::*;
393 use gpui::{ImageId, RenderImageParams};
394 use pollster::block_on;
395 use std::sync::Arc;
396
397 fn test_device_and_queue() -> anyhow::Result<(Arc<wgpu::Device>, Arc<wgpu::Queue>)> {
398 block_on(async {
399 let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
400 backends: wgpu::Backends::all(),
401 flags: wgpu::InstanceFlags::default(),
402 backend_options: wgpu::BackendOptions::default(),
403 memory_budget_thresholds: wgpu::MemoryBudgetThresholds::default(),
404 display: None,
405 });
406 let adapter = instance
407 .request_adapter(&wgpu::RequestAdapterOptions {
408 power_preference: wgpu::PowerPreference::LowPower,
409 compatible_surface: None,
410 force_fallback_adapter: false,
411 })
412 .await
413 .map_err(|error| anyhow::anyhow!("failed to request adapter: {error}"))?;
414 let (device, queue) = adapter
415 .request_device(&wgpu::DeviceDescriptor {
416 label: Some("wgpu_atlas_test_device"),
417 required_features: wgpu::Features::empty(),
418 required_limits: wgpu::Limits::downlevel_defaults()
419 .using_resolution(adapter.limits())
420 .using_alignment(adapter.limits()),
421 memory_hints: wgpu::MemoryHints::MemoryUsage,
422 trace: wgpu::Trace::Off,
423 experimental_features: wgpu::ExperimentalFeatures::disabled(),
424 })
425 .await
426 .map_err(|error| anyhow::anyhow!("failed to request device: {error}"))?;
427 Ok((Arc::new(device), Arc::new(queue)))
428 })
429 }
430
431 #[test]
432 fn before_frame_skips_uploads_for_removed_texture() -> anyhow::Result<()> {
433 let (device, queue) = test_device_and_queue()?;
434
435 let atlas = WgpuAtlas::new(device, queue, wgpu::TextureFormat::Bgra8Unorm);
436 let key = AtlasKey::Image(RenderImageParams {
437 image_id: ImageId(1),
438 frame_index: 0,
439 });
440 let size = Size {
441 width: DevicePixels(1),
442 height: DevicePixels(1),
443 };
444 let mut build = || Ok(Some((size, Cow::Owned(vec![0, 0, 0, 255]))));
445
446 // Regression test: before the fix, this panicked in flush_uploads
447 atlas
448 .get_or_insert_with(&key, &mut build)?
449 .expect("tile should be created");
450 atlas.remove(&key);
451 atlas.before_frame();
452 Ok(())
453 }
454
455 #[test]
456 fn swizzle_upload_data_preserves_bgra_uploads() {
457 let input = vec![0x10, 0x20, 0x30, 0x40];
458 assert_eq!(
459 swizzle_upload_data(&input, wgpu::TextureFormat::Bgra8Unorm),
460 input
461 );
462 }
463
464 #[test]
465 fn swizzle_upload_data_converts_bgra_to_rgba() {
466 let input = vec![0x10, 0x20, 0x30, 0x40, 0xAA, 0xBB, 0xCC, 0xDD];
467 assert_eq!(
468 swizzle_upload_data(&input, wgpu::TextureFormat::Rgba8Unorm),
469 vec![0x30, 0x20, 0x10, 0x40, 0xCC, 0xBB, 0xAA, 0xDD]
470 );
471 }
472}