Skip to main content

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