bevy_flora 0.22.0

//! Seeing what the camera follows through the trees in front of it: [`SeeThroughTrees`].

use std::sync::RwLock;

use bevy::asset::{Asset, AssetId, Assets, Handle};
use bevy::ecs::prelude::*;
use bevy::math::{Vec3, Vec4};
use bevy::reflect::Reflect;
use bevy::render::storage::ShaderBuffer;
use bevy::shader::Shader;
use bevy::transform::components::GlobalTransform;

use crate::gpu::GrowthMaterial;
use crate::impostor::ImpostorMaterial;
use crate::shading::{FloraMaterial, FloraWindBuffer, TreeShakes};

/// Put on the game's camera to keep what it follows in view: the parts of trees between the
/// camera and [`target`](Self::target) dither out, in a cone around the line of sight that is
/// [`radius`](Self::radius) wide at the target and a third of that at the camera. Third-person
/// and top-down games use it to keep the player visible in a forest.
///
/// Only that camera sees through: shadow maps keep the whole tree (its shadow still falls on the
/// player), and so do other cameras. Every way a tree is drawn takes part (`ProceduralTree`,
/// `GpuGrowth`, pooled levels and impostors), and so do the camera's prepasses (a `DepthPrepass`
/// for TAA, SSAO or depth of field would otherwise keep the trees' depth where they were cut). One
/// camera at a time: the first found with the component. Update `target` as what it follows
/// moves; the camera's own movement is read from its `GlobalTransform`.
///
/// While no camera has it, no pipeline carries the fade: the plugin turns it on in every flora
/// material when one appears (the materials are re-prepared once) and off when it goes. With
/// plain `StandardMaterial`s (no flora shading, see [`FloraPlugin`](crate::FloraPlugin)) there is
/// no fade.
#[derive(Component, Clone, Copy, Debug, PartialEq, Reflect)]
#[reflect(Component)]
pub struct SeeThroughTrees {
    /// What the camera must keep in view, in world space: the player's chest, a unit, the point
    /// under the cursor.
    pub target: Vec3,
    /// How wide the clear view is at `target`, metres (a radius). The view closes over the last
    /// tenth of the way to the target, so nothing behind it is cut.
    pub radius: f32,
}

impl Default for SeeThroughTrees {
    fn default() -> Self {
        Self { target: Vec3::ZERO, radius: 1.5 }
    }
}

/// Where the fade's two vectors sit in the shared buffer ([`FloraWindBuffer`]): after the wind,
/// the season and the shakes. `common.wesl`'s `SEE_THROUGH` is this number.
pub(crate) const SEE_THROUGH_SLOT: usize = 7 + 2 * TreeShakes::MAX;

/// The fade as the shaders read it: the camera's position and the radius at the target, and the
/// target with 1 while a camera has [`SeeThroughTrees`] (all zero otherwise).
#[derive(Resource, Clone, Copy, Debug, Default, PartialEq)]
pub(crate) struct SeeThroughCapsule(pub [Vec4; 2]);

impl SeeThroughCapsule {
    fn of(camera: Option<(&SeeThroughTrees, &GlobalTransform)>) -> Self {
        match camera {
            Some((fade, at)) => Self([at.translation().extend(fade.radius.max(0.0)), fade.target.extend(1.0)]),
            None => Self::default(),
        }
    }
}

/// Write the fading camera's capsule into the shared buffer, in place, when it changed. Runs
/// after transform propagation, so the camera's position is the one this frame is drawn from:
/// the shaders cut only in the view drawn from exactly there.
pub(crate) fn sync_see_through(
    cameras: Query<(&SeeThroughTrees, &GlobalTransform)>,
    mut capsule: ResMut<SeeThroughCapsule>,
    buffer: Option<Res<FloraWindBuffer>>,
    buffers: Option<ResMut<Assets<ShaderBuffer>>>,
) {
    let now = SeeThroughCapsule::of(cameras.iter().next());
    // Look before borrowing mutably: a mutable borrow marks it changed.
    if *capsule == now {
        return;
    }
    *capsule = now;
    let (Some(buffer), Some(mut buffers)) = (buffer, buffers) else { return };
    let Some(mut shader_buffer) = buffers.get_mut(&buffer.0) else { return };
    if let Some(slice) = shader_buffer.cast_slice_mut::<[f32; 4]>()
        && slice.len() >= SEE_THROUGH_SLOT + 2
    {
        slice[SEE_THROUGH_SLOT] = now.0[0].to_array();
        slice[SEE_THROUGH_SLOT + 1] = now.0[1].to_array();
    }
}

/// Turn the fade on in every flora material while a camera has [`SeeThroughTrees`], and off when
/// none has: the flag is part of the materials' pipeline keys, so their shaders carry the
/// fade's discard (and the prepass its fragment stage) only then. Materials made while it is on
/// (a new look, a new pool) are turned on the frame after.
pub(crate) fn arm_see_through(
    cameras: Query<(), With<SeeThroughTrees>>,
    flora: Option<ResMut<Assets<FloraMaterial>>>,
    growth: Option<ResMut<Assets<GrowthMaterial>>>,
    impostors: Option<ResMut<Assets<ImpostorMaterial>>>,
) {
    let on = !cameras.is_empty();
    if let Some(mut assets) = flora {
        arm(&mut assets, on, |m| m.extension.see_through, |m, on| m.extension.see_through = on);
    }
    if let Some(mut assets) = growth {
        arm(&mut assets, on, |m| m.extension.see_through, |m, on| m.extension.see_through = on);
    }
    if let Some(mut assets) = impostors {
        arm(&mut assets, on, |m| m.see_through, |m, on| m.see_through = on);
    }
}

/// Set the flag of the materials in `assets` whose flag is not `on` (only those: a material
/// borrowed mutably is re-prepared).
fn arm<A: Asset>(assets: &mut Assets<A>, on: bool, flag: fn(&A) -> bool, set: fn(&mut A, bool)) {
    let stale: Vec<AssetId<A>> = assets.iter().filter(|(_, material)| flag(material) != on).map(|(id, _)| id).collect();
    for id in stale {
        if let Some(mut material) = assets.get_mut(id) {
            set(&mut material, on);
        }
    }
}

/// The prepass fragment stage of a flora pipeline while the fade is on (`see_through_prepass.wesl`).
/// Pipelines are specialized without access to the world, so the handle loaded when the shaders
/// are embedded is kept here.
static PREPASS_FRAGMENT: RwLock<Option<Handle<Shader>>> = RwLock::new(None);

/// Keep the handle of the prepass fragment shader for [`prepass_fragment`].
pub(crate) fn set_prepass_fragment(shader: Handle<Shader>) {
    if let Ok(mut slot) = PREPASS_FRAGMENT.write() {
        *slot = Some(shader);
    }
}

/// The prepass fragment shader of the fade, once the shaders are embedded.
pub(crate) fn prepass_fragment() -> Option<Handle<Shader>> {
    PREPASS_FRAGMENT.read().ok()?.clone()
}

#[cfg(test)]
mod tests {
    use bevy::app::App;
    use bevy::asset::{AssetApp, AssetPlugin};
    use bevy::mesh::Mesh;
    use bevy::pbr::StandardMaterial;

    use super::*;
    use crate::shading::{FloraMaterialSupport, FloraWind};
    use crate::{FloraPlugin, ProceduralTree, TreePreset};

    fn app() -> App {
        let mut app = App::new();
        app.add_plugins((bevy::MinimalPlugins, AssetPlugin::default()));
        app.init_asset::<Mesh>();
        app.init_asset::<StandardMaterial>();
        app.add_plugins(FloraPlugin);
        // No renderer: draw the trees with flora materials anyway.
        app.insert_resource(FloraMaterialSupport);
        app
    }

    fn flags(app: &App) -> Vec<bool> {
        app.world().resource::<Assets<FloraMaterial>>().iter().map(|(_, m)| m.extension.see_through).collect()
    }

    fn slots(app: &App) -> [[f32; 4]; 2] {
        let buffer = app.world().resource::<FloraWindBuffer>();
        let data = app.world().resource::<Assets<ShaderBuffer>>().get(&buffer.0).unwrap().cast_slice::<[f32; 4]>().unwrap();
        [data[SEE_THROUGH_SLOT], data[SEE_THROUGH_SLOT + 1]]
    }

    #[test]
    fn a_camera_that_sees_through_turns_the_fade_on_in_every_flora_material_until_it_goes() {
        let mut app = app();
        app.world_mut().spawn(ProceduralTree(TreePreset::Oak.params()));
        app.update();
        assert!(!flags(&app).is_empty() && flags(&app).iter().all(|on| !on), "off without a camera");
        assert_eq!(slots(&app), [[0.0; 4]; 2]);

        let camera = app
            .world_mut()
            .spawn((
                SeeThroughTrees { target: Vec3::new(4.0, 5.0, 6.0), radius: 2.0 },
                GlobalTransform::from_translation(Vec3::new(1.0, 2.0, 3.0)),
            ))
            .id();
        app.update();
        assert!(flags(&app).iter().all(|on| *on));
        assert_eq!(slots(&app), [[1.0, 2.0, 3.0, 2.0], [4.0, 5.0, 6.0, 1.0]]);

        // A tree of another look planted meanwhile gets materials of its own, turned on too; a
        // change of wind rewrites the whole buffer and keeps the fade.
        app.world_mut().spawn(ProceduralTree(TreePreset::Birch.params()));
        app.world_mut().resource_mut::<FloraWind>().strength = 0.9;
        app.update();
        app.update();
        assert!(flags(&app).len() > 2 && flags(&app).iter().all(|on| *on));
        assert_eq!(slots(&app)[1], [4.0, 5.0, 6.0, 1.0]);

        app.world_mut().entity_mut(camera).remove::<SeeThroughTrees>();
        app.update();
        assert!(flags(&app).iter().all(|on| !on), "off again once no camera sees through");
        assert_eq!(slots(&app), [[0.0; 4]; 2]);
    }
}