Skip to main content

max / shop

11.8 KB · 355 lines History Blame Raw
1 //! RGBA image pipeline: one wgpu texture per image, one draw call per image.
2 //!
3 //! MVP scope: a fixed pool of image slots, each a `wgpu::Texture` and a
4 //! precomputed bind group. `set_image` uploads or replaces a slot's texture;
5 //! `draw_frame` draws all live slots at their placement rectangles. Instance
6 //! batching within a single pipeline is not worth it while we're expecting
7 //! ~1 image on screen at a time (yazi preview pane).
8
9 use std::collections::HashMap;
10
11 use bytemuck::{Pod, Zeroable};
12 use wgpu::util::DeviceExt;
13
14 /// Placement of an image on the screen, in physical pixels.
15 #[derive(Debug, Clone, Copy)]
16 pub struct ImagePlacement {
17 pub image_id: u32,
18 pub x: f32,
19 pub y: f32,
20 pub w: f32,
21 pub h: f32,
22 }
23
24 #[repr(C)]
25 #[derive(Clone, Copy, Pod, Zeroable)]
26 struct ImageUniforms {
27 screen: [f32; 2],
28 _pad: [f32; 2],
29 }
30
31 const IMAGE_SHADER: &str = r"
32 struct Uniforms {
33 screen: vec2<f32>,
34 _pad: vec2<f32>,
35 };
36
37 @group(0) @binding(0) var<uniform> u: Uniforms;
38 @group(0) @binding(1) var tex: texture_2d<f32>;
39 @group(0) @binding(2) var samp: sampler;
40
41 struct VOut {
42 @builtin(position) pos: vec4<f32>,
43 @location(0) uv: vec2<f32>,
44 };
45
46 @vertex
47 fn vs(
48 @location(0) corner: vec2<f32>,
49 @location(1) inst_pos: vec2<f32>,
50 @location(2) inst_size: vec2<f32>,
51 ) -> VOut {
52 let px = inst_pos + corner * inst_size;
53 let ndc = vec2<f32>(
54 (px.x / u.screen.x) * 2.0 - 1.0,
55 1.0 - (px.y / u.screen.y) * 2.0,
56 );
57 var out: VOut;
58 out.pos = vec4<f32>(ndc, 0.0, 1.0);
59 out.uv = corner;
60 return out;
61 }
62
63 @fragment
64 fn fs(in: VOut) -> @location(0) vec4<f32> {
65 return textureSample(tex, samp, in.uv);
66 }
67 ";
68
69 #[repr(C)]
70 #[derive(Clone, Copy, Pod, Zeroable)]
71 struct ImgInstance {
72 pos: [f32; 2],
73 size: [f32; 2],
74 }
75
76 const QUAD: [[f32; 2]; 4] = [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]];
77 const INDICES: [u16; 6] = [0, 1, 2, 2, 1, 3];
78
79 struct ImageSlot {
80 _texture: wgpu::Texture,
81 bind_group: wgpu::BindGroup,
82 }
83
84 pub struct ImageRenderer {
85 pipeline: wgpu::RenderPipeline,
86 bind_group_layout: wgpu::BindGroupLayout,
87 sampler: wgpu::Sampler,
88 uniforms_buf: wgpu::Buffer,
89 quad_buf: wgpu::Buffer,
90 index_buf: wgpu::Buffer,
91 instance_buf: wgpu::Buffer,
92 images: HashMap<u32, ImageSlot>,
93 }
94
95 impl ImageRenderer {
96 pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat) -> Self {
97 let uniforms_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
98 label: Some("shop.image.uniforms"),
99 contents: bytemuck::bytes_of(&ImageUniforms {
100 screen: [1.0, 1.0],
101 _pad: [0.0, 0.0],
102 }),
103 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
104 });
105 let quad_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
106 label: Some("shop.image.quad"),
107 contents: bytemuck::cast_slice(&QUAD),
108 usage: wgpu::BufferUsages::VERTEX,
109 });
110 let index_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
111 label: Some("shop.image.index"),
112 contents: bytemuck::cast_slice(&INDICES),
113 usage: wgpu::BufferUsages::INDEX,
114 });
115 let instance_buf = device.create_buffer(&wgpu::BufferDescriptor {
116 label: Some("shop.image.instance"),
117 size: std::mem::size_of::<ImgInstance>() as u64,
118 usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
119 mapped_at_creation: false,
120 });
121
122 let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
123 label: Some("shop.image.sampler"),
124 mag_filter: wgpu::FilterMode::Linear,
125 min_filter: wgpu::FilterMode::Linear,
126 ..Default::default()
127 });
128
129 let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
130 label: Some("shop.image.bgl"),
131 entries: &[
132 wgpu::BindGroupLayoutEntry {
133 binding: 0,
134 visibility: wgpu::ShaderStages::VERTEX,
135 ty: wgpu::BindingType::Buffer {
136 ty: wgpu::BufferBindingType::Uniform,
137 has_dynamic_offset: false,
138 min_binding_size: None,
139 },
140 count: None,
141 },
142 wgpu::BindGroupLayoutEntry {
143 binding: 1,
144 visibility: wgpu::ShaderStages::FRAGMENT,
145 ty: wgpu::BindingType::Texture {
146 sample_type: wgpu::TextureSampleType::Float { filterable: true },
147 view_dimension: wgpu::TextureViewDimension::D2,
148 multisampled: false,
149 },
150 count: None,
151 },
152 wgpu::BindGroupLayoutEntry {
153 binding: 2,
154 visibility: wgpu::ShaderStages::FRAGMENT,
155 ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
156 count: None,
157 },
158 ],
159 });
160
161 let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
162 label: Some("shop.image.shader"),
163 source: wgpu::ShaderSource::Wgsl(IMAGE_SHADER.into()),
164 });
165 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
166 label: Some("shop.image.pl"),
167 bind_group_layouts: &[Some(&bind_group_layout)],
168 immediate_size: 0,
169 });
170 let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
171 label: Some("shop.image.pipeline"),
172 layout: Some(&pipeline_layout),
173 vertex: wgpu::VertexState {
174 module: &module,
175 entry_point: Some("vs"),
176 compilation_options: wgpu::PipelineCompilationOptions::default(),
177 buffers: &[
178 Some(wgpu::VertexBufferLayout {
179 array_stride: 8,
180 step_mode: wgpu::VertexStepMode::Vertex,
181 attributes: &wgpu::vertex_attr_array![0 => Float32x2],
182 }),
183 Some(wgpu::VertexBufferLayout {
184 array_stride: std::mem::size_of::<ImgInstance>() as u64,
185 step_mode: wgpu::VertexStepMode::Instance,
186 attributes: &wgpu::vertex_attr_array![
187 1 => Float32x2,
188 2 => Float32x2,
189 ],
190 }),
191 ],
192 },
193 fragment: Some(wgpu::FragmentState {
194 module: &module,
195 entry_point: Some("fs"),
196 compilation_options: wgpu::PipelineCompilationOptions::default(),
197 targets: &[Some(wgpu::ColorTargetState {
198 format,
199 blend: Some(wgpu::BlendState::ALPHA_BLENDING),
200 write_mask: wgpu::ColorWrites::ALL,
201 })],
202 }),
203 primitive: wgpu::PrimitiveState::default(),
204 depth_stencil: None,
205 multisample: wgpu::MultisampleState::default(),
206 multiview_mask: None,
207 cache: None,
208 });
209
210 Self {
211 pipeline,
212 bind_group_layout,
213 sampler,
214 uniforms_buf,
215 quad_buf,
216 index_buf,
217 instance_buf,
218 images: HashMap::new(),
219 }
220 }
221
222 pub fn resize(&self, queue: &wgpu::Queue, width: u32, height: u32) {
223 let u = ImageUniforms {
224 screen: [width as f32, height as f32],
225 _pad: [0.0, 0.0],
226 };
227 queue.write_buffer(&self.uniforms_buf, 0, bytemuck::bytes_of(&u));
228 }
229
230 /// Upload an RGBA image and register it under `id`. Replaces any existing
231 /// image at the same id.
232 pub fn set_image(
233 &mut self,
234 device: &wgpu::Device,
235 queue: &wgpu::Queue,
236 id: u32,
237 rgba: &[u8],
238 width: u32,
239 height: u32,
240 ) {
241 let texture = device.create_texture(&wgpu::TextureDescriptor {
242 label: Some("shop.image.tex"),
243 size: wgpu::Extent3d {
244 width,
245 height,
246 depth_or_array_layers: 1,
247 },
248 mip_level_count: 1,
249 sample_count: 1,
250 dimension: wgpu::TextureDimension::D2,
251 format: wgpu::TextureFormat::Rgba8UnormSrgb,
252 usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
253 view_formats: &[],
254 });
255 queue.write_texture(
256 wgpu::TexelCopyTextureInfo {
257 texture: &texture,
258 mip_level: 0,
259 origin: wgpu::Origin3d::ZERO,
260 aspect: wgpu::TextureAspect::All,
261 },
262 rgba,
263 wgpu::TexelCopyBufferLayout {
264 offset: 0,
265 bytes_per_row: Some(width * 4),
266 rows_per_image: Some(height),
267 },
268 wgpu::Extent3d {
269 width,
270 height,
271 depth_or_array_layers: 1,
272 },
273 );
274 let view = texture.create_view(&wgpu::TextureViewDescriptor::default());
275 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
276 label: Some("shop.image.bg"),
277 layout: &self.bind_group_layout,
278 entries: &[
279 wgpu::BindGroupEntry {
280 binding: 0,
281 resource: self.uniforms_buf.as_entire_binding(),
282 },
283 wgpu::BindGroupEntry {
284 binding: 1,
285 resource: wgpu::BindingResource::TextureView(&view),
286 },
287 wgpu::BindGroupEntry {
288 binding: 2,
289 resource: wgpu::BindingResource::Sampler(&self.sampler),
290 },
291 ],
292 });
293 self.images.insert(
294 id,
295 ImageSlot {
296 _texture: texture,
297 bind_group,
298 },
299 );
300 }
301
302 pub fn drop_image(&mut self, id: u32) {
303 self.images.remove(&id);
304 }
305
306 pub fn drop_all(&mut self) {
307 self.images.clear();
308 }
309
310 /// Draw each placement in order (one texture bind + draw call per image).
311 pub fn draw(
312 &self,
313 queue: &wgpu::Queue,
314 view: &wgpu::TextureView,
315 encoder: &mut wgpu::CommandEncoder,
316 placements: &[ImagePlacement],
317 ) {
318 if placements.is_empty() {
319 return;
320 }
321 let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
322 label: Some("shop.image.pass"),
323 color_attachments: &[Some(wgpu::RenderPassColorAttachment {
324 view,
325 resolve_target: None,
326 ops: wgpu::Operations {
327 load: wgpu::LoadOp::Load,
328 store: wgpu::StoreOp::Store,
329 },
330 depth_slice: None,
331 })],
332 ..Default::default()
333 });
334 pass.set_pipeline(&self.pipeline);
335 pass.set_vertex_buffer(0, self.quad_buf.slice(..));
336 pass.set_index_buffer(self.index_buf.slice(..), wgpu::IndexFormat::Uint16);
337 for placement in placements {
338 let Some(slot) = self.images.get(&placement.image_id) else {
339 continue;
340 };
341 queue.write_buffer(
342 &self.instance_buf,
343 0,
344 bytemuck::bytes_of(&ImgInstance {
345 pos: [placement.x, placement.y],
346 size: [placement.w, placement.h],
347 }),
348 );
349 pass.set_vertex_buffer(1, self.instance_buf.slice(..));
350 pass.set_bind_group(0, &slot.bind_group, &[]);
351 pass.draw_indexed(0..6, 0, 0..1);
352 }
353 }
354 }
355