#[cfg(test)]
mod model;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use bevy::asset::{AssetPath, embedded_asset};
use bevy::camera::primitives::Aabb;
use bevy::camera::visibility::NoAutoAabb;
use bevy::image::{ImageAddressMode, ImageFilterMode, ImageSampler, ImageSamplerDescriptor};
use bevy::mesh::MeshTag;
use bevy::pbr::{ExtendedMaterial, MaterialExtension, MaterialPlugin};
use bevy::prelude::*;
use bevy::render::render_resource::{
AsBindGroup, Extent3d, ShaderType, TextureDimension, TextureFormat, TextureUsages,
};
use bevy::render::storage::ShaderBuffer;
use bevy::shader::ShaderRef;
use bevy::transform::TransformSystems;
use super::water_gpu;
use crate::clouds::CloudShadows;
use crate::plugin::{ElementsUnits, ElementsView, ensure_gpu_core, needs_depth_prepass};
use crate::terrain::lod::MAX_LEVELS;
use crate::terrain::{
NO_TERRAIN, TerrainParamsGpu, TerrainRuntime, TerrainState, assign_slots, lod_centre_uniform,
lod_uniform, patch_mesh, patch_transform, shadow_inputs,
};
use crate::waves::WaveState;
use crate::wind::WindState;
pub(crate) const NORMAL_TEXELS: u32 = 512;
pub(crate) const NORMAL_MIPS: u32 = 10;
pub(crate) const NORMAL_OCTAVES: u32 = 3;
pub(crate) const NORMAL_BASE_TEXELS: u32 = 16;
pub(crate) const FOAM_OCTAVES: u32 = 4;
pub(crate) const NORMAL_RMS_SLOPE: f32 = 0.15;
pub(crate) const NORMAL_SLOPE_GAIN: f32 = 1.8575;
const MAP_WIND_MS: f32 = 4.0;
pub(crate) const WATER_POOL_CAPACITY: usize = 192;
pub(crate) const MIN_PATCH_DEPTH_M: f32 = 0.003;
const WATER_AABB_PAD_M: f32 = 4.0;
const WATER_REFLECTANCE: f32 = 0.353_553_4;
const WATER_ROUGHNESS: f32 = 0.089;
pub(crate) const NO_STILL_LEVEL: f32 = -1.0e9;
pub(crate) const WATER_SKIRT_REACH_M: f32 = 100_000.0;
const SKIRT_TAG: u32 = 1000;
const EMBEDDED_SHADER: &str = "embedded://bevy_elements/hydrology/water.wesl";
#[derive(Resource, Clone, Copy, Debug, PartialEq, Reflect)]
pub struct StillWater {
pub level: f32,
}
#[derive(Resource, Clone, Debug, PartialEq, Reflect)]
#[cfg_attr(
feature = "serde",
derive(serde::Serialize, serde::Deserialize),
serde(default)
)]
pub struct WaterSettings {
pub absorption: Vec3,
pub refractive_index: f32,
pub silt_tint: Vec3,
pub deep_color: Vec3,
pub turbidity_scattering: f32,
pub clear_scattering: f32,
pub foam_speed_ms: f32,
pub foam_thickness_m: f32,
pub foam_albedo: f32,
pub foam_roughness: f32,
pub tile_m: f32,
pub flow_cycle_s: f32,
pub flow_distortion: f32,
pub wind_speed_ms: Option<f32>,
pub ripple_scale: f32,
}
impl Default for WaterSettings {
fn default() -> Self {
Self {
absorption: Vec3::new(0.45, 0.06, 0.03),
refractive_index: 1.33,
silt_tint: Vec3::new(0.45, 0.38, 0.25),
deep_color: Vec3::new(0.0, 0.005, 0.025),
turbidity_scattering: 6.0,
clear_scattering: 0.05,
foam_speed_ms: 1.5,
foam_thickness_m: 0.05,
foam_albedo: 0.9,
foam_roughness: 0.8,
tile_m: 8.0,
flow_cycle_s: 1.5,
flow_distortion: 1.0 / 3.0,
wind_speed_ms: None,
ripple_scale: 1.0,
}
}
}
pub(crate) fn ripple_slope(u10: f32) -> f32 {
(0.003 + 0.00512 * u10.max(0.0)).sqrt()
}
#[derive(Clone, Copy, Debug, PartialEq, ShaderType, Reflect)]
#[reflect(opaque)]
pub struct WaterParams {
pub(crate) terrain: TerrainParamsGpu,
pub(crate) lod: [Vec4; MAX_LEVELS],
pub(crate) top: Vec4,
pub(crate) lod_centre: Vec4,
pub(crate) shadow: Vec4,
pub(crate) light: Vec4,
pub(crate) absorption: Vec4,
pub(crate) silt: Vec4,
pub(crate) deep: Vec4,
pub(crate) flow: Vec4,
pub(crate) foam: Vec4,
pub(crate) waves: Vec4,
pub(crate) still: Vec4,
pub(crate) units: f32,
}
impl Default for WaterParams {
fn default() -> Self {
Self {
terrain: TerrainParamsGpu::default(),
lod: [Vec4::ZERO; MAX_LEVELS],
top: Vec4::ZERO,
lod_centre: Vec4::ZERO,
shadow: Vec4::ZERO,
light: Vec4::ZERO,
absorption: Vec4::ZERO,
silt: Vec4::ZERO,
deep: Vec4::ZERO,
flow: Vec4::ZERO,
foam: Vec4::ZERO,
waves: Vec4::ZERO,
still: Vec4::new(NO_STILL_LEVEL, 0.0, 0.0, 0.0),
units: 1.0,
}
}
}
#[derive(Asset, AsBindGroup, Reflect, Debug, Clone)]
pub struct WaterExtension {
#[uniform(100)]
pub(crate) params: WaterParams,
#[storage(101, read_only)]
pub(crate) cells: Handle<ShaderBuffer>,
#[storage(102, read_only)]
pub(crate) flow: Handle<ShaderBuffer>,
#[texture(103)]
#[sampler(104)]
pub(crate) normals: Handle<Image>,
#[texture(105)]
#[sampler(106)]
pub(crate) cloud_shadow: Option<Handle<Image>>,
#[texture(107, dimension = "2d_array")]
#[sampler(108)]
pub(crate) waves: Option<Handle<Image>>,
}
impl MaterialExtension for WaterExtension {
fn vertex_shader() -> ShaderRef {
ShaderRef::Path(AssetPath::parse(EMBEDDED_SHADER))
}
fn fragment_shader() -> ShaderRef {
ShaderRef::Path(AssetPath::parse(EMBEDDED_SHADER))
}
fn enable_prepass() -> bool {
false
}
fn enable_shadows() -> bool {
false
}
}
pub type WaterMaterial = ExtendedMaterial<StandardMaterial, WaterExtension>;
#[derive(Component, Clone, Copy, Debug, PartialEq, Eq)]
pub struct WaterPatch {
pub surface: bool,
}
#[derive(Component, Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct WaterSkirt {
tile: usize,
}
#[derive(Resource, Clone, Debug)]
pub struct WaterState {
pub normals: Handle<Image>,
pub normals_ready: Arc<AtomicBool>,
}
struct WaterSlot {
absorb: Entity,
surface: Entity,
patch: Option<crate::terrain::lod::Patch>,
}
#[derive(Resource, Default)]
struct WaterRuntime {
terrain: Option<Entity>,
absorb: Handle<WaterMaterial>,
surface: Handle<WaterMaterial>,
slots: Vec<WaterSlot>,
mesh: Option<Handle<Mesh>>,
warned_capacity: bool,
skirt: Vec<Entity>,
}
pub struct WaterPlugin;
impl Plugin for WaterPlugin {
fn build(&self, app: &mut App) {
ensure_gpu_core(app);
needs_depth_prepass(app);
embedded_asset!(app, "water.wesl");
embedded_asset!(app, "water_normals.wesl");
app.add_plugins(MaterialPlugin::<WaterMaterial>::default())
.init_resource::<WaterSettings>()
.init_resource::<WaterRuntime>()
.add_systems(Startup, create_normal_map)
.add_systems(
Update,
(ensure_water_pool, update_water_materials)
.chain()
.after(crate::plugin::ElementsSystems::Readback),
)
.add_systems(
PostUpdate,
(update_water_patches, update_water_skirt)
.after(crate::terrain::update_patches)
.before(TransformSystems::Propagate),
);
water_gpu::plugin(app);
}
}
fn create_normal_map(mut commands: Commands, mut images: ResMut<Assets<Image>>) {
let mut image = Image::new_uninit(
Extent3d {
width: NORMAL_TEXELS,
height: NORMAL_TEXELS,
depth_or_array_layers: 1,
},
TextureDimension::D2,
TextureFormat::Rgba8Unorm,
bevy::asset::RenderAssetUsages::RENDER_WORLD,
);
image.texture_descriptor.mip_level_count = NORMAL_MIPS;
image.texture_descriptor.usage =
TextureUsages::STORAGE_BINDING | TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_SRC;
image.sampler = ImageSampler::Descriptor(ImageSamplerDescriptor {
address_mode_u: ImageAddressMode::Repeat,
address_mode_v: ImageAddressMode::Repeat,
mag_filter: ImageFilterMode::Linear,
min_filter: ImageFilterMode::Linear,
mipmap_filter: ImageFilterMode::Linear,
anisotropy_clamp: 16,
..default()
});
commands.insert_resource(WaterState {
normals: images.add(image),
normals_ready: Arc::new(AtomicBool::new(false)),
});
}
fn pass_materials(
materials: &mut Assets<WaterMaterial>,
normals: &Handle<Image>,
) -> (Handle<WaterMaterial>, Handle<WaterMaterial>) {
let extension = WaterExtension {
params: WaterParams::default(),
cells: NO_TERRAIN,
flow: NO_TERRAIN,
normals: normals.clone(),
cloud_shadow: None,
waves: None,
};
let base = StandardMaterial {
base_color: Color::BLACK,
perceptual_roughness: WATER_ROUGHNESS,
reflectance: WATER_REFLECTANCE,
metallic: 0.0,
..default()
};
let absorb = materials.add(WaterMaterial {
base: StandardMaterial {
alpha_mode: AlphaMode::Multiply,
..base.clone()
},
extension: extension.clone(),
});
let surface = materials.add(WaterMaterial {
base: StandardMaterial {
alpha_mode: AlphaMode::Add,
depth_bias: 1.0,
..base
},
extension,
});
(absorb, surface)
}
fn ensure_water_pool(
mut commands: Commands,
state: Option<Res<TerrainState>>,
water: Option<Res<WaterState>>,
mut runtime: ResMut<WaterRuntime>,
mut materials: ResMut<Assets<WaterMaterial>>,
mut meshes: ResMut<Assets<Mesh>>,
) {
let (Some(state), Some(water)) = (state, water) else {
if runtime.terrain.take().is_some() {
runtime.slots.clear();
runtime.skirt.clear();
}
return;
};
if runtime.terrain == Some(state.entity) {
return;
}
let (absorb, surface) = pass_materials(&mut materials, &water.normals);
runtime.absorb = absorb.clone();
runtime.surface = surface.clone();
let mesh = runtime
.mesh
.get_or_insert_with(|| meshes.add(patch_mesh()))
.clone();
let hidden = Aabb::from_min_max(Vec3::ZERO, Vec3::ONE);
let mut slots = Vec::with_capacity(WATER_POOL_CAPACITY);
let mut skirt = Vec::with_capacity(8);
commands.entity(state.entity).with_children(|parent| {
for _ in 0..WATER_POOL_CAPACITY {
let mut spawn = |surface_pass: bool| {
parent
.spawn((
Name::new(if surface_pass {
"Water surface"
} else {
"Water absorb"
}),
WaterPatch {
surface: surface_pass,
},
Mesh3d(mesh.clone()),
MeshMaterial3d(if surface_pass {
surface.clone()
} else {
absorb.clone()
}),
Transform::default(),
MeshTag::new(0),
hidden,
NoAutoAabb,
Visibility::Hidden,
))
.id()
};
let absorb = spawn(false);
let surface = spawn(true);
slots.push(WaterSlot {
absorb,
surface,
patch: None,
});
}
for tile in 0..4 {
for (surface_pass, material) in [(false, &absorb), (true, &surface)] {
let id = parent
.spawn((
Name::new(if surface_pass {
"Water skirt surface"
} else {
"Water skirt absorb"
}),
WaterSkirt { tile },
Mesh3d(mesh.clone()),
MeshMaterial3d(material.clone()),
Transform::default(),
MeshTag::new(SKIRT_TAG),
hidden,
NoAutoAabb,
Visibility::Hidden,
))
.id();
skirt.push(id);
}
}
});
runtime.slots = slots;
runtime.skirt = skirt;
runtime.terrain = Some(state.entity);
}
pub(crate) fn skirt_rects(lo: Vec2, hi: Vec2, reach: f32) -> [(Vec2, Vec2); 4] {
let r = Vec2::splat(reach);
let span = hi - lo;
[
(lo - r, Vec2::new(reach, span.y + reach)),
(
Vec2::new(lo.x, lo.y - reach),
Vec2::new(span.x + reach, reach),
),
(Vec2::new(hi.x, lo.y), Vec2::new(reach, span.y + reach)),
(
Vec2::new(lo.x - reach, hi.y),
Vec2::new(span.x + reach, reach),
),
]
}
fn update_water_skirt(
terrain: Res<TerrainRuntime>,
state: Option<Res<TerrainState>>,
runtime: Res<WaterRuntime>,
units: Res<ElementsUnits>,
still: Option<Res<StillWater>>,
mut tiles: Query<(&WaterSkirt, &mut Transform, &mut Visibility, &mut Aabb)>,
) {
let (Some(state), Some(lod)) = (state, terrain.lod) else {
return;
};
if runtime.skirt.is_empty() {
return;
}
let half = state.grid.extent() * 0.5;
let rects = skirt_rects(
lod.sample_offset() - half,
lod.last_sample() - half,
WATER_SKIRT_REACH_M,
);
let level_m = still.map(|s| s.level / units.per_metre());
for entity in &runtime.skirt {
let Ok((tile, mut transform, mut visibility, mut aabb)) = tiles.get_mut(*entity) else {
continue;
};
let Some(level_m) = level_m else {
if *visibility != Visibility::Hidden {
*visibility = Visibility::Hidden;
}
continue;
};
let (corner, size) = rects[tile.tile];
let translation = units.point_to_world(Vec3::new(corner.x, 0.0, corner.y));
let scale = Vec3::new(units.world(size.x), 1.0, units.world(size.y));
if transform.translation != translation || transform.scale != scale {
transform.translation = translation;
transform.scale = scale;
}
let bounds = Aabb::from_min_max(
Vec3::new(0.0, units.world(level_m - WATER_AABB_PAD_M), 0.0),
Vec3::new(1.0, units.world(level_m + WATER_AABB_PAD_M), 1.0),
);
if *aabb != bounds {
*aabb = bounds;
}
if *visibility != Visibility::Inherited {
*visibility = Visibility::Inherited;
}
}
}
pub(crate) fn wet_patches(
selected: &[crate::terrain::lod::Patch],
depth: &crate::terrain::lod::BoundsTree,
boxes: &crate::terrain::lod::BoundsTree,
still_level_m: Option<f32>,
) -> Vec<crate::terrain::lod::Patch> {
selected
.iter()
.copied()
.filter(|p| {
depth.get(p.level, p.ix, p.iz).1 >= MIN_PATCH_DEPTH_M
|| still_level_m
.is_some_and(|l| l - boxes.get(p.level, p.ix, p.iz).0 >= MIN_PATCH_DEPTH_M)
})
.collect()
}
fn update_water_patches(
terrain: Res<TerrainRuntime>,
state: Option<Res<TerrainState>>,
view: Res<ElementsView>,
mut runtime: ResMut<WaterRuntime>,
mut patches: Query<(&mut Transform, &mut Visibility, &mut MeshTag, &mut Aabb)>,
units: Res<ElementsUnits>,
still: Option<Res<StillWater>>,
) {
let (Some(state), Some(lod)) = (state, terrain.lod) else {
return;
};
let (Some(boxes), Some(depth)) = (terrain.water_box.as_ref(), terrain.water_depth.as_ref())
else {
return;
};
if runtime.slots.is_empty() || view.entity.is_none() {
return;
}
let still_m = still.map(|s| s.level / units.per_metre());
let mut wanted = wet_patches(&terrain.selected, depth, boxes, still_m);
if wanted.len() > WATER_POOL_CAPACITY {
if !runtime.warned_capacity {
warn!(
"the water pool wants {} patches, more than its {WATER_POOL_CAPACITY}; the farthest are dropped",
wanted.len()
);
runtime.warned_capacity = true;
}
wanted.truncate(WATER_POOL_CAPACITY);
}
let current: Vec<_> = runtime.slots.iter().map(|s| s.patch).collect();
let (assigned, used) = assign_slots(wanted, ¤t);
let extent = state.grid.extent();
for (i, slot) in runtime.slots.iter_mut().enumerate() {
if used[i] {
continue;
}
if slot.patch.take().is_some() {
for entity in [slot.absorb, slot.surface] {
if let Ok((_, mut visibility, _, _)) = patches.get_mut(entity)
&& *visibility != Visibility::Hidden
{
*visibility = Visibility::Hidden;
}
}
}
}
for (patch, i) in assigned {
let slot = &mut runtime.slots[i];
slot.patch = Some(patch);
let (translation, scale) = patch_transform(&lod, extent, patch);
let (translation, scale) = (
units.point_to_world(translation),
Vec3::new(units.world(scale.x), scale.y, units.world(scale.z)),
);
let (low, high) = boxes.get(patch.level, patch.ix, patch.iz);
let high = still_m.map_or(high, |l| high.max(l));
let bounds = Aabb::from_min_max(
Vec3::new(0.0, units.world(low - WATER_AABB_PAD_M), 0.0),
Vec3::new(1.0, units.world(high + WATER_AABB_PAD_M), 1.0),
);
for entity in [slot.absorb, slot.surface] {
let Ok((mut transform, mut visibility, mut tag, mut aabb)) = patches.get_mut(entity)
else {
continue;
};
if transform.translation != translation || transform.scale != scale {
transform.translation = translation;
transform.scale = scale;
}
if tag.value != patch.level {
*tag = MeshTag::new(patch.level);
}
if *aabb != bounds {
*aabb = bounds;
}
if *visibility != Visibility::Inherited {
*visibility = Visibility::Inherited;
}
}
}
}
pub(crate) fn water_params(
settings: &WaterSettings,
terrain: TerrainParamsGpu,
lod: [Vec4; MAX_LEVELS],
top: Vec4,
shadows: (Vec4, Vec4),
wind_ms: f32,
units: f32,
) -> WaterParams {
let u10 = settings.wind_speed_ms.unwrap_or(wind_ms);
let silt = Color::srgb(
settings.silt_tint.x,
settings.silt_tint.y,
settings.silt_tint.z,
)
.to_linear();
WaterParams {
terrain,
lod,
top,
lod_centre: Vec4::ZERO,
shadow: shadows.0,
light: shadows.1,
absorption: settings.absorption.extend(settings.refractive_index),
silt: Vec3::new(silt.red, silt.green, silt.blue).extend(settings.turbidity_scattering),
deep: settings.deep_color.extend(settings.clear_scattering),
flow: Vec4::new(
settings.flow_cycle_s,
settings.tile_m,
settings.flow_distortion,
settings.ripple_scale * ripple_slope(u10) / NORMAL_RMS_SLOPE,
),
foam: Vec4::new(
settings.foam_speed_ms,
settings.foam_thickness_m,
settings.foam_albedo,
settings.foam_roughness,
),
waves: Vec4::ZERO,
still: Vec4::new(NO_STILL_LEVEL, 0.0, 0.0, 0.0),
units,
}
}
#[allow(clippy::too_many_arguments)]
fn update_water_materials(
state: Option<Res<TerrainState>>,
terrain: Res<TerrainRuntime>,
settings: Res<WaterSettings>,
shadows: Option<Res<CloudShadows>>,
wind: Option<Res<WindState>>,
waves: Option<Res<WaveState>>,
still: Option<Res<StillWater>>,
runtime: Res<WaterRuntime>,
units: Res<ElementsUnits>,
mut materials: ResMut<Assets<WaterMaterial>>,
) {
let (Some(state), Some(lod)) = (state, terrain.lod.as_ref()) else {
return;
};
let (shadow, light, map) = shadow_inputs(shadows.as_deref());
let (lod, top) = lod_uniform(lod);
let wind_ms = wind.map_or(MAP_WIND_MS, |w| (w.speed_10m * 2.0).round() * 0.5);
let mut params = water_params(
&settings,
TerrainParamsGpu::new(&state.grid, 1, state.is_live()),
lod,
top,
(shadow, light),
wind_ms,
units.per_metre(),
);
params.still.x = still.map_or(NO_STILL_LEVEL, |s| s.level / units.per_metre());
params.lod_centre = lod_centre_uniform(terrain.lod_centre);
let field = waves.as_deref().map(|w| w.field.clone());
if let Some(w) = waves.as_deref() {
params.waves = Vec4::new(w.texel_m[0], w.texels as f32, 1.0, 0.0);
}
for handle in [&runtime.absorb, &runtime.surface] {
let Some(material) = materials.get(handle) else {
continue;
};
let e = &material.extension;
if (e.params != params
|| e.cells != state.cells
|| e.flow != state.flow_cells
|| e.cloud_shadow != map
|| e.waves != field)
&& let Some(mut material) = materials.get_mut(handle)
{
material.extension.params = params;
material.extension.cells = state.cells.clone();
material.extension.flow = state.flow_cells.clone();
material.extension.cloud_shadow = map.clone();
material.extension.waves = field.clone();
}
}
}