1use std::any::{Any, TypeId};
5
6use bytemuck::{Pod, Zeroable};
7use wgpu::util::DeviceExt;
8
9use crate::canvas::{Blit, CanvasRect, RenderTexture};
10use crate::texture::{TextureId, TextureRegistry};
11
12#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
14pub enum PostStage {
15 Canvas,
20 Window,
23}
24
25impl PostStage {
26 const fn index(self) -> usize {
29 match self {
30 Self::Canvas => 0,
31 Self::Window => 1,
32 }
33 }
34}
35
36#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
38pub struct PostEffectId(u64);
39
40pub trait PostEffect: 'static {
43 fn apply(&mut self, pass: &mut PostPass<'_>);
46
47 fn uses_history(&self) -> bool {
50 false
51 }
52}
53
54pub struct PostPass<'a> {
56 pub device: &'a wgpu::Device,
57 pub queue: &'a wgpu::Queue,
58 pub encoder: &'a mut wgpu::CommandEncoder,
59 pub input: &'a wgpu::TextureView,
61 pub output: &'a wgpu::TextureView,
63 pub history: Option<&'a wgpu::TextureView>,
66 pub indices: Option<&'a wgpu::TextureView>,
73 pub format: wgpu::TextureFormat,
74 pub stage: PostStage,
75 pub size: [u32; 2],
77 pub globals: PostGlobals,
78 pub(crate) textures: &'a TextureRegistry,
80}
81
82impl<'a> PostPass<'a> {
83 pub fn texture(&self, id: TextureId) -> Option<&'a wgpu::TextureView> {
89 self.textures.get(id).map(|t| &t.view)
90 }
91
92 pub fn texture_size(&self, id: TextureId) -> Option<[u32; 2]> {
94 self.textures.get(id).map(|t| t.size)
95 }
96}
97
98#[repr(C)]
101#[derive(Copy, Clone, Debug, Default, PartialEq, Pod, Zeroable)]
102pub struct PostGlobals {
103 pub size: [f32; 2],
105 pub canvas_size: [f32; 2],
107 pub window_size: [f32; 2],
109 pub canvas_origin: [f32; 2],
112 pub canvas_extent: [f32; 2],
113 pub time: f32,
115 pub frame: u32,
117 pub indexed: u32,
119 pub source: u32,
122 pub _pad: [u32; 2],
123}
124
125pub(crate) trait AnyPostEffect: PostEffect {
127 fn as_any_mut(&mut self) -> &mut dyn Any;
128}
129
130impl<T: PostEffect> AnyPostEffect for T {
131 fn as_any_mut(&mut self) -> &mut dyn Any {
132 self
133 }
134}
135
136struct Entry {
137 id: PostEffectId,
138 enabled: bool,
139 effect: Box<dyn AnyPostEffect>,
140}
141
142struct Targets {
144 format: wgpu::TextureFormat,
145 ping: [RenderTexture; 2],
146 history: Option<RenderTexture>,
147}
148
149#[derive(Default)]
150struct Stage {
151 effects: Vec<Entry>,
152 targets: Option<Targets>,
153}
154
155#[derive(Copy, Clone)]
157pub(crate) enum StageOutput {
158 Ping(usize),
159}
160
161impl Stage {
162 fn active(&self) -> bool {
163 self.effects.iter().any(|e| e.enabled)
164 }
165
166 fn wants_history(&self) -> bool {
167 self.effects
168 .iter()
169 .any(|e| e.enabled && e.effect.uses_history())
170 }
171
172 fn ensure_targets(
173 &mut self,
174 device: &wgpu::Device,
175 blit: &Blit,
176 format: wgpu::TextureFormat,
177 size: [u32; 2],
178 label: &str,
179 ) {
180 let history = self.wants_history();
181 let fits = self.targets.as_ref().is_some_and(|t| {
182 t.ping[0].size == size && t.format == format && t.history.is_some() == history
183 });
184 if fits {
185 return;
186 }
187 let make = |l: &str| RenderTexture::new(device, blit, format, size, l);
188 self.targets = Some(Targets {
189 format,
190 ping: [make(label), make(label)],
191 history: history.then(|| make(label)),
193 });
194 }
195}
196
197#[derive(Default)]
199pub(crate) struct PostChain {
200 stages: [Stage; 2],
201 next_id: u64,
202}
203
204pub(crate) struct StageRun<'a> {
205 pub device: &'a wgpu::Device,
206 pub queue: &'a wgpu::Queue,
207 pub blit: &'a Blit,
208 pub format: wgpu::TextureFormat,
209 pub globals: PostGlobals,
210 pub indices: Option<&'a wgpu::TextureView>,
212 pub textures: &'a TextureRegistry,
213}
214
215impl PostChain {
216 pub fn is_empty(&self) -> bool {
217 self.stages.iter().all(|s| s.effects.is_empty())
218 }
219
220 pub fn active(&self, stage: PostStage) -> bool {
221 self.stages[stage.index()].active()
222 }
223
224 pub fn add(&mut self, stage: PostStage, effect: Box<dyn AnyPostEffect>) -> PostEffectId {
225 self.next_id += 1;
226 let id = PostEffectId(self.next_id);
227 self.stages[stage.index()].effects.push(Entry {
228 id,
229 enabled: true,
230 effect,
231 });
232 id
233 }
234
235 fn find(&self, id: PostEffectId) -> Option<(usize, usize)> {
236 self.stages.iter().enumerate().find_map(|(s, stage)| {
237 stage
238 .effects
239 .iter()
240 .position(|e| e.id == id)
241 .map(|i| (s, i))
242 })
243 }
244
245 pub fn remove(&mut self, id: PostEffectId) -> bool {
246 match self.find(id) {
247 Some((s, i)) => {
248 self.stages[s].effects.remove(i);
249 true
250 }
251 None => false,
252 }
253 }
254
255 pub fn set_enabled(&mut self, id: PostEffectId, enabled: bool) -> bool {
256 match self.find(id) {
257 Some((s, i)) => {
258 self.stages[s].effects[i].enabled = enabled;
259 true
260 }
261 None => false,
262 }
263 }
264
265 pub fn is_enabled(&self, id: PostEffectId) -> Option<bool> {
266 self.find(id)
267 .map(|(s, i)| self.stages[s].effects[i].enabled)
268 }
269
270 pub fn get_mut(&mut self, id: PostEffectId) -> Option<&mut dyn Any> {
271 let (s, i) = self.find(id)?;
272 Some(self.stages[s].effects[i].effect.as_any_mut())
273 }
274
275 pub fn run(
278 &mut self,
279 stage: PostStage,
280 run: &StageRun<'_>,
281 encoder: &mut wgpu::CommandEncoder,
282 input: &wgpu::TextureView,
283 size: [u32; 2],
284 ) -> Option<StageOutput> {
285 let st = &mut self.stages[stage.index()];
286 if !st.active() {
287 return None;
288 }
289 let label = match stage {
290 PostStage::Canvas => "post.canvas",
291 PostStage::Window => "post.window",
292 };
293 st.ensure_targets(run.device, run.blit, run.format, size, label);
294 let targets = st.targets.as_ref()?;
295 let mut globals = run.globals;
296 globals.size = [size[0] as f32, size[1] as f32];
297 globals.indexed = u32::from(run.indices.is_some());
298 let mut current: Option<usize> = None;
299 for entry in st.effects.iter_mut().filter(|e| e.enabled) {
300 let out = current.map_or(0, |c| 1 - c);
301 let input = current.map_or(input, |c| &targets.ping[c].view);
302 let mut pass = PostPass {
303 device: run.device,
304 queue: run.queue,
305 encoder,
306 input,
307 output: &targets.ping[out].view,
308 history: targets.history.as_ref().map(|h| &h.view),
309 indices: run.indices,
310 textures: run.textures,
311 format: run.format,
312 stage,
313 size,
314 globals,
315 };
316 entry.effect.apply(&mut pass);
317 current = Some(out);
318 }
319 let out = current?;
320 if let Some(history) = &targets.history {
321 encoder.copy_texture_to_texture(
322 targets.ping[out].texture.as_image_copy(),
323 history.texture.as_image_copy(),
324 history.texture.size(),
325 );
326 }
327 Some(StageOutput::Ping(out))
328 }
329
330 pub fn output(&self, stage: PostStage, out: StageOutput) -> Option<&RenderTexture> {
332 let targets = self.stages[stage.index()].targets.as_ref()?;
333 match out {
334 StageOutput::Ping(i) => targets.ping.get(i),
335 }
336 }
337}
338
339pub const SHADER_PRELUDE: &str = concat!(
343 include_str!("post_prelude.wgsl"),
344 "\n",
345 include_str!("bayer.wgsl"),
346 "\n",
347 include_str!("index.wgsl")
348);
349
350pub const BAYER_WGSL: &str = include_str!("bayer.wgsl");
355
356pub struct ShaderEffect {
392 pipeline: wgpu::RenderPipeline,
393 layout: wgpu::BindGroupLayout,
394 globals: wgpu::Buffer,
395 samplers: [wgpu::Sampler; 2],
396 params: Option<Params>,
397 history: bool,
398 source: Option<TextureId>,
399 source_warned: bool,
402}
403
404struct Params {
405 buffer: wgpu::Buffer,
406 bind_group: wgpu::BindGroup,
407 bytes: Vec<u8>,
408 type_id: TypeId,
410 dirty: bool,
411}
412
413impl ShaderEffect {
414 pub fn new(gfx: &crate::Graphics, label: &str, wgsl: &str) -> Self {
417 Self::build(gfx.device(), gfx.surface_format(), label, wgsl, None)
418 }
419
420 pub fn with_params<T: Pod>(gfx: &crate::Graphics, label: &str, wgsl: &str, params: &T) -> Self {
424 Self::build(
425 gfx.device(),
426 gfx.surface_format(),
427 label,
428 wgsl,
429 Some((bytemuck::bytes_of(params), TypeId::of::<T>())),
430 )
431 }
432
433 pub(crate) fn build(
434 device: &wgpu::Device,
435 format: wgpu::TextureFormat,
436 label: &str,
437 wgsl: &str,
438 params: Option<(&[u8], TypeId)>,
439 ) -> Self {
440 let source = format!("{SHADER_PRELUDE}\n{wgsl}");
441 let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
442 label: Some(label),
443 source: wgpu::ShaderSource::Wgsl(source.into()),
444 });
445 let tex = |binding| wgpu::BindGroupLayoutEntry {
446 binding,
447 visibility: wgpu::ShaderStages::FRAGMENT,
448 ty: wgpu::BindingType::Texture {
449 sample_type: wgpu::TextureSampleType::Float { filterable: true },
450 view_dimension: wgpu::TextureViewDimension::D2,
451 multisampled: false,
452 },
453 count: None,
454 };
455 let sampler = |binding| wgpu::BindGroupLayoutEntry {
456 binding,
457 visibility: wgpu::ShaderStages::FRAGMENT,
458 ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
459 count: None,
460 };
461 let uniform = |binding| wgpu::BindGroupLayoutEntry {
462 binding,
463 visibility: wgpu::ShaderStages::FRAGMENT,
464 ty: wgpu::BindingType::Buffer {
465 ty: wgpu::BufferBindingType::Uniform,
466 has_dynamic_offset: false,
467 min_binding_size: None,
468 },
469 count: None,
470 };
471 let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
472 label: Some("post.bgl"),
473 entries: &[
474 tex(0),
475 sampler(1),
476 sampler(2),
477 uniform(3),
478 tex(4),
479 tex(5),
480 tex(6),
481 ],
482 });
483 let params = params.map(|(bytes, type_id)| {
484 let params_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
485 label: Some("post.params.bgl"),
486 entries: &[uniform(0)],
487 });
488 let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
489 label: Some("post.params"),
490 contents: bytes,
491 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
492 });
493 let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
494 label: Some("post.params.bg"),
495 layout: ¶ms_layout,
496 entries: &[wgpu::BindGroupEntry {
497 binding: 0,
498 resource: buffer.as_entire_binding(),
499 }],
500 });
501 (
502 params_layout,
503 Params {
504 buffer,
505 bind_group,
506 bytes: bytes.to_vec(),
507 type_id,
508 dirty: false,
509 },
510 )
511 });
512 let mut layouts = vec![Some(&layout)];
513 if let Some((params_layout, _)) = ¶ms {
514 layouts.push(Some(params_layout));
515 }
516 let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
517 label: Some(label),
518 bind_group_layouts: &layouts,
519 immediate_size: 0,
520 });
521 let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
522 label: Some(label),
523 layout: Some(&pipeline_layout),
524 vertex: wgpu::VertexState {
525 module: &module,
526 entry_point: Some("post_vs_main"),
527 compilation_options: Default::default(),
528 buffers: &[],
529 },
530 fragment: Some(wgpu::FragmentState {
531 module: &module,
532 entry_point: Some("fs_main"),
533 compilation_options: Default::default(),
534 targets: &[Some(wgpu::ColorTargetState {
535 format,
536 blend: None,
537 write_mask: wgpu::ColorWrites::ALL,
538 })],
539 }),
540 primitive: wgpu::PrimitiveState::default(),
541 depth_stencil: None,
542 multisample: wgpu::MultisampleState::default(),
543 multiview_mask: None,
544 cache: None,
545 });
546 let globals = device.create_buffer(&wgpu::BufferDescriptor {
547 label: Some("post.globals"),
548 size: size_of::<PostGlobals>() as u64,
549 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
550 mapped_at_creation: false,
551 });
552 let make_sampler = |filter| {
553 device.create_sampler(&wgpu::SamplerDescriptor {
554 label: Some("post.sampler"),
555 mag_filter: filter,
556 min_filter: filter,
557 ..Default::default()
558 })
559 };
560 Self {
561 pipeline,
562 layout,
563 globals,
564 samplers: [
565 make_sampler(wgpu::FilterMode::Nearest),
566 make_sampler(wgpu::FilterMode::Linear),
567 ],
568 params: params.map(|(_, p)| p),
569 history: false,
570 source: None,
571 source_warned: false,
572 }
573 }
574
575 #[must_use]
589 pub fn with_source(mut self, source: TextureId) -> Self {
590 self.source = Some(source);
591 self.source_warned = false;
592 self
593 }
594
595 pub fn set_source(&mut self, source: Option<TextureId>) {
597 self.source = source;
598 self.source_warned = false;
599 }
600
601 pub fn source(&self) -> Option<TextureId> {
603 self.source
604 }
605
606 #[must_use]
608 pub fn with_history(mut self) -> Self {
609 self.history = true;
610 self
611 }
612
613 pub fn set_params<T: Pod>(&mut self, params: &T) {
616 let bytes = bytemuck::bytes_of(params);
617 match &mut self.params {
618 Some(p) if p.type_id == TypeId::of::<T>() => {
619 p.bytes.copy_from_slice(bytes);
620 p.dirty = true;
621 }
622 _ => log::warn!(
623 "ShaderEffect::set_params: effect has no params of type {}",
624 std::any::type_name::<T>()
625 ),
626 }
627 }
628}
629
630impl PostEffect for ShaderEffect {
631 fn apply(&mut self, pass: &mut PostPass<'_>) {
632 let source = self.source.and_then(|id| pass.texture(id));
633 match (self.source, source) {
634 (Some(id), None) if !self.source_warned => {
635 log::warn!("ShaderEffect: unknown source texture {id:?}; using the input");
636 self.source_warned = true;
637 }
638 (_, Some(_)) => self.source_warned = false,
639 _ => {}
640 }
641 let globals = PostGlobals {
642 source: u32::from(source.is_some()),
643 ..pass.globals
644 };
645 pass.queue
646 .write_buffer(&self.globals, 0, bytemuck::bytes_of(&globals));
647 if let Some(p) = &mut self.params
648 && p.dirty
649 {
650 pass.queue.write_buffer(&p.buffer, 0, &p.bytes);
651 p.dirty = false;
652 }
653 let history = pass.history.filter(|_| self.history).unwrap_or(pass.input);
654 let indices = pass.indices.unwrap_or(pass.input);
655 let bind_group = pass.device.create_bind_group(&wgpu::BindGroupDescriptor {
656 label: Some("post.bg"),
657 layout: &self.layout,
658 entries: &[
659 wgpu::BindGroupEntry {
660 binding: 0,
661 resource: wgpu::BindingResource::TextureView(pass.input),
662 },
663 wgpu::BindGroupEntry {
664 binding: 1,
665 resource: wgpu::BindingResource::Sampler(&self.samplers[0]),
666 },
667 wgpu::BindGroupEntry {
668 binding: 2,
669 resource: wgpu::BindingResource::Sampler(&self.samplers[1]),
670 },
671 wgpu::BindGroupEntry {
672 binding: 3,
673 resource: self.globals.as_entire_binding(),
674 },
675 wgpu::BindGroupEntry {
676 binding: 4,
677 resource: wgpu::BindingResource::TextureView(history),
678 },
679 wgpu::BindGroupEntry {
680 binding: 5,
681 resource: wgpu::BindingResource::TextureView(indices),
682 },
683 wgpu::BindGroupEntry {
684 binding: 6,
685 resource: wgpu::BindingResource::TextureView(source.unwrap_or(pass.input)),
686 },
687 ],
688 });
689 let mut rp = pass.encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
690 label: Some("post.pass"),
691 color_attachments: &[Some(wgpu::RenderPassColorAttachment {
692 view: pass.output,
693 depth_slice: None,
694 resolve_target: None,
695 ops: wgpu::Operations {
696 load: wgpu::LoadOp::Clear(wgpu::Color::BLACK),
697 store: wgpu::StoreOp::Store,
698 },
699 })],
700 depth_stencil_attachment: None,
701 occlusion_query_set: None,
702 timestamp_writes: None,
703 multiview_mask: None,
704 });
705 rp.set_pipeline(&self.pipeline);
706 rp.set_bind_group(0, &bind_group, &[]);
707 if let Some(p) = &self.params {
708 rp.set_bind_group(1, &p.bind_group, &[]);
709 }
710 rp.draw(0..3, 0..1);
711 }
712
713 fn uses_history(&self) -> bool {
714 self.history
715 }
716}
717
718pub(crate) fn globals(
720 canvas_size: [u32; 2],
721 window_size: [u32; 2],
722 rect: CanvasRect,
723 time: f32,
724 frame: u32,
725) -> PostGlobals {
726 PostGlobals {
727 size: [0.0; 2],
728 canvas_size: [canvas_size[0] as f32, canvas_size[1] as f32],
729 window_size: [window_size[0] as f32, window_size[1] as f32],
730 canvas_origin: rect.origin,
731 canvas_extent: rect.size,
732 time,
733 frame,
734 indexed: 0,
735 source: 0,
736 _pad: [0; 2],
737 }
738}
739
740#[cfg(test)]
741mod tests {
742 use super::*;
743
744 #[test]
745 fn globals_are_uniform_sized() {
746 assert_eq!(size_of::<PostGlobals>() % 16, 0);
748 }
749}