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