bevy_pbr/light_probe/
generate.rs

1//! Like [`EnvironmentMapLight`], but filtered in realtime from a cubemap.
2//!
3//! An environment map needs to be processed to be able to support uses beyond a simple skybox,
4//! such as reflections, and ambient light contribution.
5//! This process is called filtering, and can either be done ahead of time (prefiltering), or
6//! in realtime, although at a reduced quality. Prefiltering is preferred, but not always possible:
7//! sometimes you only gain access to an environment map at runtime, for whatever reason.
8//! Typically this is from realtime reflection probes, but can also be from other sources.
9//!
10//! In any case, Bevy supports both modes of filtering.
11//! This module provides realtime filtering via [`bevy_light::GeneratedEnvironmentMapLight`].
12//! For prefiltered environment maps, see [`bevy_light::EnvironmentMapLight`].
13//! These components are intended to be added to a camera.
14use bevy_app::{App, Plugin, Update};
15use bevy_asset::{embedded_asset, load_embedded_asset, AssetServer, Assets, RenderAssetUsages};
16use bevy_core_pipeline::core_3d::graph::{Core3d, Node3d};
17use bevy_ecs::{
18    component::Component,
19    entity::Entity,
20    query::{QueryState, With, Without},
21    resource::Resource,
22    schedule::IntoScheduleConfigs,
23    system::{lifetimeless::Read, Commands, Query, Res, ResMut},
24    world::{FromWorld, World},
25};
26use bevy_image::Image;
27use bevy_math::{Quat, UVec2, Vec2};
28use bevy_render::{
29    diagnostic::RecordDiagnostics,
30    render_asset::RenderAssets,
31    render_graph::{Node, NodeRunError, RenderGraphContext, RenderGraphExt, RenderLabel},
32    render_resource::{
33        binding_types::*, AddressMode, BindGroup, BindGroupEntries, BindGroupLayout,
34        BindGroupLayoutEntries, CachedComputePipelineId, ComputePassDescriptor,
35        ComputePipelineDescriptor, DownlevelFlags, Extent3d, FilterMode, PipelineCache, Sampler,
36        SamplerBindingType, SamplerDescriptor, ShaderStages, ShaderType, StorageTextureAccess,
37        Texture, TextureAspect, TextureDescriptor, TextureDimension, TextureFormat,
38        TextureFormatFeatureFlags, TextureSampleType, TextureUsages, TextureView,
39        TextureViewDescriptor, TextureViewDimension, UniformBuffer,
40    },
41    renderer::{RenderAdapter, RenderContext, RenderDevice, RenderQueue},
42    settings::WgpuFeatures,
43    sync_component::SyncComponentPlugin,
44    sync_world::RenderEntity,
45    texture::{CachedTexture, GpuImage, TextureCache},
46    Extract, ExtractSchedule, Render, RenderApp, RenderStartup, RenderSystems,
47};
48
49// Implementation: generate diffuse and specular cubemaps required by PBR
50// from a given high-res cubemap by
51//
52// 1. Copying the base mip (level 0) of the source cubemap into an intermediate
53//    storage texture.
54// 2. Generating mipmaps using [single-pass down-sampling] (SPD).
55// 3. Convolving the mip chain twice:
56//    * a [Lambertian convolution] for the 32 × 32 diffuse cubemap
57//    * a [GGX convolution], once per mip level, for the specular cubemap.
58//
59// [single-pass down-sampling]: https://gpuopen.com/fidelityfx-spd/
60// [Lambertian convolution]: https://bruop.github.io/ibl/#:~:text=Lambertian%20Diffuse%20Component
61// [GGX convolution]: https://gpuopen.com/download/Bounded_VNDF_Sampling_for_Smith-GGX_Reflections.pdf
62
63use bevy_light::{EnvironmentMapLight, GeneratedEnvironmentMapLight};
64use bevy_shader::ShaderDefVal;
65use core::cmp::min;
66use tracing::info;
67
68use crate::Bluenoise;
69
70/// Labels for the environment map generation nodes
71#[derive(PartialEq, Eq, Debug, Copy, Clone, Hash, RenderLabel)]
72pub enum GeneratorNode {
73    Downsampling,
74    Filtering,
75}
76
77/// Stores the bind group layouts for the environment map generation pipelines
78#[derive(Resource)]
79pub struct GeneratorBindGroupLayouts {
80    pub downsampling_first: BindGroupLayout,
81    pub downsampling_second: BindGroupLayout,
82    pub radiance: BindGroupLayout,
83    pub irradiance: BindGroupLayout,
84    pub copy: BindGroupLayout,
85}
86
87/// Samplers for the environment map generation pipelines
88#[derive(Resource)]
89pub struct GeneratorSamplers {
90    pub linear: Sampler,
91}
92
93/// Pipelines for the environment map generation pipelines
94#[derive(Resource)]
95pub struct GeneratorPipelines {
96    pub downsample_first: CachedComputePipelineId,
97    pub downsample_second: CachedComputePipelineId,
98    pub copy: CachedComputePipelineId,
99    pub radiance: CachedComputePipelineId,
100    pub irradiance: CachedComputePipelineId,
101}
102
103/// Configuration for downsampling strategy based on device limits
104#[derive(Resource, Clone, Copy, Debug, PartialEq, Eq)]
105pub struct DownsamplingConfig {
106    // can bind ≥12 storage textures and use read-write storage textures
107    pub combine_bind_group: bool,
108}
109
110pub struct EnvironmentMapGenerationPlugin;
111
112impl Plugin for EnvironmentMapGenerationPlugin {
113    fn build(&self, _: &mut App) {}
114    fn finish(&self, app: &mut App) {
115        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
116            let adapter = render_app.world().resource::<RenderAdapter>();
117            let device = render_app.world().resource::<RenderDevice>();
118
119            // Cubemap SPD requires at least 6 storage textures
120            let limit_support = device.limits().max_storage_textures_per_shader_stage >= 6
121                && device.limits().max_compute_workgroup_storage_size != 0
122                && device.limits().max_compute_workgroup_size_x != 0;
123
124            let downlevel_support = adapter
125                .get_downlevel_capabilities()
126                .flags
127                .contains(DownlevelFlags::COMPUTE_SHADERS);
128
129            if !limit_support || !downlevel_support {
130                info!("Disabling EnvironmentMapGenerationPlugin because compute is not supported on this platform. This is safe to ignore if you are not using EnvironmentMapGenerationPlugin.");
131                return;
132            }
133        } else {
134            return;
135        }
136
137        embedded_asset!(app, "environment_filter.wgsl");
138        embedded_asset!(app, "downsample.wgsl");
139        embedded_asset!(app, "copy.wgsl");
140
141        app.add_plugins(SyncComponentPlugin::<GeneratedEnvironmentMapLight>::default())
142            .add_systems(Update, generate_environment_map_light);
143
144        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
145            return;
146        };
147
148        render_app
149            .add_render_graph_node::<DownsamplingNode>(Core3d, GeneratorNode::Downsampling)
150            .add_render_graph_node::<FilteringNode>(Core3d, GeneratorNode::Filtering)
151            .add_render_graph_edges(
152                Core3d,
153                (
154                    Node3d::EndPrepasses,
155                    GeneratorNode::Downsampling,
156                    GeneratorNode::Filtering,
157                    Node3d::StartMainPass,
158                ),
159            )
160            .add_systems(
161                ExtractSchedule,
162                extract_generated_environment_map_entities.after(generate_environment_map_light),
163            )
164            .add_systems(
165                Render,
166                prepare_generated_environment_map_bind_groups
167                    .in_set(RenderSystems::PrepareBindGroups),
168            )
169            .add_systems(
170                Render,
171                prepare_generated_environment_map_intermediate_textures
172                    .in_set(RenderSystems::PrepareResources),
173            )
174            .add_systems(
175                RenderStartup,
176                initialize_generated_environment_map_resources,
177            );
178    }
179}
180
181// The number of storage textures required to combine the bind group
182const REQUIRED_STORAGE_TEXTURES: u32 = 12;
183
184/// Initializes all render-world resources used by the environment-map generator once on
185/// [`bevy_render::RenderStartup`].
186pub fn initialize_generated_environment_map_resources(
187    mut commands: Commands,
188    render_device: Res<RenderDevice>,
189    render_adapter: Res<RenderAdapter>,
190    pipeline_cache: Res<PipelineCache>,
191    asset_server: Res<AssetServer>,
192) {
193    // Determine whether we can use a single, large bind group for all mip outputs
194    let storage_texture_limit = render_device.limits().max_storage_textures_per_shader_stage;
195
196    // Determine whether we can read and write to the same rgba16f storage texture
197    let read_write_support = render_adapter
198        .get_texture_format_features(TextureFormat::Rgba16Float)
199        .flags
200        .contains(TextureFormatFeatureFlags::STORAGE_READ_WRITE);
201
202    // Combine the bind group and use read-write storage if it is supported
203    let combine_bind_group =
204        storage_texture_limit >= REQUIRED_STORAGE_TEXTURES && read_write_support;
205
206    // Output mips are write-only
207    let mips =
208        texture_storage_2d_array(TextureFormat::Rgba16Float, StorageTextureAccess::WriteOnly);
209
210    // Bind group layouts
211    let (downsampling_first, downsampling_second) = if combine_bind_group {
212        // One big bind group layout containing all outputs 1–12
213        let downsampling = render_device.create_bind_group_layout(
214            "downsampling_bind_group_layout_combined",
215            &BindGroupLayoutEntries::sequential(
216                ShaderStages::COMPUTE,
217                (
218                    sampler(SamplerBindingType::Filtering),
219                    uniform_buffer::<DownsamplingConstants>(false),
220                    texture_2d_array(TextureSampleType::Float { filterable: true }),
221                    mips, // 1
222                    mips, // 2
223                    mips, // 3
224                    mips, // 4
225                    mips, // 5
226                    texture_storage_2d_array(
227                        TextureFormat::Rgba16Float,
228                        StorageTextureAccess::ReadWrite,
229                    ), // 6
230                    mips, // 7
231                    mips, // 8
232                    mips, // 9
233                    mips, // 10
234                    mips, // 11
235                    mips, // 12
236                ),
237            ),
238        );
239
240        (downsampling.clone(), downsampling)
241    } else {
242        // Split layout: first pass outputs 1–6, second pass outputs 7–12 (input mip6 read-only)
243
244        let downsampling_first = render_device.create_bind_group_layout(
245            "downsampling_first_bind_group_layout",
246            &BindGroupLayoutEntries::sequential(
247                ShaderStages::COMPUTE,
248                (
249                    sampler(SamplerBindingType::Filtering),
250                    uniform_buffer::<DownsamplingConstants>(false),
251                    // Input mip 0
252                    texture_2d_array(TextureSampleType::Float { filterable: true }),
253                    mips, // 1
254                    mips, // 2
255                    mips, // 3
256                    mips, // 4
257                    mips, // 5
258                    mips, // 6
259                ),
260            ),
261        );
262
263        let downsampling_second = render_device.create_bind_group_layout(
264            "downsampling_second_bind_group_layout",
265            &BindGroupLayoutEntries::sequential(
266                ShaderStages::COMPUTE,
267                (
268                    sampler(SamplerBindingType::Filtering),
269                    uniform_buffer::<DownsamplingConstants>(false),
270                    // Input mip 6
271                    texture_2d_array(TextureSampleType::Float { filterable: true }),
272                    mips, // 7
273                    mips, // 8
274                    mips, // 9
275                    mips, // 10
276                    mips, // 11
277                    mips, // 12
278                ),
279            ),
280        );
281
282        (downsampling_first, downsampling_second)
283    };
284    let radiance = render_device.create_bind_group_layout(
285        "radiance_bind_group_layout",
286        &BindGroupLayoutEntries::sequential(
287            ShaderStages::COMPUTE,
288            (
289                // Source environment cubemap
290                texture_2d_array(TextureSampleType::Float { filterable: true }),
291                sampler(SamplerBindingType::Filtering), // Source sampler
292                // Output specular map
293                texture_storage_2d_array(
294                    TextureFormat::Rgba16Float,
295                    StorageTextureAccess::WriteOnly,
296                ),
297                uniform_buffer::<FilteringConstants>(false), // Uniforms
298                texture_2d_array(TextureSampleType::Float { filterable: true }), // Blue noise texture
299            ),
300        ),
301    );
302
303    let irradiance = render_device.create_bind_group_layout(
304        "irradiance_bind_group_layout",
305        &BindGroupLayoutEntries::sequential(
306            ShaderStages::COMPUTE,
307            (
308                // Source environment cubemap
309                texture_2d_array(TextureSampleType::Float { filterable: true }),
310                sampler(SamplerBindingType::Filtering), // Source sampler
311                // Output irradiance map
312                texture_storage_2d_array(
313                    TextureFormat::Rgba16Float,
314                    StorageTextureAccess::WriteOnly,
315                ),
316                uniform_buffer::<FilteringConstants>(false), // Uniforms
317                texture_2d_array(TextureSampleType::Float { filterable: true }), // Blue noise texture
318            ),
319        ),
320    );
321
322    let copy = render_device.create_bind_group_layout(
323        "copy_bind_group_layout",
324        &BindGroupLayoutEntries::sequential(
325            ShaderStages::COMPUTE,
326            (
327                // Source cubemap
328                texture_2d_array(TextureSampleType::Float { filterable: true }),
329                // Destination mip0
330                texture_storage_2d_array(
331                    TextureFormat::Rgba16Float,
332                    StorageTextureAccess::WriteOnly,
333                ),
334            ),
335        ),
336    );
337
338    let layouts = GeneratorBindGroupLayouts {
339        downsampling_first,
340        downsampling_second,
341        radiance,
342        irradiance,
343        copy,
344    };
345
346    // Samplers
347    let linear = render_device.create_sampler(&SamplerDescriptor {
348        label: Some("generator_linear_sampler"),
349        address_mode_u: AddressMode::ClampToEdge,
350        address_mode_v: AddressMode::ClampToEdge,
351        address_mode_w: AddressMode::ClampToEdge,
352        mag_filter: FilterMode::Linear,
353        min_filter: FilterMode::Linear,
354        mipmap_filter: FilterMode::Linear,
355        ..Default::default()
356    });
357
358    let samplers = GeneratorSamplers { linear };
359
360    // Pipelines
361    let features = render_device.features();
362    let mut shader_defs = vec![];
363    if features.contains(WgpuFeatures::SUBGROUP) {
364        shader_defs.push(ShaderDefVal::Int("SUBGROUP_SUPPORT".into(), 1));
365    }
366    if combine_bind_group {
367        shader_defs.push(ShaderDefVal::Int("COMBINE_BIND_GROUP".into(), 1));
368    }
369    #[cfg(feature = "bluenoise_texture")]
370    {
371        shader_defs.push(ShaderDefVal::Int("HAS_BLUE_NOISE".into(), 1));
372    }
373
374    let downsampling_shader = load_embedded_asset!(asset_server.as_ref(), "downsample.wgsl");
375    let env_filter_shader = load_embedded_asset!(asset_server.as_ref(), "environment_filter.wgsl");
376    let copy_shader = load_embedded_asset!(asset_server.as_ref(), "copy.wgsl");
377
378    // First pass for base mip Levels (0-5)
379    let downsample_first = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
380        label: Some("downsampling_first_pipeline".into()),
381        layout: vec![layouts.downsampling_first.clone()],
382        push_constant_ranges: vec![],
383        shader: downsampling_shader.clone(),
384        shader_defs: {
385            let mut defs = shader_defs.clone();
386            if !combine_bind_group {
387                defs.push(ShaderDefVal::Int("FIRST_PASS".into(), 1));
388            }
389            defs
390        },
391        entry_point: Some("downsample_first".into()),
392        zero_initialize_workgroup_memory: false,
393    });
394
395    let downsample_second = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
396        label: Some("downsampling_second_pipeline".into()),
397        layout: vec![layouts.downsampling_second.clone()],
398        push_constant_ranges: vec![],
399        shader: downsampling_shader,
400        shader_defs: {
401            let mut defs = shader_defs.clone();
402            if !combine_bind_group {
403                defs.push(ShaderDefVal::Int("SECOND_PASS".into(), 1));
404            }
405            defs
406        },
407        entry_point: Some("downsample_second".into()),
408        zero_initialize_workgroup_memory: false,
409    });
410
411    // Radiance map for specular environment maps
412    let radiance = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
413        label: Some("radiance_pipeline".into()),
414        layout: vec![layouts.radiance.clone()],
415        push_constant_ranges: vec![],
416        shader: env_filter_shader.clone(),
417        shader_defs: shader_defs.clone(),
418        entry_point: Some("generate_radiance_map".into()),
419        zero_initialize_workgroup_memory: false,
420    });
421
422    // Irradiance map for diffuse environment maps
423    let irradiance = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
424        label: Some("irradiance_pipeline".into()),
425        layout: vec![layouts.irradiance.clone()],
426        push_constant_ranges: vec![],
427        shader: env_filter_shader,
428        shader_defs: shader_defs.clone(),
429        entry_point: Some("generate_irradiance_map".into()),
430        zero_initialize_workgroup_memory: false,
431    });
432
433    // Copy pipeline handles format conversion and populates mip0 when formats differ
434    let copy_pipeline = pipeline_cache.queue_compute_pipeline(ComputePipelineDescriptor {
435        label: Some("copy_pipeline".into()),
436        layout: vec![layouts.copy.clone()],
437        push_constant_ranges: vec![],
438        shader: copy_shader,
439        shader_defs: vec![],
440        entry_point: Some("copy".into()),
441        zero_initialize_workgroup_memory: false,
442    });
443
444    let pipelines = GeneratorPipelines {
445        downsample_first,
446        downsample_second,
447        radiance,
448        irradiance,
449        copy: copy_pipeline,
450    };
451
452    // Insert all resources into the render world
453    commands.insert_resource(layouts);
454    commands.insert_resource(samplers);
455    commands.insert_resource(pipelines);
456    commands.insert_resource(DownsamplingConfig { combine_bind_group });
457}
458
459pub fn extract_generated_environment_map_entities(
460    query: Extract<
461        Query<(
462            RenderEntity,
463            &GeneratedEnvironmentMapLight,
464            &EnvironmentMapLight,
465        )>,
466    >,
467    mut commands: Commands,
468    render_images: Res<RenderAssets<GpuImage>>,
469) {
470    for (entity, filtered_env_map, env_map_light) in query.iter() {
471        let Some(env_map) = render_images.get(&filtered_env_map.environment_map) else {
472            continue;
473        };
474
475        let diffuse_map = render_images.get(&env_map_light.diffuse_map);
476        let specular_map = render_images.get(&env_map_light.specular_map);
477
478        // continue if the diffuse map is not found
479        if diffuse_map.is_none() || specular_map.is_none() {
480            continue;
481        }
482
483        let diffuse_map = diffuse_map.unwrap();
484        let specular_map = specular_map.unwrap();
485
486        let render_filtered_env_map = RenderEnvironmentMap {
487            environment_map: env_map.clone(),
488            diffuse_map: diffuse_map.clone(),
489            specular_map: specular_map.clone(),
490            intensity: filtered_env_map.intensity,
491            rotation: filtered_env_map.rotation,
492            affects_lightmapped_mesh_diffuse: filtered_env_map.affects_lightmapped_mesh_diffuse,
493        };
494        commands
495            .get_entity(entity)
496            .expect("Entity not synced to render world")
497            .insert(render_filtered_env_map);
498    }
499}
500
501// A render-world specific version of FilteredEnvironmentMapLight that uses CachedTexture
502#[derive(Component, Clone)]
503pub struct RenderEnvironmentMap {
504    pub environment_map: GpuImage,
505    pub diffuse_map: GpuImage,
506    pub specular_map: GpuImage,
507    pub intensity: f32,
508    pub rotation: Quat,
509    pub affects_lightmapped_mesh_diffuse: bool,
510}
511
512#[derive(Component)]
513pub struct IntermediateTextures {
514    pub environment_map: CachedTexture,
515}
516
517/// Returns the total number of mip levels for the provided square texture size.
518/// `size` must be a power of two greater than zero. For example, `size = 512` → `9`.
519#[inline]
520fn compute_mip_count(size: u32) -> u32 {
521    debug_assert!(size.is_power_of_two());
522    32 - size.leading_zeros()
523}
524
525/// Prepares textures needed for single pass downsampling
526pub fn prepare_generated_environment_map_intermediate_textures(
527    light_probes: Query<(Entity, &RenderEnvironmentMap)>,
528    render_device: Res<RenderDevice>,
529    mut texture_cache: ResMut<TextureCache>,
530    mut commands: Commands,
531) {
532    for (entity, env_map_light) in &light_probes {
533        let base_size = env_map_light.environment_map.size.width;
534        let mip_level_count = compute_mip_count(base_size);
535
536        let environment_map = texture_cache.get(
537            &render_device,
538            TextureDescriptor {
539                label: Some("intermediate_environment_map"),
540                size: Extent3d {
541                    width: base_size,
542                    height: base_size,
543                    depth_or_array_layers: 6, // Cubemap faces
544                },
545                mip_level_count,
546                sample_count: 1,
547                dimension: TextureDimension::D2,
548                format: TextureFormat::Rgba16Float,
549                usage: TextureUsages::TEXTURE_BINDING
550                    | TextureUsages::STORAGE_BINDING
551                    | TextureUsages::COPY_DST,
552                view_formats: &[],
553            },
554        );
555
556        commands
557            .entity(entity)
558            .insert(IntermediateTextures { environment_map });
559    }
560}
561
562/// Shader constants for downsampling algorithm
563#[derive(Clone, Copy, ShaderType)]
564#[repr(C)]
565pub struct DownsamplingConstants {
566    mips: u32,
567    inverse_input_size: Vec2,
568    _padding: u32,
569}
570
571/// Constants for filtering
572#[derive(Clone, Copy, ShaderType)]
573#[repr(C)]
574pub struct FilteringConstants {
575    mip_level: f32,
576    sample_count: u32,
577    roughness: f32,
578    noise_size_bits: UVec2,
579}
580
581/// Stores bind groups for the environment map generation pipelines
582#[derive(Component)]
583pub struct GeneratorBindGroups {
584    pub downsampling_first: BindGroup,
585    pub downsampling_second: BindGroup,
586    pub radiance: Vec<BindGroup>, // One per mip level
587    pub irradiance: BindGroup,
588    pub copy: BindGroup,
589}
590
591/// Prepares bind groups for environment map generation pipelines
592pub fn prepare_generated_environment_map_bind_groups(
593    light_probes: Query<
594        (Entity, &IntermediateTextures, &RenderEnvironmentMap),
595        With<RenderEnvironmentMap>,
596    >,
597    render_device: Res<RenderDevice>,
598    queue: Res<RenderQueue>,
599    layouts: Res<GeneratorBindGroupLayouts>,
600    samplers: Res<GeneratorSamplers>,
601    render_images: Res<RenderAssets<GpuImage>>,
602    bluenoise: Res<Bluenoise>,
603    config: Res<DownsamplingConfig>,
604    mut commands: Commands,
605) {
606    // Skip until the blue-noise texture is available to avoid panicking.
607    // The system will retry next frame once the asset has loaded.
608    let Some(stbn_texture) = render_images.get(&bluenoise.texture) else {
609        return;
610    };
611
612    assert!(stbn_texture.size.width.is_power_of_two());
613    assert!(stbn_texture.size.height.is_power_of_two());
614    let noise_size_bits = UVec2::new(
615        stbn_texture.size.width.trailing_zeros(),
616        stbn_texture.size.height.trailing_zeros(),
617    );
618
619    for (entity, textures, env_map_light) in &light_probes {
620        // Determine mip chain based on input size
621        let base_size = env_map_light.environment_map.size.width;
622        let mip_count = compute_mip_count(base_size);
623        let last_mip = mip_count - 1;
624        let env_map_texture = env_map_light.environment_map.texture.clone();
625
626        // Create downsampling constants
627        let downsampling_constants = DownsamplingConstants {
628            mips: mip_count - 1, // Number of mips we are generating (excluding mip 0)
629            inverse_input_size: Vec2::new(1.0 / base_size as f32, 1.0 / base_size as f32),
630            _padding: 0,
631        };
632
633        let mut downsampling_constants_buffer = UniformBuffer::from(downsampling_constants);
634        downsampling_constants_buffer.write_buffer(&render_device, &queue);
635
636        let input_env_map_first = env_map_texture.clone().create_view(&TextureViewDescriptor {
637            dimension: Some(TextureViewDimension::D2Array),
638            ..Default::default()
639        });
640
641        // Utility closure to get a unique storage view for a given mip level.
642        let mip_storage = |level: u32| {
643            if level <= last_mip {
644                create_storage_view(&textures.environment_map.texture, level, &render_device)
645            } else {
646                // Return a fresh 1×1 placeholder view so each binding has its own sub-resource and cannot alias.
647                create_placeholder_storage_view(&render_device)
648            }
649        };
650
651        // Depending on device limits, build either a combined or split bind group layout
652        let (downsampling_first_bind_group, downsampling_second_bind_group) =
653            if config.combine_bind_group {
654                // Combined layout expects destinations 1–12 in both bind groups
655                let bind_group = render_device.create_bind_group(
656                    "downsampling_bind_group_combined_first",
657                    &layouts.downsampling_first,
658                    &BindGroupEntries::sequential((
659                        &samplers.linear,
660                        &downsampling_constants_buffer,
661                        &input_env_map_first,
662                        &mip_storage(1),
663                        &mip_storage(2),
664                        &mip_storage(3),
665                        &mip_storage(4),
666                        &mip_storage(5),
667                        &mip_storage(6),
668                        &mip_storage(7),
669                        &mip_storage(8),
670                        &mip_storage(9),
671                        &mip_storage(10),
672                        &mip_storage(11),
673                        &mip_storage(12),
674                    )),
675                );
676
677                (bind_group.clone(), bind_group)
678            } else {
679                // Split path requires a separate view for mip6 input
680                let input_env_map_second = env_map_texture.create_view(&TextureViewDescriptor {
681                    dimension: Some(TextureViewDimension::D2Array),
682                    base_mip_level: min(6, last_mip),
683                    mip_level_count: Some(1),
684                    ..Default::default()
685                });
686
687                // Split layout (current behavior)
688                let first = render_device.create_bind_group(
689                    "downsampling_first_bind_group",
690                    &layouts.downsampling_first,
691                    &BindGroupEntries::sequential((
692                        &samplers.linear,
693                        &downsampling_constants_buffer,
694                        &input_env_map_first,
695                        &mip_storage(1),
696                        &mip_storage(2),
697                        &mip_storage(3),
698                        &mip_storage(4),
699                        &mip_storage(5),
700                        &mip_storage(6),
701                    )),
702                );
703
704                let second = render_device.create_bind_group(
705                    "downsampling_second_bind_group",
706                    &layouts.downsampling_second,
707                    &BindGroupEntries::sequential((
708                        &samplers.linear,
709                        &downsampling_constants_buffer,
710                        &input_env_map_second,
711                        &mip_storage(7),
712                        &mip_storage(8),
713                        &mip_storage(9),
714                        &mip_storage(10),
715                        &mip_storage(11),
716                        &mip_storage(12),
717                    )),
718                );
719
720                (first, second)
721            };
722
723        // create a 2d array view of the bluenoise texture
724        let stbn_texture_view = stbn_texture
725            .texture
726            .clone()
727            .create_view(&TextureViewDescriptor {
728                dimension: Some(TextureViewDimension::D2Array),
729                ..Default::default()
730            });
731
732        // Create radiance map bind groups for each mip level
733        let num_mips = mip_count as usize;
734        let mut radiance_bind_groups = Vec::with_capacity(num_mips);
735
736        for mip in 0..num_mips {
737            // Calculate roughness from 0.0 (mip 0) to 0.889 (mip 8)
738            // We don't need roughness=1.0 as a mip level because it's handled by the separate diffuse irradiance map
739            let roughness = mip as f32 / (num_mips - 1) as f32;
740            let sample_count = 32u32 * 2u32.pow((roughness * 4.0) as u32);
741
742            let radiance_constants = FilteringConstants {
743                mip_level: mip as f32,
744                sample_count,
745                roughness,
746                noise_size_bits,
747            };
748
749            let mut radiance_constants_buffer = UniformBuffer::from(radiance_constants);
750            radiance_constants_buffer.write_buffer(&render_device, &queue);
751
752            let mip_storage_view = create_storage_view(
753                &env_map_light.specular_map.texture,
754                mip as u32,
755                &render_device,
756            );
757            let bind_group = render_device.create_bind_group(
758                Some(format!("radiance_bind_group_mip_{mip}").as_str()),
759                &layouts.radiance,
760                &BindGroupEntries::sequential((
761                    &textures.environment_map.default_view,
762                    &samplers.linear,
763                    &mip_storage_view,
764                    &radiance_constants_buffer,
765                    &stbn_texture_view,
766                )),
767            );
768
769            radiance_bind_groups.push(bind_group);
770        }
771
772        // Create irradiance bind group
773        let irradiance_constants = FilteringConstants {
774            mip_level: 0.0,
775            // 32 phi, 32 theta = 1024 samples total
776            sample_count: 1024,
777            roughness: 1.0,
778            noise_size_bits,
779        };
780
781        let mut irradiance_constants_buffer = UniformBuffer::from(irradiance_constants);
782        irradiance_constants_buffer.write_buffer(&render_device, &queue);
783
784        // create a 2d array view
785        let irradiance_map =
786            env_map_light
787                .diffuse_map
788                .texture
789                .create_view(&TextureViewDescriptor {
790                    dimension: Some(TextureViewDimension::D2Array),
791                    ..Default::default()
792                });
793
794        let irradiance_bind_group = render_device.create_bind_group(
795            "irradiance_bind_group",
796            &layouts.irradiance,
797            &BindGroupEntries::sequential((
798                &textures.environment_map.default_view,
799                &samplers.linear,
800                &irradiance_map,
801                &irradiance_constants_buffer,
802                &stbn_texture_view,
803            )),
804        );
805
806        // Create copy bind group (source env map → destination mip0)
807        let src_view = env_map_light
808            .environment_map
809            .texture
810            .create_view(&TextureViewDescriptor {
811                dimension: Some(TextureViewDimension::D2Array),
812                ..Default::default()
813            });
814
815        let dst_view = create_storage_view(&textures.environment_map.texture, 0, &render_device);
816
817        let copy_bind_group = render_device.create_bind_group(
818            "copy_bind_group",
819            &layouts.copy,
820            &BindGroupEntries::with_indices(((0, &src_view), (1, &dst_view))),
821        );
822
823        commands.entity(entity).insert(GeneratorBindGroups {
824            downsampling_first: downsampling_first_bind_group,
825            downsampling_second: downsampling_second_bind_group,
826            radiance: radiance_bind_groups,
827            irradiance: irradiance_bind_group,
828            copy: copy_bind_group,
829        });
830    }
831}
832
833/// Helper function to create a storage texture view for a specific mip level
834fn create_storage_view(texture: &Texture, mip: u32, _render_device: &RenderDevice) -> TextureView {
835    texture.create_view(&TextureViewDescriptor {
836        label: Some(format!("storage_view_mip_{mip}").as_str()),
837        format: Some(texture.format()),
838        dimension: Some(TextureViewDimension::D2Array),
839        aspect: TextureAspect::All,
840        base_mip_level: mip,
841        mip_level_count: Some(1),
842        base_array_layer: 0,
843        array_layer_count: Some(texture.depth_or_array_layers()),
844        usage: Some(TextureUsages::STORAGE_BINDING),
845    })
846}
847
848/// To ensure compatibility in web browsers, each call returns a unique resource so that multiple missing mip
849/// bindings in the same bind-group never alias.
850fn create_placeholder_storage_view(render_device: &RenderDevice) -> TextureView {
851    let tex = render_device.create_texture(&TextureDescriptor {
852        label: Some("lightprobe_placeholder"),
853        size: Extent3d {
854            width: 1,
855            height: 1,
856            depth_or_array_layers: 6,
857        },
858        mip_level_count: 1,
859        sample_count: 1,
860        dimension: TextureDimension::D2,
861        format: TextureFormat::Rgba16Float,
862        usage: TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING,
863        view_formats: &[],
864    });
865
866    tex.create_view(&TextureViewDescriptor::default())
867}
868
869/// Downsampling node implementation that handles all parts of the mip chain
870pub struct DownsamplingNode {
871    query: QueryState<(
872        Entity,
873        Read<GeneratorBindGroups>,
874        Read<RenderEnvironmentMap>,
875    )>,
876}
877
878impl FromWorld for DownsamplingNode {
879    fn from_world(world: &mut World) -> Self {
880        Self {
881            query: QueryState::new(world),
882        }
883    }
884}
885
886impl Node for DownsamplingNode {
887    fn update(&mut self, world: &mut World) {
888        self.query.update_archetypes(world);
889    }
890
891    fn run(
892        &self,
893        _graph: &mut RenderGraphContext,
894        render_context: &mut RenderContext,
895        world: &World,
896    ) -> Result<(), NodeRunError> {
897        let pipeline_cache = world.resource::<PipelineCache>();
898        let pipelines = world.resource::<GeneratorPipelines>();
899
900        let Some(downsample_first_pipeline) =
901            pipeline_cache.get_compute_pipeline(pipelines.downsample_first)
902        else {
903            return Ok(());
904        };
905
906        let Some(downsample_second_pipeline) =
907            pipeline_cache.get_compute_pipeline(pipelines.downsample_second)
908        else {
909            return Ok(());
910        };
911
912        let diagnostics = render_context.diagnostic_recorder();
913
914        for (_, bind_groups, env_map_light) in self.query.iter_manual(world) {
915            // Copy base mip using compute shader with pre-built bind group
916            let Some(copy_pipeline) = pipeline_cache.get_compute_pipeline(pipelines.copy) else {
917                return Ok(());
918            };
919
920            {
921                let mut compute_pass =
922                    render_context
923                        .command_encoder()
924                        .begin_compute_pass(&ComputePassDescriptor {
925                            label: Some("lightprobe_copy"),
926                            timestamp_writes: None,
927                        });
928
929                let pass_span = diagnostics.pass_span(&mut compute_pass, "lightprobe_copy");
930
931                compute_pass.set_pipeline(copy_pipeline);
932                compute_pass.set_bind_group(0, &bind_groups.copy, &[]);
933
934                let tex_size = env_map_light.environment_map.size;
935                let wg_x = tex_size.width.div_ceil(8);
936                let wg_y = tex_size.height.div_ceil(8);
937                compute_pass.dispatch_workgroups(wg_x, wg_y, 6);
938
939                pass_span.end(&mut compute_pass);
940            }
941
942            // First pass - process mips 0-5
943            {
944                let mut compute_pass =
945                    render_context
946                        .command_encoder()
947                        .begin_compute_pass(&ComputePassDescriptor {
948                            label: Some("lightprobe_downsampling_first_pass"),
949                            timestamp_writes: None,
950                        });
951
952                let pass_span =
953                    diagnostics.pass_span(&mut compute_pass, "lightprobe_downsampling_first_pass");
954
955                compute_pass.set_pipeline(downsample_first_pipeline);
956                compute_pass.set_bind_group(0, &bind_groups.downsampling_first, &[]);
957
958                let tex_size = env_map_light.environment_map.size;
959                let wg_x = tex_size.width.div_ceil(64);
960                let wg_y = tex_size.height.div_ceil(64);
961                compute_pass.dispatch_workgroups(wg_x, wg_y, 6); // 6 faces
962
963                pass_span.end(&mut compute_pass);
964            }
965
966            // Second pass - process mips 6-12
967            {
968                let mut compute_pass =
969                    render_context
970                        .command_encoder()
971                        .begin_compute_pass(&ComputePassDescriptor {
972                            label: Some("lightprobe_downsampling_second_pass"),
973                            timestamp_writes: None,
974                        });
975
976                let pass_span =
977                    diagnostics.pass_span(&mut compute_pass, "lightprobe_downsampling_second_pass");
978
979                compute_pass.set_pipeline(downsample_second_pipeline);
980                compute_pass.set_bind_group(0, &bind_groups.downsampling_second, &[]);
981
982                let tex_size = env_map_light.environment_map.size;
983                let wg_x = tex_size.width.div_ceil(256);
984                let wg_y = tex_size.height.div_ceil(256);
985                compute_pass.dispatch_workgroups(wg_x, wg_y, 6);
986
987                pass_span.end(&mut compute_pass);
988            }
989        }
990
991        Ok(())
992    }
993}
994
995/// Radiance map node for generating specular environment maps
996pub struct FilteringNode {
997    query: QueryState<(
998        Entity,
999        Read<GeneratorBindGroups>,
1000        Read<RenderEnvironmentMap>,
1001    )>,
1002}
1003
1004impl FromWorld for FilteringNode {
1005    fn from_world(world: &mut World) -> Self {
1006        Self {
1007            query: QueryState::new(world),
1008        }
1009    }
1010}
1011
1012impl Node for FilteringNode {
1013    fn update(&mut self, world: &mut World) {
1014        self.query.update_archetypes(world);
1015    }
1016
1017    fn run(
1018        &self,
1019        _graph: &mut RenderGraphContext,
1020        render_context: &mut RenderContext,
1021        world: &World,
1022    ) -> Result<(), NodeRunError> {
1023        let pipeline_cache = world.resource::<PipelineCache>();
1024        let pipelines = world.resource::<GeneratorPipelines>();
1025
1026        let Some(radiance_pipeline) = pipeline_cache.get_compute_pipeline(pipelines.radiance)
1027        else {
1028            return Ok(());
1029        };
1030        let Some(irradiance_pipeline) = pipeline_cache.get_compute_pipeline(pipelines.irradiance)
1031        else {
1032            return Ok(());
1033        };
1034
1035        let diagnostics = render_context.diagnostic_recorder();
1036
1037        for (_, bind_groups, env_map_light) in self.query.iter_manual(world) {
1038            let mut compute_pass =
1039                render_context
1040                    .command_encoder()
1041                    .begin_compute_pass(&ComputePassDescriptor {
1042                        label: Some("lightprobe_radiance_map"),
1043                        timestamp_writes: None,
1044                    });
1045
1046            let pass_span = diagnostics.pass_span(&mut compute_pass, "lightprobe_radiance_map");
1047
1048            compute_pass.set_pipeline(radiance_pipeline);
1049
1050            let base_size = env_map_light.specular_map.size.width;
1051
1052            // Radiance convolution pass
1053            // Process each mip at different roughness levels
1054            for (mip, bind_group) in bind_groups.radiance.iter().enumerate() {
1055                compute_pass.set_bind_group(0, bind_group, &[]);
1056
1057                // Calculate dispatch size based on mip level
1058                let mip_size = base_size >> mip;
1059                let workgroup_count = mip_size.div_ceil(8);
1060
1061                // Dispatch for all 6 faces
1062                compute_pass.dispatch_workgroups(workgroup_count, workgroup_count, 6);
1063            }
1064            pass_span.end(&mut compute_pass);
1065            // End the compute pass before starting the next one
1066            drop(compute_pass);
1067
1068            // Irradiance convolution pass
1069            // Generate the diffuse environment map
1070            {
1071                let mut compute_pass =
1072                    render_context
1073                        .command_encoder()
1074                        .begin_compute_pass(&ComputePassDescriptor {
1075                            label: Some("lightprobe_irradiance_map"),
1076                            timestamp_writes: None,
1077                        });
1078
1079                let irr_span =
1080                    diagnostics.pass_span(&mut compute_pass, "lightprobe_irradiance_map");
1081
1082                compute_pass.set_pipeline(irradiance_pipeline);
1083                compute_pass.set_bind_group(0, &bind_groups.irradiance, &[]);
1084
1085                // 32×32 texture processed with 8×8 workgroups for all 6 faces
1086                compute_pass.dispatch_workgroups(4, 4, 6);
1087
1088                irr_span.end(&mut compute_pass);
1089            }
1090        }
1091
1092        Ok(())
1093    }
1094}
1095
1096/// System that generates an `EnvironmentMapLight` component based on the `GeneratedEnvironmentMapLight` component
1097pub fn generate_environment_map_light(
1098    mut commands: Commands,
1099    mut images: ResMut<Assets<Image>>,
1100    query: Query<(Entity, &GeneratedEnvironmentMapLight), Without<EnvironmentMapLight>>,
1101) {
1102    for (entity, filtered_env_map) in &query {
1103        // Validate and fetch the source cubemap so we can size our targets correctly
1104        let Some(src_image) = images.get(&filtered_env_map.environment_map) else {
1105            // Texture not ready yet – try again next frame
1106            continue;
1107        };
1108
1109        let base_size = src_image.texture_descriptor.size.width;
1110
1111        // Sanity checks – square, power-of-two, ≤ 8192
1112        if src_image.texture_descriptor.size.height != base_size
1113            || !base_size.is_power_of_two()
1114            || base_size > 8192
1115        {
1116            panic!(
1117                "GeneratedEnvironmentMapLight source cubemap must be square power-of-two ≤ 8192, got {}×{}",
1118                base_size, src_image.texture_descriptor.size.height
1119            );
1120        }
1121
1122        let mip_count = compute_mip_count(base_size);
1123
1124        // Create a placeholder for the irradiance map
1125        let mut diffuse = Image::new_fill(
1126            Extent3d {
1127                width: 32,
1128                height: 32,
1129                depth_or_array_layers: 6,
1130            },
1131            TextureDimension::D2,
1132            &[0; 8],
1133            TextureFormat::Rgba16Float,
1134            RenderAssetUsages::all(),
1135        );
1136
1137        diffuse.texture_descriptor.usage =
1138            TextureUsages::TEXTURE_BINDING | TextureUsages::STORAGE_BINDING;
1139
1140        diffuse.texture_view_descriptor = Some(TextureViewDescriptor {
1141            dimension: Some(TextureViewDimension::Cube),
1142            ..Default::default()
1143        });
1144
1145        let diffuse_handle = images.add(diffuse);
1146
1147        // Create a placeholder for the specular map. It matches the input cubemap resolution.
1148        let mut specular = Image::new_fill(
1149            Extent3d {
1150                width: base_size,
1151                height: base_size,
1152                depth_or_array_layers: 6,
1153            },
1154            TextureDimension::D2,
1155            &[0; 8],
1156            TextureFormat::Rgba16Float,
1157            RenderAssetUsages::all(),
1158        );
1159
1160        // Set up for mipmaps
1161        specular.texture_descriptor.usage =
1162            TextureUsages::TEXTURE_BINDING | TextureUsages::STORAGE_BINDING;
1163        specular.texture_descriptor.mip_level_count = mip_count;
1164
1165        // When setting mip_level_count, we need to allocate appropriate data size
1166        // For GPU-generated mipmaps, we can set data to None since the GPU will generate the data
1167        specular.data = None;
1168
1169        specular.texture_view_descriptor = Some(TextureViewDescriptor {
1170            dimension: Some(TextureViewDimension::Cube),
1171            mip_level_count: Some(mip_count),
1172            ..Default::default()
1173        });
1174
1175        let specular_handle = images.add(specular);
1176
1177        // Add the EnvironmentMapLight component with the placeholder handles
1178        commands.entity(entity).insert(EnvironmentMapLight {
1179            diffuse_map: diffuse_handle,
1180            specular_map: specular_handle,
1181            intensity: filtered_env_map.intensity,
1182            rotation: filtered_env_map.rotation,
1183            affects_lightmapped_mesh_diffuse: filtered_env_map.affects_lightmapped_mesh_diffuse,
1184        });
1185    }
1186}