Krolyxon
Commit 8b10288

Branches

Tags

Refs →

remagesrcrenderatlas.rs

404 lines · 13.8 KB · Rust

Raw Blame History Download
1//! Thumbnail atlas: square slots in texture pages, evicted least-recently-drawn first.
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5
6use bytemuck::{Pod, Zeroable};
7use wgpu::util::DeviceExt;
8
9use crate::thumbs::Thumb;
10
11const FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba8UnormSrgb;
12/// Preferred page side; smaller if the GPU can't do it, larger for huge cells.
13const PAGE: u32 = 2048;
14
15/// A thumbnail to draw: `rect` in physical pixels.
16pub struct ThumbDraw<'a> {
17 pub path: &'a Path,
18 pub rect: [f32; 4],
19 pub alpha: f32,
20}
21
22#[repr(C)]
23#[derive(Clone, Copy, Pod, Zeroable)]
24struct Instance {
25 rect: [f32; 4],
26 uv: [f32; 4],
27 params: [f32; 4],
28}
29
30const ATTRS: [wgpu::VertexAttribute; 3] =
31 wgpu::vertex_attr_array![0 => Float32x4, 1 => Float32x4, 2 => Float32x4];
32
33struct Page {
34 texture: wgpu::Texture,
35 bind_group: wgpu::BindGroup,
36}
37
38#[derive(Clone, Copy)]
39struct Slot {
40 page: usize,
41 index: u32,
42 size: [u32; 2],
43 used: u64,
44}
45
46pub struct Atlas {
47 pipeline: wgpu::RenderPipeline,
48 globals: wgpu::Buffer,
49 globals_group: wgpu::BindGroup,
50 page_layout: wgpu::BindGroupLayout,
51 instances: wgpu::Buffer,
52 capacity: usize,
53 /// Per page: instance range drawn this frame.
54 batches: Vec<(usize, std::ops::Range<u32>)>,
55 max_dim: u32,
56 page_side: u32,
57 /// Slot side: the cell plus a 1 px transparent border all round.
58 slot: u32,
59 pages: Vec<Page>,
60 slots: HashMap<PathBuf, Slot>,
61 free: Vec<(usize, u32)>,
62 frame: u64,
63}
64
65fn instance_buffer(device: &wgpu::Device, capacity: usize) -> wgpu::Buffer {
66 device.create_buffer(&wgpu::BufferDescriptor {
67 label: Some("thumb instances"),
68 size: (capacity * size_of::<Instance>()) as u64,
69 usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
70 mapped_at_creation: false,
71 })
72}
73
74impl Atlas {
75 pub fn new(device: &wgpu::Device, target: wgpu::TextureFormat, max_dim: u32) -> Self {
76 let globals = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
77 label: Some("thumb globals"),
78 contents: bytemuck::bytes_of(&[0.0f32; 4]),
79 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
80 });
81 let globals_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
82 label: Some("thumbs"),
83 entries: &[
84 wgpu::BindGroupLayoutEntry {
85 binding: 0,
86 visibility: wgpu::ShaderStages::VERTEX,
87 ty: wgpu::BindingType::Buffer {
88 ty: wgpu::BufferBindingType::Uniform,
89 has_dynamic_offset: false,
90 min_binding_size: None,
91 },
92 count: None,
93 },
94 wgpu::BindGroupLayoutEntry {
95 binding: 1,
96 visibility: wgpu::ShaderStages::FRAGMENT,
97 ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
98 count: None,
99 },
100 ],
101 });
102 let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
103 label: Some("thumbs"),
104 mag_filter: wgpu::FilterMode::Linear,
105 min_filter: wgpu::FilterMode::Linear,
106 ..Default::default()
107 });
108 let globals_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
109 label: Some("thumbs"),
110 layout: &globals_layout,
111 entries: &[
112 wgpu::BindGroupEntry {
113 binding: 0,
114 resource: globals.as_entire_binding(),
115 },
116 wgpu::BindGroupEntry {
117 binding: 1,
118 resource: wgpu::BindingResource::Sampler(&sampler),
119 },
120 ],
121 });
122 let page_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
123 label: Some("thumb page"),
124 entries: &[wgpu::BindGroupLayoutEntry {
125 binding: 0,
126 visibility: wgpu::ShaderStages::FRAGMENT,
127 ty: wgpu::BindingType::Texture {
128 sample_type: wgpu::TextureSampleType::Float { filterable: true },
129 view_dimension: wgpu::TextureViewDimension::D2,
130 multisampled: false,
131 },
132 count: None,
133 }],
134 });
135 let shader = device.create_shader_module(wgpu::include_wgsl!("shaders/thumbs.wgsl"));
136 let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
137 label: Some("thumbs"),
138 bind_group_layouts: &[Some(&globals_layout), Some(&page_layout)],
139 immediate_size: 0,
140 });
141 let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
142 label: Some("thumbs"),
143 layout: Some(&layout),
144 vertex: wgpu::VertexState {
145 module: &shader,
146 entry_point: Some("vs"),
147 compilation_options: Default::default(),
148 buffers: &[Some(wgpu::VertexBufferLayout {
149 array_stride: size_of::<Instance>() as u64,
150 step_mode: wgpu::VertexStepMode::Instance,
151 attributes: &ATTRS,
152 })],
153 },
154 primitive: wgpu::PrimitiveState::default(),
155 depth_stencil: None,
156 multisample: wgpu::MultisampleState::default(),
157 fragment: Some(wgpu::FragmentState {
158 module: &shader,
159 entry_point: Some("fs"),
160 compilation_options: Default::default(),
161 targets: &[Some(wgpu::ColorTargetState {
162 format: target,
163 blend: Some(wgpu::BlendState::ALPHA_BLENDING),
164 write_mask: wgpu::ColorWrites::ALL,
165 })],
166 }),
167 multiview_mask: None,
168 cache: None,
169 });
170 let capacity = 256;
171 Self {
172 pipeline,
173 globals,
174 globals_group,
175 page_layout,
176 instances: instance_buffer(device, capacity),
177 capacity,
178 batches: Vec::new(),
179 max_dim,
180 page_side: PAGE.min(max_dim),
181 slot: 0,
182 pages: Vec::new(),
183 slots: HashMap::new(),
184 free: Vec::new(),
185 frame: 0,
186 }
187 }
188
189 /// Start over for thumbnails of `cell` pixels (a new grid size or display scale).
190 pub fn reset(&mut self, cell: u32) {
191 let slot = (cell + 2).min(self.max_dim);
192 self.slot = slot;
193 self.page_side = PAGE.max(slot.next_power_of_two()).min(self.max_dim);
194 self.pages.clear();
195 self.slots.clear();
196 self.free.clear();
197 }
198
199 pub fn contains(&self, path: &Path) -> bool {
200 self.slots.contains_key(path)
201 }
202
203 fn per_row(&self) -> u32 {
204 (self.page_side / self.slot.max(1)).max(1)
205 }
206
207 fn add_page(&mut self, device: &wgpu::Device) {
208 let texture = device.create_texture(&wgpu::TextureDescriptor {
209 label: Some("thumb page"),
210 size: wgpu::Extent3d {
211 width: self.page_side,
212 height: self.page_side,
213 depth_or_array_layers: 1,
214 },
215 mip_level_count: 1,
216 sample_count: 1,
217 dimension: wgpu::TextureDimension::D2,
218 format: FORMAT,
219 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
220 view_formats: &[],
221 });
222 let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
223 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
224 label: Some("thumb page"),
225 layout: &self.page_layout,
226 entries: &[wgpu::BindGroupEntry {
227 binding: 0,
228 resource: wgpu::BindingResource::TextureView(&view),
229 }],
230 });
231 let page = self.pages.len();
232 self.pages.push(Page {
233 texture,
234 bind_group,
235 });
236 let n = self.per_row() * self.per_row();
237 self.free.extend((0..n).rev().map(|i| (page, i)));
238 log::debug!(
239 "thumbnail atlas: page {} ({n} slots of {} px)",
240 page + 1,
241 self.slot
242 );
243 }
244
245 /// Store `thumb` for `path`. When full, the least recently drawn thumbnail that isn't
246 /// `wanted` makes room; if every one is wanted, the atlas grows. Returns evicted paths.
247 pub fn insert(
248 &mut self,
249 device: &wgpu::Device,
250 queue: &wgpu::Queue,
251 path: &Path,
252 thumb: &Thumb,
253 wanted: impl Fn(&Path) -> bool,
254 ) -> Vec<PathBuf> {
255 let mut evicted = Vec::new();
256 let (page, index) = if let Some(s) = self.slots.remove(path) {
257 (s.page, s.index)
258 } else if let Some(free) = self.free.pop() {
259 free
260 } else {
261 let victim = self
262 .slots
263 .iter()
264 .filter(|(p, _)| !wanted(p))
265 .min_by_key(|(_, s)| s.used)
266 .map(|(p, _)| p.clone());
267 match victim.and_then(|v| self.slots.remove(&v).map(|s| (v, s))) {
268 Some((v, s)) => {
269 evicted.push(v);
270 (s.page, s.index)
271 }
272 None => {
273 self.add_page(device);
274 match self.free.pop() {
275 Some(f) => f,
276 None => return evicted,
277 }
278 }
279 }
280 };
281 // Clamp to the slot, keeping a 1 px clear border so filtering never picks up a
282 // neighbour or a previous occupant.
283 let inner = self.slot.saturating_sub(2).max(1);
284 let w = thumb.size[0].min(inner);
285 let h = thumb.size[1].min(inner);
286 let (pw, ph) = (w + 2, h + 2);
287 let mut padded = vec![0u8; (pw * ph * 4) as usize];
288 let src_stride = thumb.size[0] as usize * 4;
289 for y in 0..h as usize {
290 let src = &thumb.pixels[y * src_stride..][..w as usize * 4];
291 let dst = ((y + 1) * pw as usize + 1) * 4;
292 padded[dst..dst + src.len()].copy_from_slice(src);
293 }
294 let per_row = self.per_row();
295 let origin = [(index % per_row) * self.slot, (index / per_row) * self.slot];
296 queue.write_texture(
297 wgpu::TexelCopyTextureInfo {
298 texture: &self.pages[page].texture,
299 mip_level: 0,
300 origin: wgpu::Origin3d {
301 x: origin[0],
302 y: origin[1],
303 z: 0,
304 },
305 aspect: wgpu::TextureAspect::All,
306 },
307 &padded,
308 wgpu::TexelCopyBufferLayout {
309 offset: 0,
310 bytes_per_row: Some(pw * 4),
311 rows_per_image: Some(ph),
312 },
313 wgpu::Extent3d {
314 width: pw,
315 height: ph,
316 depth_or_array_layers: 1,
317 },
318 );
319 self.slots.insert(
320 path.to_path_buf(),
321 Slot {
322 page,
323 index,
324 size: [w, h],
325 used: self.frame,
326 },
327 );
328 evicted
329 }
330
331 /// Build this frame's instances. Thumbnails not in the atlas are skipped.
332 pub fn prepare(
333 &mut self,
334 device: &wgpu::Device,
335 queue: &wgpu::Queue,
336 draws: &[ThumbDraw],
337 viewport: [f32; 2],
338 ) {
339 self.frame += 1;
340 self.batches.clear();
341 if draws.is_empty() {
342 return;
343 }
344 queue.write_buffer(
345 &self.globals,
346 0,
347 bytemuck::bytes_of(&[viewport[0], viewport[1], 0.0, 0.0]),
348 );
349 let side = self.page_side as f32;
350 let per_row = self.per_row();
351 let mut list: Vec<(usize, Instance)> = Vec::with_capacity(draws.len());
352 for d in draws {
353 let Some(slot) = self.slots.get_mut(d.path) else {
354 continue;
355 };
356 slot.used = self.frame;
357 let x = (slot.index % per_row) * self.slot + 1;
358 let y = (slot.index / per_row) * self.slot + 1;
359 list.push((
360 slot.page,
361 Instance {
362 rect: d.rect,
363 uv: [
364 x as f32 / side,
365 y as f32 / side,
366 slot.size[0] as f32 / side,
367 slot.size[1] as f32 / side,
368 ],
369 params: [d.alpha, 0.0, 0.0, 0.0],
370 },
371 ));
372 }
373 list.sort_by_key(|(p, _)| *p);
374 if list.len() > self.capacity {
375 self.capacity = list.len().next_power_of_two();
376 self.instances = instance_buffer(device, self.capacity);
377 }
378 let instances: Vec<Instance> = list.iter().map(|(_, i)| *i).collect();
379 if !instances.is_empty() {
380 queue.write_buffer(&self.instances, 0, bytemuck::cast_slice(&instances));
381 }
382 let mut start = 0;
383 for (i, (page, _)) in list.iter().enumerate() {
384 let end = i + 1;
385 if end == list.len() || list[end].0 != *page {
386 self.batches.push((*page, start as u32..end as u32));
387 start = end;
388 }
389 }
390 }
391
392 pub fn draw(&self, pass: &mut wgpu::RenderPass<'_>) {
393 if self.batches.is_empty() {
394 return;
395 }
396 pass.set_pipeline(&self.pipeline);
397 pass.set_bind_group(0, &self.globals_group, &[]);
398 pass.set_vertex_buffer(0, self.instances.slice(..));
399 for (page, range) in &self.batches {
400 pass.set_bind_group(1, &self.pages[*page].bind_group, &[]);
401 pass.draw(0..6, range.clone());
402 }
403 }
404}