bevy_pbr/ssao/
mod.rs

1use crate::NodePbr;
2use bevy_app::{App, Plugin};
3use bevy_asset::{load_internal_asset, weak_handle, Handle};
4use bevy_core_pipeline::{
5    core_3d::graph::{Core3d, Node3d},
6    prelude::Camera3d,
7    prepass::{DepthPrepass, NormalPrepass, ViewPrepassTextures},
8};
9use bevy_ecs::{
10    prelude::{Component, Entity},
11    query::{Has, QueryItem, With},
12    reflect::ReflectComponent,
13    resource::Resource,
14    schedule::IntoScheduleConfigs,
15    system::{Commands, Query, Res, ResMut},
16    world::{FromWorld, World},
17};
18use bevy_reflect::{std_traits::ReflectDefault, Reflect};
19use bevy_render::{
20    camera::{ExtractedCamera, TemporalJitter},
21    extract_component::ExtractComponent,
22    globals::{GlobalsBuffer, GlobalsUniform},
23    prelude::Camera,
24    render_graph::{NodeRunError, RenderGraphApp, RenderGraphContext, ViewNode, ViewNodeRunner},
25    render_resource::{
26        binding_types::{
27            sampler, texture_2d, texture_depth_2d, texture_storage_2d, uniform_buffer,
28        },
29        *,
30    },
31    renderer::{RenderAdapter, RenderContext, RenderDevice, RenderQueue},
32    sync_component::SyncComponentPlugin,
33    sync_world::RenderEntity,
34    texture::{CachedTexture, TextureCache},
35    view::{Msaa, ViewUniform, ViewUniformOffset, ViewUniforms},
36    Extract, ExtractSchedule, Render, RenderApp, RenderSet,
37};
38use bevy_utils::prelude::default;
39use core::mem;
40use tracing::{error, warn};
41
42const PREPROCESS_DEPTH_SHADER_HANDLE: Handle<Shader> =
43    weak_handle!("b7f2cc3d-c935-4f5c-9ae2-43d6b0d5659a");
44const SSAO_SHADER_HANDLE: Handle<Shader> = weak_handle!("9ea355d7-37a2-4cc4-b4d1-5d8ab47b07f5");
45const SPATIAL_DENOISE_SHADER_HANDLE: Handle<Shader> =
46    weak_handle!("0f2764a0-b343-471b-b7ce-ef5d636f4fc3");
47const SSAO_UTILS_SHADER_HANDLE: Handle<Shader> =
48    weak_handle!("da53c78d-f318-473e-bdff-b388bc50ada2");
49
50/// Plugin for screen space ambient occlusion.
51pub struct ScreenSpaceAmbientOcclusionPlugin;
52
53impl Plugin for ScreenSpaceAmbientOcclusionPlugin {
54    fn build(&self, app: &mut App) {
55        load_internal_asset!(
56            app,
57            PREPROCESS_DEPTH_SHADER_HANDLE,
58            "preprocess_depth.wgsl",
59            Shader::from_wgsl
60        );
61        load_internal_asset!(app, SSAO_SHADER_HANDLE, "ssao.wgsl", Shader::from_wgsl);
62        load_internal_asset!(
63            app,
64            SPATIAL_DENOISE_SHADER_HANDLE,
65            "spatial_denoise.wgsl",
66            Shader::from_wgsl
67        );
68        load_internal_asset!(
69            app,
70            SSAO_UTILS_SHADER_HANDLE,
71            "ssao_utils.wgsl",
72            Shader::from_wgsl
73        );
74
75        app.register_type::<ScreenSpaceAmbientOcclusion>();
76
77        app.add_plugins(SyncComponentPlugin::<ScreenSpaceAmbientOcclusion>::default());
78    }
79
80    fn finish(&self, app: &mut App) {
81        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
82            return;
83        };
84
85        if !render_app
86            .world()
87            .resource::<RenderAdapter>()
88            .get_texture_format_features(TextureFormat::R16Float)
89            .allowed_usages
90            .contains(TextureUsages::STORAGE_BINDING)
91        {
92            warn!("ScreenSpaceAmbientOcclusionPlugin not loaded. GPU lacks support: TextureFormat::R16Float does not support TextureUsages::STORAGE_BINDING.");
93            return;
94        }
95
96        if render_app
97            .world()
98            .resource::<RenderDevice>()
99            .limits()
100            .max_storage_textures_per_shader_stage
101            < 5
102        {
103            warn!("ScreenSpaceAmbientOcclusionPlugin not loaded. GPU lacks support: Limits::max_storage_textures_per_shader_stage is less than 5.");
104            return;
105        }
106
107        render_app
108            .init_resource::<SsaoPipelines>()
109            .init_resource::<SpecializedComputePipelines<SsaoPipelines>>()
110            .add_systems(ExtractSchedule, extract_ssao_settings)
111            .add_systems(
112                Render,
113                (
114                    prepare_ssao_pipelines.in_set(RenderSet::Prepare),
115                    prepare_ssao_textures.in_set(RenderSet::PrepareResources),
116                    prepare_ssao_bind_groups.in_set(RenderSet::PrepareBindGroups),
117                ),
118            )
119            .add_render_graph_node::<ViewNodeRunner<SsaoNode>>(
120                Core3d,
121                NodePbr::ScreenSpaceAmbientOcclusion,
122            )
123            .add_render_graph_edges(
124                Core3d,
125                (
126                    // END_PRE_PASSES -> SCREEN_SPACE_AMBIENT_OCCLUSION -> MAIN_PASS
127                    Node3d::EndPrepasses,
128                    NodePbr::ScreenSpaceAmbientOcclusion,
129                    Node3d::StartMainPass,
130                ),
131            );
132    }
133}
134
135/// Component to apply screen space ambient occlusion to a 3d camera.
136///
137/// Screen space ambient occlusion (SSAO) approximates small-scale,
138/// local occlusion of _indirect_ diffuse light between objects, based on what's visible on-screen.
139/// SSAO does not apply to direct lighting, such as point or directional lights.
140///
141/// This darkens creases, e.g. on staircases, and gives nice contact shadows
142/// where objects meet, giving entities a more "grounded" feel.
143///
144/// # Usage Notes
145///
146/// Requires that you add [`ScreenSpaceAmbientOcclusionPlugin`] to your app.
147///
148/// It strongly recommended that you use SSAO in conjunction with
149/// TAA ([`bevy_core_pipeline::experimental::taa::TemporalAntiAliasing`]).
150/// Doing so greatly reduces SSAO noise.
151///
152/// SSAO is not supported on `WebGL2`, and is not currently supported on `WebGPU`.
153#[derive(Component, ExtractComponent, Reflect, PartialEq, Clone, Debug)]
154#[reflect(Component, Debug, Default, PartialEq, Clone)]
155#[require(DepthPrepass, NormalPrepass)]
156#[doc(alias = "Ssao")]
157pub struct ScreenSpaceAmbientOcclusion {
158    /// Quality of the SSAO effect.
159    pub quality_level: ScreenSpaceAmbientOcclusionQualityLevel,
160    /// A constant estimated thickness of objects.
161    ///
162    /// This value is used to decide how far behind an object a ray of light needs to be in order
163    /// to pass behind it. Any ray closer than that will be occluded.
164    pub constant_object_thickness: f32,
165}
166
167impl Default for ScreenSpaceAmbientOcclusion {
168    fn default() -> Self {
169        Self {
170            quality_level: ScreenSpaceAmbientOcclusionQualityLevel::default(),
171            constant_object_thickness: 0.25,
172        }
173    }
174}
175
176#[derive(Reflect, PartialEq, Eq, Hash, Clone, Copy, Default, Debug)]
177#[reflect(PartialEq, Hash, Clone, Default)]
178pub enum ScreenSpaceAmbientOcclusionQualityLevel {
179    Low,
180    Medium,
181    #[default]
182    High,
183    Ultra,
184    Custom {
185        /// Higher slice count means less noise, but worse performance.
186        slice_count: u32,
187        /// Samples per slice side is also tweakable, but recommended to be left at 2 or 3.
188        samples_per_slice_side: u32,
189    },
190}
191
192impl ScreenSpaceAmbientOcclusionQualityLevel {
193    fn sample_counts(&self) -> (u32, u32) {
194        match self {
195            Self::Low => (1, 2),    // 4 spp (1 * (2 * 2)), plus optional temporal samples
196            Self::Medium => (2, 2), // 8 spp (2 * (2 * 2)), plus optional temporal samples
197            Self::High => (3, 3),   // 18 spp (3 * (3 * 2)), plus optional temporal samples
198            Self::Ultra => (9, 3),  // 54 spp (9 * (3 * 2)), plus optional temporal samples
199            Self::Custom {
200                slice_count: slices,
201                samples_per_slice_side,
202            } => (*slices, *samples_per_slice_side),
203        }
204    }
205}
206
207#[derive(Default)]
208struct SsaoNode {}
209
210impl ViewNode for SsaoNode {
211    type ViewQuery = (
212        &'static ExtractedCamera,
213        &'static SsaoPipelineId,
214        &'static SsaoBindGroups,
215        &'static ViewUniformOffset,
216    );
217
218    fn run(
219        &self,
220        _graph: &mut RenderGraphContext,
221        render_context: &mut RenderContext,
222        (camera, pipeline_id, bind_groups, view_uniform_offset): QueryItem<Self::ViewQuery>,
223        world: &World,
224    ) -> Result<(), NodeRunError> {
225        let pipelines = world.resource::<SsaoPipelines>();
226        let pipeline_cache = world.resource::<PipelineCache>();
227        let (
228            Some(camera_size),
229            Some(preprocess_depth_pipeline),
230            Some(spatial_denoise_pipeline),
231            Some(ssao_pipeline),
232        ) = (
233            camera.physical_viewport_size,
234            pipeline_cache.get_compute_pipeline(pipelines.preprocess_depth_pipeline),
235            pipeline_cache.get_compute_pipeline(pipelines.spatial_denoise_pipeline),
236            pipeline_cache.get_compute_pipeline(pipeline_id.0),
237        )
238        else {
239            return Ok(());
240        };
241
242        render_context.command_encoder().push_debug_group("ssao");
243
244        {
245            let mut preprocess_depth_pass =
246                render_context
247                    .command_encoder()
248                    .begin_compute_pass(&ComputePassDescriptor {
249                        label: Some("ssao_preprocess_depth_pass"),
250                        timestamp_writes: None,
251                    });
252            preprocess_depth_pass.set_pipeline(preprocess_depth_pipeline);
253            preprocess_depth_pass.set_bind_group(0, &bind_groups.preprocess_depth_bind_group, &[]);
254            preprocess_depth_pass.set_bind_group(
255                1,
256                &bind_groups.common_bind_group,
257                &[view_uniform_offset.offset],
258            );
259            preprocess_depth_pass.dispatch_workgroups(
260                camera_size.x.div_ceil(16),
261                camera_size.y.div_ceil(16),
262                1,
263            );
264        }
265
266        {
267            let mut ssao_pass =
268                render_context
269                    .command_encoder()
270                    .begin_compute_pass(&ComputePassDescriptor {
271                        label: Some("ssao_ssao_pass"),
272                        timestamp_writes: None,
273                    });
274            ssao_pass.set_pipeline(ssao_pipeline);
275            ssao_pass.set_bind_group(0, &bind_groups.ssao_bind_group, &[]);
276            ssao_pass.set_bind_group(
277                1,
278                &bind_groups.common_bind_group,
279                &[view_uniform_offset.offset],
280            );
281            ssao_pass.dispatch_workgroups(camera_size.x.div_ceil(8), camera_size.y.div_ceil(8), 1);
282        }
283
284        {
285            let mut spatial_denoise_pass =
286                render_context
287                    .command_encoder()
288                    .begin_compute_pass(&ComputePassDescriptor {
289                        label: Some("ssao_spatial_denoise_pass"),
290                        timestamp_writes: None,
291                    });
292            spatial_denoise_pass.set_pipeline(spatial_denoise_pipeline);
293            spatial_denoise_pass.set_bind_group(0, &bind_groups.spatial_denoise_bind_group, &[]);
294            spatial_denoise_pass.set_bind_group(
295                1,
296                &bind_groups.common_bind_group,
297                &[view_uniform_offset.offset],
298            );
299            spatial_denoise_pass.dispatch_workgroups(
300                camera_size.x.div_ceil(8),
301                camera_size.y.div_ceil(8),
302                1,
303            );
304        }
305
306        render_context.command_encoder().pop_debug_group();
307        Ok(())
308    }
309}
310
311#[derive(Resource)]
312struct SsaoPipelines {
313    preprocess_depth_pipeline: CachedComputePipelineId,
314    spatial_denoise_pipeline: CachedComputePipelineId,
315
316    common_bind_group_layout: BindGroupLayout,
317    preprocess_depth_bind_group_layout: BindGroupLayout,
318    ssao_bind_group_layout: BindGroupLayout,
319    spatial_denoise_bind_group_layout: BindGroupLayout,
320
321    hilbert_index_lut: TextureView,
322    point_clamp_sampler: Sampler,
323    linear_clamp_sampler: Sampler,
324}
325
326impl FromWorld for SsaoPipelines {
327    fn from_world(world: &mut World) -> Self {
328        let render_device = world.resource::<RenderDevice>();
329        let render_queue = world.resource::<RenderQueue>();
330        let pipeline_cache = world.resource::<PipelineCache>();
331
332        let hilbert_index_lut = render_device
333            .create_texture_with_data(
334                render_queue,
335                &(TextureDescriptor {
336                    label: Some("ssao_hilbert_index_lut"),
337                    size: Extent3d {
338                        width: HILBERT_WIDTH as u32,
339                        height: HILBERT_WIDTH as u32,
340                        depth_or_array_layers: 1,
341                    },
342                    mip_level_count: 1,
343                    sample_count: 1,
344                    dimension: TextureDimension::D2,
345                    format: TextureFormat::R16Uint,
346                    usage: TextureUsages::TEXTURE_BINDING,
347                    view_formats: &[],
348                }),
349                TextureDataOrder::default(),
350                bytemuck::cast_slice(&generate_hilbert_index_lut()),
351            )
352            .create_view(&TextureViewDescriptor::default());
353
354        let point_clamp_sampler = render_device.create_sampler(&SamplerDescriptor {
355            min_filter: FilterMode::Nearest,
356            mag_filter: FilterMode::Nearest,
357            mipmap_filter: FilterMode::Nearest,
358            address_mode_u: AddressMode::ClampToEdge,
359            address_mode_v: AddressMode::ClampToEdge,
360            ..Default::default()
361        });
362        let linear_clamp_sampler = render_device.create_sampler(&SamplerDescriptor {
363            min_filter: FilterMode::Linear,
364            mag_filter: FilterMode::Linear,
365            mipmap_filter: FilterMode::Nearest,
366            address_mode_u: AddressMode::ClampToEdge,
367            address_mode_v: AddressMode::ClampToEdge,
368            ..Default::default()
369        });
370
371        let common_bind_group_layout = render_device.create_bind_group_layout(
372            "ssao_common_bind_group_layout",
373            &BindGroupLayoutEntries::sequential(
374                ShaderStages::COMPUTE,
375                (
376                    sampler(SamplerBindingType::NonFiltering),
377                    sampler(SamplerBindingType::Filtering),
378                    uniform_buffer::<ViewUniform>(true),
379                ),
380            ),
381        );
382
383        let preprocess_depth_bind_group_layout = render_device.create_bind_group_layout(
384            "ssao_preprocess_depth_bind_group_layout",
385            &BindGroupLayoutEntries::sequential(
386                ShaderStages::COMPUTE,
387                (
388                    texture_depth_2d(),
389                    texture_storage_2d(TextureFormat::R16Float, StorageTextureAccess::WriteOnly),
390                    texture_storage_2d(TextureFormat::R16Float, StorageTextureAccess::WriteOnly),
391                    texture_storage_2d(TextureFormat::R16Float, StorageTextureAccess::WriteOnly),
392                    texture_storage_2d(TextureFormat::R16Float, StorageTextureAccess::WriteOnly),
393                    texture_storage_2d(TextureFormat::R16Float, StorageTextureAccess::WriteOnly),
394                ),
395            ),
396        );
397
398        let ssao_bind_group_layout = render_device.create_bind_group_layout(
399            "ssao_ssao_bind_group_layout",
400            &BindGroupLayoutEntries::sequential(
401                ShaderStages::COMPUTE,
402                (
403                    texture_2d(TextureSampleType::Float { filterable: true }),
404                    texture_2d(TextureSampleType::Float { filterable: false }),
405                    texture_2d(TextureSampleType::Uint),
406                    texture_storage_2d(TextureFormat::R16Float, StorageTextureAccess::WriteOnly),
407                    texture_storage_2d(TextureFormat::R32Uint, StorageTextureAccess::WriteOnly),
408                    uniform_buffer::<GlobalsUniform>(false),
409                    uniform_buffer::<f32>(false),
410                ),
411            ),
412        );
413
414        let spatial_denoise_bind_group_layout = render_device.create_bind_group_layout(
415            "ssao_spatial_denoise_bind_group_layout",
416            &BindGroupLayoutEntries::sequential(
417                ShaderStages::COMPUTE,
418                (
419                    texture_2d(TextureSampleType::Float { filterable: false }),
420                    texture_2d(TextureSampleType::Uint),
421                    texture_storage_2d(TextureFormat::R16Float, StorageTextureAccess::WriteOnly),
422                ),
423            ),
424        );
425
426        let preprocess_depth_pipeline =
427            pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
428                label: Some("ssao_preprocess_depth_pipeline".into()),
429                layout: vec![
430                    preprocess_depth_bind_group_layout.clone(),
431                    common_bind_group_layout.clone(),
432                ],
433                push_constant_ranges: vec![],
434                shader: PREPROCESS_DEPTH_SHADER_HANDLE,
435                shader_defs: Vec::new(),
436                entry_point: "preprocess_depth".into(),
437                zero_initialize_workgroup_memory: false,
438            });
439
440        let spatial_denoise_pipeline =
441            pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
442                label: Some("ssao_spatial_denoise_pipeline".into()),
443                layout: vec![
444                    spatial_denoise_bind_group_layout.clone(),
445                    common_bind_group_layout.clone(),
446                ],
447                push_constant_ranges: vec![],
448                shader: SPATIAL_DENOISE_SHADER_HANDLE,
449                shader_defs: Vec::new(),
450                entry_point: "spatial_denoise".into(),
451                zero_initialize_workgroup_memory: false,
452            });
453
454        Self {
455            preprocess_depth_pipeline,
456            spatial_denoise_pipeline,
457
458            common_bind_group_layout,
459            preprocess_depth_bind_group_layout,
460            ssao_bind_group_layout,
461            spatial_denoise_bind_group_layout,
462
463            hilbert_index_lut,
464            point_clamp_sampler,
465            linear_clamp_sampler,
466        }
467    }
468}
469
470#[derive(PartialEq, Eq, Hash, Clone)]
471struct SsaoPipelineKey {
472    quality_level: ScreenSpaceAmbientOcclusionQualityLevel,
473    temporal_jitter: bool,
474}
475
476impl SpecializedComputePipeline for SsaoPipelines {
477    type Key = SsaoPipelineKey;
478
479    fn specialize(&self, key: Self::Key) -> ComputePipelineDescriptor {
480        let (slice_count, samples_per_slice_side) = key.quality_level.sample_counts();
481
482        let mut shader_defs = vec![
483            ShaderDefVal::Int("SLICE_COUNT".to_string(), slice_count as i32),
484            ShaderDefVal::Int(
485                "SAMPLES_PER_SLICE_SIDE".to_string(),
486                samples_per_slice_side as i32,
487            ),
488        ];
489
490        if key.temporal_jitter {
491            shader_defs.push("TEMPORAL_JITTER".into());
492        }
493
494        ComputePipelineDescriptor {
495            label: Some("ssao_ssao_pipeline".into()),
496            layout: vec![
497                self.ssao_bind_group_layout.clone(),
498                self.common_bind_group_layout.clone(),
499            ],
500            push_constant_ranges: vec![],
501            shader: SSAO_SHADER_HANDLE,
502            shader_defs,
503            entry_point: "ssao".into(),
504            zero_initialize_workgroup_memory: false,
505        }
506    }
507}
508
509fn extract_ssao_settings(
510    mut commands: Commands,
511    cameras: Extract<
512        Query<
513            (RenderEntity, &Camera, &ScreenSpaceAmbientOcclusion, &Msaa),
514            (With<Camera3d>, With<DepthPrepass>, With<NormalPrepass>),
515        >,
516    >,
517) {
518    for (entity, camera, ssao_settings, msaa) in &cameras {
519        if *msaa != Msaa::Off {
520            error!(
521                "SSAO is being used which requires Msaa::Off, but Msaa is currently set to Msaa::{:?}",
522                *msaa
523            );
524            return;
525        }
526        let mut entity_commands = commands
527            .get_entity(entity)
528            .expect("SSAO entity wasn't synced.");
529        if camera.is_active {
530            entity_commands.insert(ssao_settings.clone());
531        } else {
532            entity_commands.remove::<ScreenSpaceAmbientOcclusion>();
533        }
534    }
535}
536
537#[derive(Component)]
538pub struct ScreenSpaceAmbientOcclusionResources {
539    preprocessed_depth_texture: CachedTexture,
540    ssao_noisy_texture: CachedTexture, // Pre-spatially denoised texture
541    pub screen_space_ambient_occlusion_texture: CachedTexture, // Spatially denoised texture
542    depth_differences_texture: CachedTexture,
543    thickness_buffer: Buffer,
544}
545
546fn prepare_ssao_textures(
547    mut commands: Commands,
548    mut texture_cache: ResMut<TextureCache>,
549    render_device: Res<RenderDevice>,
550    views: Query<(Entity, &ExtractedCamera, &ScreenSpaceAmbientOcclusion)>,
551) {
552    for (entity, camera, ssao_settings) in &views {
553        let Some(physical_viewport_size) = camera.physical_viewport_size else {
554            continue;
555        };
556        let size = Extent3d {
557            width: physical_viewport_size.x,
558            height: physical_viewport_size.y,
559            depth_or_array_layers: 1,
560        };
561
562        let preprocessed_depth_texture = texture_cache.get(
563            &render_device,
564            TextureDescriptor {
565                label: Some("ssao_preprocessed_depth_texture"),
566                size,
567                mip_level_count: 5,
568                sample_count: 1,
569                dimension: TextureDimension::D2,
570                format: TextureFormat::R16Float,
571                usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
572                view_formats: &[],
573            },
574        );
575
576        let ssao_noisy_texture = texture_cache.get(
577            &render_device,
578            TextureDescriptor {
579                label: Some("ssao_noisy_texture"),
580                size,
581                mip_level_count: 1,
582                sample_count: 1,
583                dimension: TextureDimension::D2,
584                format: TextureFormat::R16Float,
585                usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
586                view_formats: &[],
587            },
588        );
589
590        let ssao_texture = texture_cache.get(
591            &render_device,
592            TextureDescriptor {
593                label: Some("ssao_texture"),
594                size,
595                mip_level_count: 1,
596                sample_count: 1,
597                dimension: TextureDimension::D2,
598                format: TextureFormat::R16Float,
599                usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
600                view_formats: &[],
601            },
602        );
603
604        let depth_differences_texture = texture_cache.get(
605            &render_device,
606            TextureDescriptor {
607                label: Some("ssao_depth_differences_texture"),
608                size,
609                mip_level_count: 1,
610                sample_count: 1,
611                dimension: TextureDimension::D2,
612                format: TextureFormat::R32Uint,
613                usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
614                view_formats: &[],
615            },
616        );
617
618        let thickness_buffer = render_device.create_buffer_with_data(&BufferInitDescriptor {
619            label: Some("thickness_buffer"),
620            contents: &ssao_settings.constant_object_thickness.to_le_bytes(),
621            usage: BufferUsages::UNIFORM,
622        });
623
624        commands
625            .entity(entity)
626            .insert(ScreenSpaceAmbientOcclusionResources {
627                preprocessed_depth_texture,
628                ssao_noisy_texture,
629                screen_space_ambient_occlusion_texture: ssao_texture,
630                depth_differences_texture,
631                thickness_buffer,
632            });
633    }
634}
635
636#[derive(Component)]
637struct SsaoPipelineId(CachedComputePipelineId);
638
639fn prepare_ssao_pipelines(
640    mut commands: Commands,
641    pipeline_cache: Res<PipelineCache>,
642    mut pipelines: ResMut<SpecializedComputePipelines<SsaoPipelines>>,
643    pipeline: Res<SsaoPipelines>,
644    views: Query<(Entity, &ScreenSpaceAmbientOcclusion, Has<TemporalJitter>)>,
645) {
646    for (entity, ssao_settings, temporal_jitter) in &views {
647        let pipeline_id = pipelines.specialize(
648            &pipeline_cache,
649            &pipeline,
650            SsaoPipelineKey {
651                quality_level: ssao_settings.quality_level,
652                temporal_jitter,
653            },
654        );
655
656        commands.entity(entity).insert(SsaoPipelineId(pipeline_id));
657    }
658}
659
660#[derive(Component)]
661struct SsaoBindGroups {
662    common_bind_group: BindGroup,
663    preprocess_depth_bind_group: BindGroup,
664    ssao_bind_group: BindGroup,
665    spatial_denoise_bind_group: BindGroup,
666}
667
668fn prepare_ssao_bind_groups(
669    mut commands: Commands,
670    render_device: Res<RenderDevice>,
671    pipelines: Res<SsaoPipelines>,
672    view_uniforms: Res<ViewUniforms>,
673    global_uniforms: Res<GlobalsBuffer>,
674    views: Query<(
675        Entity,
676        &ScreenSpaceAmbientOcclusionResources,
677        &ViewPrepassTextures,
678    )>,
679) {
680    let (Some(view_uniforms), Some(globals_uniforms)) = (
681        view_uniforms.uniforms.binding(),
682        global_uniforms.buffer.binding(),
683    ) else {
684        return;
685    };
686
687    for (entity, ssao_resources, prepass_textures) in &views {
688        let common_bind_group = render_device.create_bind_group(
689            "ssao_common_bind_group",
690            &pipelines.common_bind_group_layout,
691            &BindGroupEntries::sequential((
692                &pipelines.point_clamp_sampler,
693                &pipelines.linear_clamp_sampler,
694                view_uniforms.clone(),
695            )),
696        );
697
698        let create_depth_view = |mip_level| {
699            ssao_resources
700                .preprocessed_depth_texture
701                .texture
702                .create_view(&TextureViewDescriptor {
703                    label: Some("ssao_preprocessed_depth_texture_mip_view"),
704                    base_mip_level: mip_level,
705                    format: Some(TextureFormat::R16Float),
706                    dimension: Some(TextureViewDimension::D2),
707                    mip_level_count: Some(1),
708                    ..default()
709                })
710        };
711
712        let preprocess_depth_bind_group = render_device.create_bind_group(
713            "ssao_preprocess_depth_bind_group",
714            &pipelines.preprocess_depth_bind_group_layout,
715            &BindGroupEntries::sequential((
716                prepass_textures.depth_view().unwrap(),
717                &create_depth_view(0),
718                &create_depth_view(1),
719                &create_depth_view(2),
720                &create_depth_view(3),
721                &create_depth_view(4),
722            )),
723        );
724
725        let ssao_bind_group = render_device.create_bind_group(
726            "ssao_ssao_bind_group",
727            &pipelines.ssao_bind_group_layout,
728            &BindGroupEntries::sequential((
729                &ssao_resources.preprocessed_depth_texture.default_view,
730                prepass_textures.normal_view().unwrap(),
731                &pipelines.hilbert_index_lut,
732                &ssao_resources.ssao_noisy_texture.default_view,
733                &ssao_resources.depth_differences_texture.default_view,
734                globals_uniforms.clone(),
735                ssao_resources.thickness_buffer.as_entire_binding(),
736            )),
737        );
738
739        let spatial_denoise_bind_group = render_device.create_bind_group(
740            "ssao_spatial_denoise_bind_group",
741            &pipelines.spatial_denoise_bind_group_layout,
742            &BindGroupEntries::sequential((
743                &ssao_resources.ssao_noisy_texture.default_view,
744                &ssao_resources.depth_differences_texture.default_view,
745                &ssao_resources
746                    .screen_space_ambient_occlusion_texture
747                    .default_view,
748            )),
749        );
750
751        commands.entity(entity).insert(SsaoBindGroups {
752            common_bind_group,
753            preprocess_depth_bind_group,
754            ssao_bind_group,
755            spatial_denoise_bind_group,
756        });
757    }
758}
759
760fn generate_hilbert_index_lut() -> [[u16; 64]; 64] {
761    use core::array::from_fn;
762    from_fn(|x| from_fn(|y| hilbert_index(x as u16, y as u16)))
763}
764
765// https://www.shadertoy.com/view/3tB3z3
766const HILBERT_WIDTH: u16 = 64;
767fn hilbert_index(mut x: u16, mut y: u16) -> u16 {
768    let mut index = 0;
769
770    let mut level: u16 = HILBERT_WIDTH / 2;
771    while level > 0 {
772        let region_x = (x & level > 0) as u16;
773        let region_y = (y & level > 0) as u16;
774        index += level * level * ((3 * region_x) ^ region_y);
775
776        if region_y == 0 {
777            if region_x == 1 {
778                x = HILBERT_WIDTH - 1 - x;
779                y = HILBERT_WIDTH - 1 - y;
780            }
781
782            mem::swap(&mut x, &mut y);
783        }
784
785        level /= 2;
786    }
787
788    index
789}