Skip to main content

wde_pbr/deferred/lights/
lights_binding.rs

1use wde_logger::prelude::*;
2
3use bevy::{ecs::system::SystemParamItem, prelude::*};
4use wde_camera::prelude::ActiveCamera;
5use wde_renderer::prelude::*;
6
7use crate::deferred::lights::*;
8#[cfg(feature = "atmosphere")]
9use crate::prelude::{AtmosphereParamsUniform, AtmosphereSettings};
10use crate::prelude::{PbrParamsUniform, PbrSettings};
11
12/// Maximum number of lights rendered: beyond, the lights closest to the camera are kept.
13pub const MAX_LIGHTS: usize = 128;
14
15pub(crate) struct LightsBindingPlugin;
16impl Plugin for LightsBindingPlugin {
17    fn build(&self, app: &mut App) {
18        // Add the render bindings
19        #[cfg(feature = "atmosphere")]
20        app.init_resource::<AtmosphereSettings>()
21            .register_type::<AtmosphereSettings>();
22
23        app.init_resource::<PbrSettings>()
24            .register_type::<PbrSettings>()
25            .add_plugins(RenderDataRegisterPlugin::<LightsData>::default());
26
27        // Add the systems to extract the lights from the world, and to update the lights buffer
28        app.get_sub_app_mut(RenderApp)
29            .unwrap()
30            .init_resource::<ExtractedLights>();
31
32        #[cfg(feature = "atmosphere")]
33        app.get_sub_app_mut(RenderApp)
34            .unwrap()
35            .init_resource::<AtmosphereParamsUniform>()
36            .add_systems(Extract, extract_atmosphere_params);
37
38        app.get_sub_app_mut(RenderApp)
39            .unwrap()
40            .init_resource::<PbrParamsUniform>()
41            .add_systems(Extract, extract)
42            .add_systems(Extract, extract_pbr_params)
43            .add_systems(Render, update_lights_buffer.in_set(RenderSet::Prepare));
44    }
45}
46
47#[derive(Asset, Clone, TypePath, Default)]
48pub struct LightsData;
49impl LightsData {
50    pub const LIGHTS_BUFFER_IDX: u32 = 0;
51    pub const LIGHTS_BUFFER_STAGING_IDX: u32 = 1;
52    pub const ATMOSPHERE_PARAMS_BUFFER_IDX: u32 = 2;
53    pub const PBR_PARAMS_BUFFER_IDX: u32 = 3;
54}
55impl RenderData for LightsData {
56    type Params = ();
57
58    fn describe(_params: &mut SystemParamItem<Self::Params>, builder: &mut RenderDataBuilder) {
59        builder
60            .add_buffer(
61                Self::LIGHTS_BUFFER_IDX,
62                Buffer {
63                    label: "lights-buffer-gpu".to_string(),
64                    size: std::mem::size_of::<LightsStorageElement>() * MAX_LIGHTS,
65                    usage: BufferUsage::STORAGE | BufferUsage::COPY_DST,
66                    content: None
67                }
68            )
69            .add_buffer(
70                Self::LIGHTS_BUFFER_STAGING_IDX,
71                Buffer {
72                    label: "lights-buffer-staging".to_string(),
73                    size: std::mem::size_of::<LightsStorageElement>() * MAX_LIGHTS,
74                    usage: BufferUsage::COPY_SRC | BufferUsage::COPY_DST,
75                    content: None
76                }
77            );
78        #[cfg(feature = "atmosphere")]
79        builder.add_buffer(
80            Self::ATMOSPHERE_PARAMS_BUFFER_IDX,
81            Buffer {
82                label: "atmosphere-params".to_string(),
83                size: std::mem::size_of::<AtmosphereParamsUniform>(),
84                usage: BufferUsage::UNIFORM | BufferUsage::COPY_DST,
85                content: Some(bytemuck::cast_slice(&[AtmosphereParamsUniform::default()]).to_vec())
86            }
87        );
88        builder.add_buffer(
89            Self::PBR_PARAMS_BUFFER_IDX,
90            Buffer {
91                label: "pbr-params".to_string(),
92                size: std::mem::size_of::<PbrParamsUniform>(),
93                usage: BufferUsage::UNIFORM | BufferUsage::COPY_DST,
94                content: Some(bytemuck::cast_slice(&[PbrParamsUniform::default()]).to_vec())
95            }
96        );
97    }
98}
99
100#[derive(Resource, Default)]
101struct ExtractedLights {
102    pub directional_lights: Vec<DirectionalLight>,
103    pub point_lights: Vec<PointLight>,
104    pub spot_lights: Vec<SpotLight>
105}
106
107/// Extracts the lights in world space (lights on entities with a transform follow it), keeping
108/// the [`MAX_LIGHTS`] most relevant: directional lights, then the closest to the camera.
109fn extract(
110    lights_directional: ExtractWorld<Query<&DirectionalLight>>,
111    lights_point: ExtractWorld<Query<(&PointLight, Option<&GlobalTransform>)>>,
112    lights_spot: ExtractWorld<Query<(&SpotLight, Option<&GlobalTransform>)>>,
113    camera: ExtractWorld<Query<&GlobalTransform, With<ActiveCamera>>>,
114    mut extracted_lights: ResMut<ExtractedLights>
115) {
116    extracted_lights.directional_lights = lights_directional.iter().copied().collect();
117    extracted_lights.point_lights = lights_point
118        .iter()
119        .map(|(light, transform)| PointLight {
120            position: transform.map_or(light.position, |t| t.transform_point(light.position)),
121            ..*light
122        })
123        .collect();
124    extracted_lights.spot_lights = lights_spot
125        .iter()
126        .map(|(light, transform)| match transform {
127            Some(t) => SpotLight {
128                position: t.transform_point(light.position),
129                direction: t.rotation() * light.direction,
130                ..*light
131            },
132            None => *light
133        })
134        .collect();
135
136    let budget = MAX_LIGHTS.saturating_sub(extracted_lights.directional_lights.len());
137    if extracted_lights.point_lights.len() + extracted_lights.spot_lights.len() > budget {
138        // Distance from the camera to the light's reach
139        let eye = camera.iter().next().map_or(Vec3::ZERO, |t| t.translation());
140        let gap = |position: Vec3, range: f32| (position.distance(eye) - range).max(0.0);
141        let mut by_gap: Vec<(f32, bool, usize)> = extracted_lights
142            .point_lights
143            .iter()
144            .enumerate()
145            .map(|(i, l)| (gap(l.position, l.range), false, i))
146            .chain(
147                extracted_lights
148                    .spot_lights
149                    .iter()
150                    .enumerate()
151                    .map(|(i, l)| (gap(l.position, l.range), true, i))
152            )
153            .collect();
154        by_gap.sort_unstable_by(|a, b| a.0.total_cmp(&b.0));
155        by_gap.truncate(budget);
156        let points = &extracted_lights.point_lights;
157        let spots = &extracted_lights.spot_lights;
158        let (kept_points, kept_spots) = (
159            by_gap
160                .iter()
161                .filter(|(_, spot, _)| !spot)
162                .map(|(_, _, i)| points[*i])
163                .collect(),
164            by_gap
165                .iter()
166                .filter(|(_, spot, _)| *spot)
167                .map(|(_, _, i)| spots[*i])
168                .collect()
169        );
170        extracted_lights.point_lights = kept_points;
171        extracted_lights.spot_lights = kept_spots;
172    }
173}
174
175#[cfg(feature = "atmosphere")]
176fn extract_atmosphere_params(
177    settings: ExtractWorld<Res<AtmosphereSettings>>,
178    time: ExtractWorld<Res<Time>>,
179    mut uniform: ResMut<AtmosphereParamsUniform>
180) {
181    let day_length = settings.day_length_seconds.max(1.0);
182    let time_of_day = if settings.animate_time {
183        (time.elapsed_secs() / day_length).fract()
184    } else {
185        settings.time_of_day.fract()
186    };
187    *uniform = AtmosphereParamsUniform::from_settings(&settings, time_of_day);
188}
189
190fn extract_pbr_params(
191    settings: ExtractWorld<Res<PbrSettings>>,
192    mut uniform: ResMut<PbrParamsUniform>
193) {
194    *uniform = PbrParamsUniform::from_settings(&settings);
195}
196
197fn update_lights_buffer(
198    lights_buffer: ResRenderData<LightsData>,
199    buffers: Res<RenderAssets<GpuBuffer>>,
200    render_instance: Res<RenderInstance>,
201    extracted_lights: Res<ExtractedLights>,
202    #[cfg(feature = "atmosphere")] atmosphere_uniform: Res<AtmosphereParamsUniform>,
203    pbr_uniform: Res<PbrParamsUniform>
204) {
205    // Get the lights buffer
206    let lights_buffer_cpu = match lights_buffer.iter().next() {
207        Some((_, buffer)) => match buffers.get(
208            &buffer
209                .get_buffer(LightsData::LIGHTS_BUFFER_STAGING_IDX)
210                .unwrap()
211        ) {
212            Some(lights_buffer) => lights_buffer,
213            None => return
214        },
215        None => return
216    };
217
218    // All the lights in one write, closed by an end marker so that the shaders never read the
219    // lights of previous frames
220    let render_instance = render_instance.0.read().unwrap();
221    let mut lights: Vec<LightsStorageElement> = extracted_lights
222        .directional_lights
223        .iter()
224        .map(LightsStorageElement::from_directional)
225        .chain(
226            extracted_lights
227                .point_lights
228                .iter()
229                .map(LightsStorageElement::from_point)
230        )
231        .chain(
232            extracted_lights
233                .spot_lights
234                .iter()
235                .map(LightsStorageElement::from_spot)
236        )
237        .collect();
238    if lights.len() > MAX_LIGHTS {
239        warn!(
240            "Number of lights exceeded the maximum of {}. Some lights will be ignored in rendering.",
241            MAX_LIGHTS
242        );
243        lights.truncate(MAX_LIGHTS);
244    }
245    if lights.len() < MAX_LIGHTS {
246        lights.push(LightsStorageElement::end_marker());
247    }
248    lights_buffer_cpu
249        .buffer
250        .write(&render_instance, bytemuck::cast_slice(&lights), 0);
251
252    // Update the buffer
253    let lights_buffer_gpu = match buffers.get(
254        &lights_buffer
255            .iter()
256            .next()
257            .unwrap()
258            .1
259            .get_buffer(LightsData::LIGHTS_BUFFER_IDX)
260            .unwrap()
261    ) {
262        Some(buffer) => buffer,
263        None => return
264    };
265    lights_buffer_gpu
266        .buffer
267        .copy_from_buffer(&render_instance, &lights_buffer_cpu.buffer);
268
269    // Update atmosphere uniform buffer stored in the same render data as lights.
270    #[cfg(feature = "atmosphere")]
271    {
272        let atmosphere_buffer = match buffers.get(
273            &lights_buffer
274                .iter()
275                .next()
276                .unwrap()
277                .1
278                .get_buffer(LightsData::ATMOSPHERE_PARAMS_BUFFER_IDX)
279                .unwrap()
280        ) {
281            Some(buffer) => buffer,
282            None => return
283        };
284        atmosphere_buffer.buffer.write(
285            &render_instance,
286            bytemuck::cast_slice(&[*atmosphere_uniform]),
287            0
288        );
289    }
290
291    let pbr_buffer = match buffers.get(
292        &lights_buffer
293            .iter()
294            .next()
295            .unwrap()
296            .1
297            .get_buffer(LightsData::PBR_PARAMS_BUFFER_IDX)
298            .unwrap()
299    ) {
300        Some(buffer) => buffer,
301        None => return
302    };
303    pbr_buffer
304        .buffer
305        .write(&render_instance, bytemuck::cast_slice(&[*pbr_uniform]), 0);
306}