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