mod gpu;
pub(crate) mod model;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use bevy::asset::{RenderAssetUsages, embedded_asset};
use bevy::image::{ImageAddressMode, ImageFilterMode, ImageSampler, ImageSamplerDescriptor};
use bevy::prelude::*;
use bevy::render::render_resource::{
Extent3d, ShaderType, TextureDimension, TextureFormat, TextureUsages,
};
use bevy::render::storage::ShaderBuffer;
use crate::celestial::SunState;
use crate::celestial::{CelestialFrame, WorldClock};
use crate::climate::{SurfaceWeathering, Weather};
use crate::clouds::CloudShadows;
use crate::cover::{CoverState, SnowSettings};
use crate::dust::DustSettings;
use crate::gpu::grid::zeroed_buffer;
use crate::gpu::readback::{ReadbackTick, Readbacks, decode};
use crate::plugin::{
ElementsQuality, ElementsSystems, ElementsUnits, ElementsView, ensure_gpu_core,
};
use crate::terrain::{TerrainGpu, TerrainParamsGpu, TerrainState};
use model::{ImprintGpu, MAX_IMPRINTS, MAX_SWEEP_M, window_origin};
#[derive(Component, Clone, Copy, Debug, PartialEq, Reflect)]
#[require(Transform)]
pub struct Imprint {
pub half_extents: Vec2,
pub load_n: f32,
}
impl Imprint {
pub fn new(half_extents: Vec2, load_n: f32) -> Self {
Self {
half_extents,
load_n,
}
}
}
#[derive(Resource, Clone)]
pub struct TrailGround {
height: Arc<dyn Fn(Vec2) -> f32 + Send + Sync>,
sand: bool,
}
impl TrailGround {
pub fn new(height: impl Fn(Vec2) -> f32 + Send + Sync + 'static) -> Self {
Self {
height: Arc::new(height),
sand: false,
}
}
pub fn flat(y: f32) -> Self {
Self::new(move |_| y)
}
pub fn with_sand(mut self) -> Self {
self.sand = true;
self
}
pub fn is_sand(&self) -> bool {
self.sand
}
pub fn height(&self, xz: Vec2) -> f32 {
(self.height)(xz)
}
}
impl std::fmt::Debug for TrailGround {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TrailGround")
.field("sand", &self.sand)
.finish_non_exhaustive()
}
}
#[derive(Resource, Clone, Debug, PartialEq, Reflect)]
#[cfg_attr(
feature = "serde",
derive(serde::Serialize, serde::Deserialize),
serde(default)
)]
pub struct TrailSettings {
pub enabled: bool,
pub texel_m: Option<f32>,
pub min_depth_m: f32,
}
impl Default for TrailSettings {
fn default() -> Self {
Self {
enabled: true,
texel_m: None,
min_depth_m: 0.0,
}
}
}
impl TrailSettings {
fn window(&self, quality: ElementsQuality) -> (u32, f32) {
let (size, texel) = quality.trail_window();
let texel = self
.texel_m
.filter(|t| t.is_finite() && *t > 0.0)
.unwrap_or(texel);
(size, texel)
}
}
#[derive(Resource, Clone, Debug)]
pub struct TrailState {
pub field: Handle<Image>,
pub texel_m: f32,
pub texels: u32,
pub centre: Vec2,
pub(crate) direct_share: f32,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, ShaderType)]
pub(crate) struct TrailParams {
terrain: TerrainParamsGpu,
origin: IVec2,
prev_origin: IVec2,
counts: UVec4,
sim: Vec4,
ground: Vec4,
mode: UVec4,
press: Vec4,
}
#[derive(Resource, Clone)]
pub(crate) struct TrailGpu {
pub state_a: Handle<ShaderBuffer>,
pub state_b: Handle<ShaderBuffer>,
pub stamp: Handle<ShaderBuffer>,
pub blocks: Handle<ShaderBuffer>,
pub cells: Handle<ShaderBuffer>,
pub cover: Handle<ShaderBuffer>,
pub readback: Handle<ShaderBuffer>,
pub field: Handle<Image>,
pub params: TrailParams,
pub imprints: Vec<ImprintGpu>,
pub run: bool,
pub read: bool,
pub executed: Arc<AtomicU64>,
pub ready: Arc<AtomicBool>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Source {
Terrain(AssetId<ShaderBuffer>),
Game,
}
#[derive(Resource, Default)]
struct TrailRuntime {
built: Option<(u32, f32)>,
source: Option<Source>,
last: HashMap<Entity, Vec2>,
}
pub struct TrailsPlugin;
impl Plugin for TrailsPlugin {
fn build(&self, app: &mut App) {
ensure_gpu_core(app);
embedded_asset!(app, "trails.wesl");
app.init_resource::<TrailSettings>()
.init_resource::<SnowSettings>()
.init_resource::<TrailRuntime>()
.add_systems(
Update,
(
(ensure_trails, run_trails)
.chain()
.in_set(ElementsSystems::Simulation),
request_readback.in_set(ElementsSystems::Readback),
),
);
gpu::plugin(app);
}
}
fn field_image(images: &mut Assets<Image>, size: u32) -> Handle<Image> {
let mut image = Image::new_fill(
Extent3d {
width: size,
height: size,
depth_or_array_layers: 1,
},
TextureDimension::D2,
&[0u8; 8],
TextureFormat::Rgba16Float,
RenderAssetUsages::RENDER_WORLD,
);
image.texture_descriptor.usage = TextureUsages::STORAGE_BINDING
| TextureUsages::TEXTURE_BINDING
| TextureUsages::COPY_DST
| 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,
..default()
});
images.add(image)
}
pub(crate) fn readback_cell(size: u32) -> u32 {
(size / 512).max(2)
}
#[allow(clippy::too_many_arguments)]
fn ensure_trails(
mut commands: Commands,
settings: Res<TrailSettings>,
quality: Res<ElementsQuality>,
game: Option<Res<TrailGround>>,
terrain: Option<Res<TerrainState>>,
cover: Option<Res<CoverState>>,
gpu: Option<Res<TrailGpu>>,
mut runtime: ResMut<TrailRuntime>,
mut buffers: ResMut<Assets<ShaderBuffer>>,
mut images: ResMut<Assets<Image>>,
) {
let maps = match (&game, &terrain, &cover) {
(Some(_), _, _) => Some((Source::Game, None)),
(None, Some(t), Some(c)) => Some((
Source::Terrain(t.water.id()),
Some((t.cells.clone(), c.cover_cells.clone())),
)),
_ => None,
};
let (Some((source, maps)), true) = (maps, settings.enabled) else {
if gpu.is_some() {
commands.remove_resource::<TrailGpu>();
commands.remove_resource::<TrailState>();
*runtime = TrailRuntime::default();
}
return;
};
let window = settings.window(*quality);
if gpu.is_some() && runtime.built == Some(window) && runtime.source == Some(source) {
return;
}
let (size, texel_m) = window;
let texels = u64::from(size) * u64::from(size);
let cells = u64::from(size / readback_cell(size));
let field = field_image(&mut images, size);
let (cells_buffer, cover_buffer) = maps.unwrap_or_else(|| {
(
zeroed_buffer(&mut buffers, "elements_trails_no_terrain", 16),
zeroed_buffer(&mut buffers, "elements_trails_no_cover", 16),
)
});
*runtime = TrailRuntime {
built: Some(window),
source: Some(source),
..TrailRuntime::default()
};
commands.insert_resource(TrailState {
field: field.clone(),
texel_m,
texels: size,
centre: Vec2::ZERO,
direct_share: 0.0,
});
commands.insert_resource(TrailGpu {
state_a: zeroed_buffer(&mut buffers, "elements_trails_state_a", texels * 8),
state_b: zeroed_buffer(&mut buffers, "elements_trails_state_b", texels * 8),
stamp: zeroed_buffer(&mut buffers, "elements_trails_stamp", texels * 8),
blocks: zeroed_buffer(&mut buffers, "elements_trails_blocks", texels / 64 * 4),
cells: cells_buffer,
cover: cover_buffer,
readback: zeroed_buffer(
&mut buffers,
"elements_trails_readback",
(4 + cells * cells) * 4,
),
field,
params: TrailParams::default(),
imprints: vec![ImprintGpu::default()],
run: false,
read: false,
executed: Arc::new(AtomicU64::new(0)),
ready: Arc::new(AtomicBool::new(false)),
});
}
fn yaw_of(transform: &GlobalTransform) -> f32 {
let x = transform.right();
x.z.atan2(x.x)
}
fn direct_share(sun: &SunState, frame: &CelestialFrame, shadows: Option<&CloudShadows>) -> f32 {
let beam = shadows
.map_or(1.0, |s| s.mean_transmittance)
.clamp(0.0, 1.0);
if sun.altitude_deg > 0.0 {
let t = ((sun.altitude_deg + 2.0) / 12.0).clamp(0.0, 1.0);
let diffuse = 1.0 + (0.15 - 1.0) * t * t * (3.0 - 2.0 * t);
(1.0 - diffuse) * beam
} else if sun.moon_altitude_deg > 0.0 && frame.moon_lux > 0.0 {
beam
} else {
0.0
}
}
fn game_ground(ground: &TrailGround, weathering: Option<&SurfaceWeathering>) -> Vec4 {
let (snow, density, wetness) = weathering.map_or((0.0, 100.0, 0.0), |w| {
(w.snow_depth_m.max(0.0), w.snow_density(), w.wetness)
});
Vec4::new(
snow,
density,
wetness.clamp(0.0, 1.0),
if ground.sand { 1.0 } else { 0.0 },
)
}
type GroundSources<'w> = (
Option<Res<'w, TrailGround>>,
Option<Res<'w, SurfaceWeathering>>,
Option<Res<'w, TerrainState>>,
Option<Res<'w, TerrainGpu>>,
);
#[allow(clippy::too_many_arguments)]
fn run_trails(
view: Res<ElementsView>,
weather: Res<Weather>,
clock: Res<WorldClock>,
time: Res<Time<Virtual>>,
(sun, frame, shadows): (
Res<SunState>,
Res<CelestialFrame>,
Option<Res<CloudShadows>>,
),
snow: Res<SnowSettings>,
tick: Res<ReadbackTick>,
dust: Option<Res<DustSettings>>,
(game, weathering, terrain, terrain_gpu): GroundSources<'_>,
settings: Res<TrailSettings>,
gpu: Option<ResMut<TrailGpu>>,
state: Option<ResMut<TrailState>>,
mut runtime: ResMut<TrailRuntime>,
imprints: Query<(Entity, &Imprint, Ref<GlobalTransform>)>,
units: Res<ElementsUnits>,
) {
let (Some(mut gpu), Some(mut state)) = (gpu, state) else {
return;
};
let Some((size, texel_m)) = runtime.built else {
return;
};
let game = game.as_deref();
let (terrain_params, ground, live) = match (game, terrain.as_deref(), terrain_gpu) {
(Some(g), _, _) => (
TerrainParamsGpu::default(),
game_ground(g, weathering.as_deref()),
true,
),
(None, Some(t), Some(t_gpu)) => (
TerrainParamsGpu::new(&t.grid, 1, true),
Vec4::ZERO,
t.is_live() && t_gpu.generated.load(Ordering::Relaxed),
),
_ => return,
};
state.direct_share = direct_share(&sun, &frame, shadows.as_deref());
let camera = Vec2::new(view.position.x, view.position.z);
let origin = window_origin(camera, texel_m, size);
state.centre = (camera / texel_m).floor() * texel_m;
let mut seen = HashMap::with_capacity(runtime.last.len());
let mut keyed: Vec<(f32, Entity, ImprintGpu)> = Vec::new();
for (entity, imprint, transform) in &imprints {
if transform.is_added() {
continue;
}
let transform_m = units.to_metres(&transform);
let p = transform_m.translation();
let xz = Vec2::new(p.x, p.z);
let from = runtime
.last
.get(&entity)
.copied()
.filter(|from| from.distance(xz) <= MAX_SWEEP_M)
.unwrap_or(xz);
seen.insert(entity, xz);
let sole_y = match game {
Some(g) => {
let world = transform.translation();
let ground_y = g.height(Vec2::new(world.x, world.z)) / units.per_metre();
p.y - ground_y - ground.x
}
None => p.y,
};
if let Some(upload) = ImprintGpu::new(
imprint.half_extents,
yaw_of(&transform_m),
imprint.load_n,
from,
xz,
sole_y,
texel_m,
) {
keyed.push((xz.distance_squared(camera), entity, upload));
}
}
runtime.last = seen;
if keyed.len() > MAX_IMPRINTS {
keyed.select_nth_unstable_by(MAX_IMPRINTS, |a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
keyed.truncate(MAX_IMPRINTS);
}
let tile = keyed
.iter()
.map(|(_, _, i)| i.tile.w as u32)
.max()
.unwrap_or(0)
.div_ceil(16)
* 16;
let u = weather.wind_speed_ms;
let snow_d = if snow.drift {
snow.drift_coefficient * (u - snow.drift_threshold_ms).max(0.0).powi(3)
/ snow.new_density.max(1.0)
} else {
0.0
};
let sand_d = dust.as_deref().filter(|d| d.enabled).map_or(0.0, |d| {
d.flux_coefficient * u * u * (u - d.threshold_ms).max(0.0)
});
let dt = if clock.paused {
0.0
} else {
time.delta_secs() * clock.time_scale.max(0.0)
};
gpu.run = gpu.ready.load(Ordering::Relaxed) && live;
gpu.read = gpu.run && tick.is_due();
gpu.params = TrailParams {
terrain: terrain_params,
origin,
prev_origin: origin,
counts: UVec4::new(size, keyed.len() as u32, 0, tile),
sim: Vec4::new(texel_m, dt, snow_d, sand_d),
ground,
mode: UVec4::new(u32::from(game.is_some()), readback_cell(size), 0, 0),
press: Vec4::new(settings.min_depth_m.max(0.0), 0.0, 0.0, 0.0),
};
gpu.imprints = if keyed.is_empty() {
vec![ImprintGpu::default()]
} else {
keyed.into_iter().map(|(_, _, i)| i).collect()
};
}
fn request_readback(
gpu: Option<Res<TrailGpu>>,
state: Option<Res<TrailState>>,
mut readbacks: Readbacks,
) {
let (Some(gpu), Some(state)) = (gpu, state) else {
return;
};
if !gpu.read {
return;
}
let (texel_m, texels) = (state.texel_m, state.texels);
readbacks.request("trails", gpu.readback.clone(), move |bytes, snapshot| {
if let Some(values) = decode::<f32>(bytes)
&& let Some(trails) = TrailSnapshot::decode(&values, texel_m, texels)
{
snapshot.trails = Some(trails);
}
});
}
#[derive(Clone, Debug)]
pub(crate) struct TrailSnapshot {
origin: IVec2,
texel_m: f32,
texels: u32,
cell: u32,
deepest: Arc<[f32]>,
}
impl TrailSnapshot {
pub(crate) fn decode(values: &[f32], texel_m: f32, texels: u32) -> Option<Self> {
let (header, rest) = values.split_at_checked(4)?;
let cell = readback_cell(texels);
let cells = (texels / cell) as usize;
if header[2] != cells as f32 || header[3] != cell as f32 || rest.len() < cells * cells {
return None;
}
Some(Self {
origin: IVec2::new(header[0] as i32, header[1] as i32),
texel_m,
texels,
cell,
deepest: rest[..cells * cells].into(),
})
}
pub(crate) fn depth_at(&self, xz: Vec2) -> Option<f32> {
let w = (xz / self.texel_m).floor().as_ivec2();
let r = w - self.origin;
let n = self.texels as i32;
if r.min_element() < 0 || r.max_element() >= n {
return None;
}
let slot = w.rem_euclid(IVec2::splat(n));
let c = slot / self.cell as i32;
let cells = (self.texels / self.cell) as i32;
let h = self.deepest.get((c.y * cells + c.x) as usize)?;
Some((-h).max(0.0))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn sample() -> Vec<f32> {
let texels = 16u32;
let cell = readback_cell(texels);
let cells = (texels / cell) as i32;
let origin = IVec2::new(-5, 7);
let mut values = vec![origin.x as f32, origin.y as f32, cells as f32, cell as f32];
let mut deepest = vec![0.0f32; (cells * cells) as usize];
let slot = IVec2::new(-3, 9).rem_euclid(IVec2::splat(texels as i32));
let c = slot / cell as i32;
deepest[(c.y * cells + c.x) as usize] = -0.12;
values.extend(deepest);
values
}
#[test]
fn the_readback_places_a_print_where_the_window_held_it() {
let texel = 0.05;
let snapshot = TrailSnapshot::decode(&sample(), texel, 16).expect("a matching buffer");
let at = |w: IVec2| (w.as_vec2() + 0.5) * texel;
assert_eq!(snapshot.depth_at(at(IVec2::new(-3, 9))), Some(0.12));
assert_eq!(snapshot.depth_at(at(IVec2::new(-4, 8))), Some(0.12));
assert_eq!(snapshot.depth_at(at(IVec2::new(5, 15))), Some(0.0));
assert_eq!(snapshot.depth_at(at(IVec2::new(-6, 9))), None);
assert_eq!(snapshot.depth_at(at(IVec2::new(11, 9))), None);
}
#[test]
fn a_readback_of_another_window_is_refused() {
assert!(TrailSnapshot::decode(&sample(), 0.05, 32).is_none());
assert!(TrailSnapshot::decode(&sample()[..10], 0.05, 16).is_none());
}
}