1use bevy_color::ColorToComponents;
2use bevy_light::atmosphere::{ScatteringMedium, ScatteringTerm};
3
4use bevy_app::{App, Plugin};
5use bevy_asset::AssetId;
6use bevy_ecs::{
7 resource::Resource,
8 system::{Commands, Res, SystemParamItem},
9};
10use bevy_math::{ops, Vec4};
11use bevy_render::{
12 render_asset::{PrepareAssetError, RenderAsset, RenderAssetPlugin},
13 render_resource::{
14 Extent3d, FilterMode, Sampler, SamplerDescriptor, Texture, TextureDataOrder,
15 TextureDescriptor, TextureDimension, TextureFormat, TextureUsages, TextureView,
16 TextureViewDescriptor,
17 },
18 renderer::{RenderDevice, RenderQueue},
19 RenderApp, RenderStartup,
20};
21use smallvec::SmallVec;
22
23#[doc(hidden)]
24pub struct ScatteringMediumPlugin;
25
26impl Plugin for ScatteringMediumPlugin {
27 fn build(&self, app: &mut App) {
28 app.add_plugins(RenderAssetPlugin::<GpuScatteringMedium>::default());
29
30 if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
31 render_app.add_systems(RenderStartup, init_scattering_medium_sampler);
32 }
33 }
34}
35
36pub struct GpuScatteringMedium {
38 pub terms: SmallVec<[ScatteringTerm; 1]>,
40 pub falloff_resolution: u32,
43 pub phase_resolution: u32,
46 pub density_lut: Texture,
52 pub density_lut_view: TextureView,
54 pub scattering_lut: Texture,
60 pub scattering_lut_view: TextureView,
62}
63
64impl RenderAsset for GpuScatteringMedium {
65 type SourceAsset = ScatteringMedium;
66
67 type Param = (Res<'static, RenderDevice>, Res<'static, RenderQueue>);
68
69 fn prepare_asset(
70 source_asset: Self::SourceAsset,
71 _asset_id: AssetId<Self::SourceAsset>,
72 (render_device, render_queue): &mut SystemParamItem<Self::Param>,
73 _previous_asset: Option<&Self>,
74 ) -> Result<Self, PrepareAssetError<Self::SourceAsset>> {
75 let mut density: Vec<Vec4> =
76 Vec::with_capacity(2 * source_asset.falloff_resolution as usize);
77
78 density.extend((0..source_asset.falloff_resolution).map(|i| {
79 let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
80
81 source_asset
82 .terms
83 .iter()
84 .map(|term| term.absorption.extend(0.0) * term.falloff.sample(falloff))
85 .sum::<Vec4>()
86 }));
87
88 density.extend((0..source_asset.falloff_resolution).map(|i| {
89 let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
90
91 source_asset
92 .terms
93 .iter()
94 .map(|term| term.scattering.extend(0.0) * term.falloff.sample(falloff))
95 .sum::<Vec4>()
96 }));
97
98 let mut scattering: Vec<Vec4> = Vec::with_capacity(
99 source_asset.falloff_resolution as usize * source_asset.phase_resolution as usize,
100 );
101
102 const PHASE_MAPPING_N: f32 = 0.5;
104 let inv_n = 1.0 / PHASE_MAPPING_N;
105
106 scattering.extend(
107 (0..source_asset.falloff_resolution * source_asset.phase_resolution).map(|raw_i| {
108 let i = raw_i % source_asset.phase_resolution;
109 let j = raw_i / source_asset.phase_resolution;
110 let falloff = (i as f32 + 0.5) / source_asset.falloff_resolution as f32;
111 let phase = (j as f32 + 0.5) / source_asset.phase_resolution as f32;
112
113 let s = 2.0 * phase - 1.0;
118 let neg_l_dot_v = s.signum() * (1.0 - ops::powf(1.0 - s.abs(), inv_n));
119
120 source_asset
121 .terms
122 .iter()
123 .filter_map(|term| {
124 let f = term.falloff.sample(falloff);
125 term.phase
126 .sample(neg_l_dot_v)
127 .map(|phase| (term.scattering * phase.to_vec3() * f).extend(0.0))
128 })
129 .sum::<Vec4>()
130 }),
131 );
132
133 let density_lut = render_device.create_texture_with_data(
134 render_queue,
135 &TextureDescriptor {
136 label: source_asset
137 .label
138 .as_deref()
139 .map(|label| format!("{}_density_lut", label))
140 .as_deref()
141 .or(Some("scattering_medium_density_lut")),
142 size: Extent3d {
143 width: source_asset.falloff_resolution,
144 height: 2,
145 depth_or_array_layers: 1,
146 },
147 mip_level_count: 1,
148 sample_count: 1,
149 dimension: TextureDimension::D2,
150 format: TextureFormat::Rgba32Float,
151 usage: TextureUsages::TEXTURE_BINDING,
152 view_formats: &[],
153 },
154 TextureDataOrder::LayerMajor,
155 bytemuck::cast_slice(density.as_slice()),
156 );
157
158 let density_lut_view = density_lut.create_view(&TextureViewDescriptor {
159 label: source_asset
160 .label
161 .as_deref()
162 .map(|label| format!("{}_density_lut_view", label))
163 .as_deref()
164 .or(Some("scattering_medium_density_lut_view")),
165 ..Default::default()
166 });
167
168 let scattering_lut = render_device.create_texture_with_data(
169 render_queue,
170 &TextureDescriptor {
171 label: source_asset
172 .label
173 .as_deref()
174 .map(|label| format!("{}_scattering_lut", label))
175 .as_deref()
176 .or(Some("scattering_medium_scattering_lut")),
177 size: Extent3d {
178 width: source_asset.falloff_resolution,
179 height: source_asset.phase_resolution,
180 depth_or_array_layers: 1,
181 },
182 mip_level_count: 1,
183 sample_count: 1,
184 dimension: TextureDimension::D2,
185 format: TextureFormat::Rgba32Float,
186 usage: TextureUsages::TEXTURE_BINDING,
187 view_formats: &[],
188 },
189 TextureDataOrder::LayerMajor,
190 bytemuck::cast_slice(scattering.as_slice()),
191 );
192
193 let scattering_lut_view = scattering_lut.create_view(&TextureViewDescriptor {
194 label: source_asset
195 .label
196 .as_deref()
197 .map(|label| format!("{}_scattering_lut", label))
198 .as_deref()
199 .or(Some("scattering_medium_scattering_lut_view")),
200 ..Default::default()
201 });
202
203 Ok(Self {
204 terms: source_asset.terms,
205 falloff_resolution: source_asset.falloff_resolution,
206 phase_resolution: source_asset.phase_resolution,
207 density_lut,
208 density_lut_view,
209 scattering_lut,
210 scattering_lut_view,
211 })
212 }
213}
214
215#[derive(Resource)]
219pub struct ScatteringMediumSampler(Sampler);
220
221impl ScatteringMediumSampler {
222 pub fn sampler(&self) -> &Sampler {
223 &self.0
224 }
225}
226
227fn init_scattering_medium_sampler(mut commands: Commands, render_device: Res<RenderDevice>) {
228 let sampler = render_device.create_sampler(&SamplerDescriptor {
229 label: Some("scattering_medium_sampler"),
230 mag_filter: FilterMode::Linear,
231 min_filter: FilterMode::Linear,
232 ..Default::default()
233 });
234
235 commands.insert_resource(ScatteringMediumSampler(sampler));
236}