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