use bevy_app::{App, Plugin};
use bevy_asset::AssetId;
use bevy_camera::{
primitives::{Aabb, Frustum},
Camera3d,
};
use bevy_derive::{Deref, DerefMut};
use bevy_ecs::{
component::Component,
entity::Entity,
query::With,
resource::Resource,
schedule::IntoScheduleConfigs,
system::{Commands, Local, Query, Res, ResMut},
};
use bevy_image::Image;
use bevy_light::{EnvironmentMapLight, IrradianceVolume, LightProbe};
use bevy_math::{Affine3A, FloatOrd, Mat4, Vec3A, Vec4};
use bevy_platform::collections::HashMap;
use bevy_render::{
extract_instances::ExtractInstancesPlugin,
render_asset::RenderAssets,
render_resource::{DynamicUniformBuffer, Sampler, ShaderType, TextureView},
renderer::{RenderAdapter, RenderAdapterInfo, RenderDevice, RenderQueue, WgpuWrapper},
settings::WgpuFeatures,
sync_world::RenderEntity,
texture::{FallbackImage, GpuImage},
view::ExtractedView,
Extract, ExtractSchedule, Render, RenderApp, RenderSystems,
};
use bevy_shader::load_shader_library;
use bevy_transform::{components::Transform, prelude::GlobalTransform};
use tracing::error;
use core::{hash::Hash, ops::Deref};
use crate::{
generate::EnvironmentMapGenerationPlugin, light_probe::environment_map::EnvironmentMapIds,
};
pub mod environment_map;
pub mod generate;
pub mod irradiance_volume;
pub const MAX_VIEW_LIGHT_PROBES: usize = 8;
const STANDARD_MATERIAL_FRAGMENT_SHADER_MIN_TEXTURE_BINDINGS: usize = 16;
pub struct LightProbePlugin;
#[derive(Clone, Copy, ShaderType, Default)]
struct RenderLightProbe {
light_from_world_transposed: [Vec4; 3],
texture_index: i32,
intensity: f32,
affects_lightmapped_mesh_diffuse: u32,
}
#[derive(ShaderType)]
pub struct LightProbesUniform {
reflection_probes: [RenderLightProbe; MAX_VIEW_LIGHT_PROBES],
irradiance_volumes: [RenderLightProbe; MAX_VIEW_LIGHT_PROBES],
reflection_probe_count: i32,
irradiance_volume_count: i32,
view_cubemap_index: i32,
smallest_specular_mip_level_for_view: u32,
intensity_for_view: f32,
view_environment_map_affects_lightmapped_mesh_diffuse: u32,
}
#[derive(Resource, Default, Deref, DerefMut)]
pub struct LightProbesBuffer(DynamicUniformBuffer<LightProbesUniform>);
#[derive(Component, Default, Deref, DerefMut)]
pub struct ViewLightProbesUniformOffset(u32);
struct LightProbeInfo<C>
where
C: LightProbeComponent,
{
light_from_world: [Vec4; 3],
world_from_light: Affine3A,
intensity: f32,
affects_lightmapped_mesh_diffuse: bool,
asset_id: C::AssetId,
}
#[derive(Component, Default)]
pub struct RenderViewLightProbes<C>
where
C: LightProbeComponent,
{
binding_index_to_textures: Vec<C::AssetId>,
cubemap_to_binding_index: HashMap<C::AssetId, u32>,
render_light_probes: Vec<RenderLightProbe>,
view_light_probe_info: C::ViewLightProbeInfo,
}
pub trait LightProbeComponent: Send + Sync + Component + Sized {
type AssetId: Send + Sync + Clone + Eq + Hash;
type ViewLightProbeInfo: Send + Sync + Default;
fn id(&self, image_assets: &RenderAssets<GpuImage>) -> Option<Self::AssetId>;
fn intensity(&self) -> f32;
fn affects_lightmapped_mesh_diffuse(&self) -> bool;
fn create_render_view_light_probes(
view_component: Option<&Self>,
image_assets: &RenderAssets<GpuImage>,
) -> RenderViewLightProbes<Self>;
}
#[derive(Component, ShaderType, Clone)]
pub struct EnvironmentMapUniform {
transform: Mat4,
}
impl Default for EnvironmentMapUniform {
fn default() -> Self {
EnvironmentMapUniform {
transform: Mat4::IDENTITY,
}
}
}
#[derive(Resource, Default, Deref, DerefMut)]
pub struct EnvironmentMapUniformBuffer(pub DynamicUniformBuffer<EnvironmentMapUniform>);
#[derive(Component, Default, Deref, DerefMut)]
pub struct ViewEnvironmentMapUniformOffset(u32);
impl Plugin for LightProbePlugin {
fn build(&self, app: &mut App) {
load_shader_library!(app, "light_probe.wgsl");
load_shader_library!(app, "environment_map.wgsl");
load_shader_library!(app, "irradiance_volume.wgsl");
app.add_plugins((
EnvironmentMapGenerationPlugin,
ExtractInstancesPlugin::<EnvironmentMapIds>::new(),
));
let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
return;
};
render_app
.init_resource::<LightProbesBuffer>()
.init_resource::<EnvironmentMapUniformBuffer>()
.add_systems(ExtractSchedule, gather_environment_map_uniform)
.add_systems(ExtractSchedule, gather_light_probes::<EnvironmentMapLight>)
.add_systems(ExtractSchedule, gather_light_probes::<IrradianceVolume>)
.add_systems(
Render,
(upload_light_probes, prepare_environment_uniform_buffer)
.in_set(RenderSystems::PrepareResources),
);
}
}
fn gather_environment_map_uniform(
view_query: Extract<Query<(RenderEntity, Option<&EnvironmentMapLight>), With<Camera3d>>>,
mut commands: Commands,
) {
for (view_entity, environment_map_light) in view_query.iter() {
let environment_map_uniform = if let Some(environment_map_light) = environment_map_light {
EnvironmentMapUniform {
transform: Transform::from_rotation(environment_map_light.rotation)
.to_matrix()
.inverse(),
}
} else {
EnvironmentMapUniform::default()
};
commands
.get_entity(view_entity)
.expect("Environment map light entity wasn't synced.")
.insert(environment_map_uniform);
}
}
fn gather_light_probes<C>(
image_assets: Res<RenderAssets<GpuImage>>,
light_probe_query: Extract<Query<(&GlobalTransform, &C), With<LightProbe>>>,
view_query: Extract<
Query<(RenderEntity, &GlobalTransform, &Frustum, Option<&C>), With<Camera3d>>,
>,
mut reflection_probes: Local<Vec<LightProbeInfo<C>>>,
mut view_reflection_probes: Local<Vec<LightProbeInfo<C>>>,
mut commands: Commands,
) where
C: LightProbeComponent,
{
reflection_probes.clear();
reflection_probes.extend(
light_probe_query
.iter()
.filter_map(|query_row| LightProbeInfo::new(query_row, &image_assets)),
);
for (view_entity, view_transform, view_frustum, view_component) in view_query.iter() {
view_reflection_probes.clear();
view_reflection_probes.extend(
reflection_probes
.iter()
.filter(|light_probe_info| light_probe_info.frustum_cull(view_frustum))
.cloned(),
);
view_reflection_probes.sort_by_cached_key(|light_probe_info| {
light_probe_info.camera_distance_sort_key(view_transform)
});
let mut render_view_light_probes =
C::create_render_view_light_probes(view_component, &image_assets);
render_view_light_probes.maybe_gather_light_probes(&view_reflection_probes);
if render_view_light_probes.is_empty() {
commands
.get_entity(view_entity)
.expect("View entity wasn't synced.")
.remove::<RenderViewLightProbes<C>>();
} else {
commands
.get_entity(view_entity)
.expect("View entity wasn't synced.")
.insert(render_view_light_probes);
}
}
}
pub fn prepare_environment_uniform_buffer(
mut commands: Commands,
views: Query<(Entity, Option<&EnvironmentMapUniform>), With<ExtractedView>>,
mut environment_uniform_buffer: ResMut<EnvironmentMapUniformBuffer>,
render_device: Res<RenderDevice>,
render_queue: Res<RenderQueue>,
) {
let Some(mut writer) =
environment_uniform_buffer.get_writer(views.iter().len(), &render_device, &render_queue)
else {
return;
};
for (view, environment_uniform) in views.iter() {
let uniform_offset = match environment_uniform {
None => 0,
Some(environment_uniform) => writer.write(environment_uniform),
};
commands
.entity(view)
.insert(ViewEnvironmentMapUniformOffset(uniform_offset));
}
}
fn upload_light_probes(
mut commands: Commands,
views: Query<Entity, With<ExtractedView>>,
mut light_probes_buffer: ResMut<LightProbesBuffer>,
mut view_light_probes_query: Query<(
Option<&RenderViewLightProbes<EnvironmentMapLight>>,
Option<&RenderViewLightProbes<IrradianceVolume>>,
)>,
render_device: Res<RenderDevice>,
render_queue: Res<RenderQueue>,
) {
if views.is_empty() {
return;
}
let mut writer = light_probes_buffer
.get_writer(views.iter().len(), &render_device, &render_queue)
.unwrap();
for view_entity in views.iter() {
let Ok((render_view_environment_maps, render_view_irradiance_volumes)) =
view_light_probes_query.get_mut(view_entity)
else {
error!("Failed to find `RenderViewLightProbes` for the view!");
continue;
};
let mut light_probes_uniform = LightProbesUniform {
reflection_probes: [RenderLightProbe::default(); MAX_VIEW_LIGHT_PROBES],
irradiance_volumes: [RenderLightProbe::default(); MAX_VIEW_LIGHT_PROBES],
reflection_probe_count: render_view_environment_maps
.map(RenderViewLightProbes::len)
.unwrap_or_default()
.min(MAX_VIEW_LIGHT_PROBES) as i32,
irradiance_volume_count: render_view_irradiance_volumes
.map(RenderViewLightProbes::len)
.unwrap_or_default()
.min(MAX_VIEW_LIGHT_PROBES) as i32,
view_cubemap_index: render_view_environment_maps
.map(|maps| maps.view_light_probe_info.cubemap_index)
.unwrap_or(-1),
smallest_specular_mip_level_for_view: render_view_environment_maps
.map(|maps| maps.view_light_probe_info.smallest_specular_mip_level)
.unwrap_or(0),
intensity_for_view: render_view_environment_maps
.map(|maps| maps.view_light_probe_info.intensity)
.unwrap_or(1.0),
view_environment_map_affects_lightmapped_mesh_diffuse: render_view_environment_maps
.map(|maps| maps.view_light_probe_info.affects_lightmapped_mesh_diffuse as u32)
.unwrap_or(1),
};
if let Some(render_view_environment_maps) = render_view_environment_maps {
render_view_environment_maps.add_to_uniform(
&mut light_probes_uniform.reflection_probes,
&mut light_probes_uniform.reflection_probe_count,
);
}
if let Some(render_view_irradiance_volumes) = render_view_irradiance_volumes {
render_view_irradiance_volumes.add_to_uniform(
&mut light_probes_uniform.irradiance_volumes,
&mut light_probes_uniform.irradiance_volume_count,
);
}
let uniform_offset = writer.write(&light_probes_uniform);
commands
.entity(view_entity)
.insert(ViewLightProbesUniformOffset(uniform_offset));
}
}
impl Default for LightProbesUniform {
fn default() -> Self {
Self {
reflection_probes: [RenderLightProbe::default(); MAX_VIEW_LIGHT_PROBES],
irradiance_volumes: [RenderLightProbe::default(); MAX_VIEW_LIGHT_PROBES],
reflection_probe_count: 0,
irradiance_volume_count: 0,
view_cubemap_index: -1,
smallest_specular_mip_level_for_view: 0,
intensity_for_view: 1.0,
view_environment_map_affects_lightmapped_mesh_diffuse: 1,
}
}
}
impl<C> LightProbeInfo<C>
where
C: LightProbeComponent,
{
fn new(
(light_probe_transform, environment_map): (&GlobalTransform, &C),
image_assets: &RenderAssets<GpuImage>,
) -> Option<LightProbeInfo<C>> {
let light_from_world_transposed =
Mat4::from(light_probe_transform.affine().inverse()).transpose();
environment_map.id(image_assets).map(|id| LightProbeInfo {
world_from_light: light_probe_transform.affine(),
light_from_world: [
light_from_world_transposed.x_axis,
light_from_world_transposed.y_axis,
light_from_world_transposed.z_axis,
],
asset_id: id,
intensity: environment_map.intensity(),
affects_lightmapped_mesh_diffuse: environment_map.affects_lightmapped_mesh_diffuse(),
})
}
fn frustum_cull(&self, view_frustum: &Frustum) -> bool {
view_frustum.intersects_obb(
&Aabb {
center: Vec3A::default(),
half_extents: Vec3A::splat(0.5),
},
&self.world_from_light,
true,
false,
)
}
fn camera_distance_sort_key(&self, view_transform: &GlobalTransform) -> FloatOrd {
FloatOrd(
(self.world_from_light.translation - view_transform.translation_vec3a())
.length_squared(),
)
}
}
impl<C> RenderViewLightProbes<C>
where
C: LightProbeComponent,
{
fn new() -> RenderViewLightProbes<C> {
RenderViewLightProbes {
binding_index_to_textures: vec![],
cubemap_to_binding_index: HashMap::default(),
render_light_probes: vec![],
view_light_probe_info: C::ViewLightProbeInfo::default(),
}
}
pub(crate) fn is_empty(&self) -> bool {
self.binding_index_to_textures.is_empty()
}
pub(crate) fn len(&self) -> usize {
self.binding_index_to_textures.len()
}
pub(crate) fn get_or_insert_cubemap(&mut self, cubemap_id: &C::AssetId) -> u32 {
*self
.cubemap_to_binding_index
.entry((*cubemap_id).clone())
.or_insert_with(|| {
let index = self.binding_index_to_textures.len() as u32;
self.binding_index_to_textures.push((*cubemap_id).clone());
index
})
}
fn add_to_uniform(
&self,
render_light_probes: &mut [RenderLightProbe; MAX_VIEW_LIGHT_PROBES],
render_light_probe_count: &mut i32,
) {
render_light_probes[0..self.render_light_probes.len()]
.copy_from_slice(&self.render_light_probes[..]);
*render_light_probe_count = self.render_light_probes.len() as i32;
}
fn maybe_gather_light_probes(&mut self, light_probes: &[LightProbeInfo<C>]) {
for light_probe in light_probes.iter().take(MAX_VIEW_LIGHT_PROBES) {
let cubemap_index = self.get_or_insert_cubemap(&light_probe.asset_id);
self.render_light_probes.push(RenderLightProbe {
light_from_world_transposed: light_probe.light_from_world,
texture_index: cubemap_index as i32,
intensity: light_probe.intensity,
affects_lightmapped_mesh_diffuse: light_probe.affects_lightmapped_mesh_diffuse
as u32,
});
}
}
}
impl<C> Clone for LightProbeInfo<C>
where
C: LightProbeComponent,
{
fn clone(&self) -> Self {
Self {
light_from_world: self.light_from_world,
world_from_light: self.world_from_light,
intensity: self.intensity,
affects_lightmapped_mesh_diffuse: self.affects_lightmapped_mesh_diffuse,
asset_id: self.asset_id.clone(),
}
}
}
pub(crate) fn add_cubemap_texture_view<'a>(
texture_views: &mut Vec<&'a <TextureView as Deref>::Target>,
sampler: &mut Option<&'a Sampler>,
image_id: AssetId<Image>,
images: &'a RenderAssets<GpuImage>,
fallback_image: &'a FallbackImage,
) {
match images.get(image_id) {
None => {
texture_views.push(&*fallback_image.cube.texture_view);
}
Some(image) => {
if sampler.is_none() {
*sampler = Some(&image.sampler);
}
texture_views.push(&*image.texture_view);
}
}
}
pub(crate) fn binding_arrays_are_usable(
render_device: &RenderDevice,
render_adapter: &RenderAdapter,
) -> bool {
let adapter_info = RenderAdapterInfo(WgpuWrapper::new(render_adapter.get_info()));
!cfg!(feature = "shader_format_glsl")
&& bevy_render::get_adreno_model(&adapter_info).is_none_or(|model| model > 610)
&& render_device.limits().max_storage_textures_per_shader_stage
>= (STANDARD_MATERIAL_FRAGMENT_SHADER_MIN_TEXTURE_BINDINGS + MAX_VIEW_LIGHT_PROBES)
as u32
&& render_device.features().contains(
WgpuFeatures::TEXTURE_BINDING_ARRAY
| WgpuFeatures::SAMPLED_TEXTURE_AND_STORAGE_BUFFER_ARRAY_NON_UNIFORM_INDEXING,
)
}