bevy_exponential_height_fog 0.1.0

use bevy::asset::{AssetServer, Handle, load_embedded_asset};
use bevy::core_pipeline::FullscreenShader;
use bevy::prelude::*;
use bevy::render::render_resource::{
    BindGroupLayoutDescriptor, BindGroupLayoutEntries, CachedRenderPipelineId, ColorTargetState,
    ColorWrites, FragmentState, PipelineCache, RenderPipelineDescriptor, Sampler,
    SamplerBindingType, SamplerDescriptor, ShaderStages, SpecializedRenderPipeline,
    SpecializedRenderPipelines, TextureFormat, TextureSampleType,
    binding_types::{
        sampler, texture_2d, texture_depth_2d, texture_depth_2d_multisampled, uniform_buffer,
    },
};
use bevy::render::renderer::RenderDevice;
use bevy::render::view::{ExtractedView, Msaa, ViewTarget, ViewUniform};
use bevy::shader::Shader;
use bevy::shader::ShaderDefVal;

use crate::component::ExponentialHeightFogUniform;

#[derive(Resource)]
pub struct ExponentialHeightFogPipeline {
    pub sampler: Sampler,
    pub layout: BindGroupLayoutDescriptor,
    pub layout_msaa: BindGroupLayoutDescriptor,
    pub fullscreen_shader: FullscreenShader,
    pub fragment_shader: Handle<Shader>,
}

impl ExponentialHeightFogPipeline {
    pub fn new(
        render_device: &RenderDevice,
        fullscreen_shader: FullscreenShader,
        fragment_shader: Handle<Shader>,
    ) -> Self {
        let layout = BindGroupLayoutDescriptor::new(
            "exponential_height_fog_layout",
            &BindGroupLayoutEntries::sequential(
                ShaderStages::FRAGMENT,
                (
                    uniform_buffer::<ViewUniform>(true),
                    texture_depth_2d(),
                    texture_2d(TextureSampleType::Float { filterable: true }),
                    sampler(SamplerBindingType::Filtering),
                    uniform_buffer::<ExponentialHeightFogUniform>(true),
                ),
            ),
        );
        let layout_msaa = BindGroupLayoutDescriptor::new(
            "exponential_height_fog_layout_msaa",
            &BindGroupLayoutEntries::sequential(
                ShaderStages::FRAGMENT,
                (
                    uniform_buffer::<ViewUniform>(true),
                    texture_depth_2d_multisampled(),
                    texture_2d(TextureSampleType::Float { filterable: true }),
                    sampler(SamplerBindingType::Filtering),
                    uniform_buffer::<ExponentialHeightFogUniform>(true),
                ),
            ),
        );
        let sampler = render_device.create_sampler(&SamplerDescriptor::default());

        Self {
            sampler,
            layout,
            layout_msaa,
            fullscreen_shader,
            fragment_shader,
        }
    }
}

pub fn init_exponential_height_fog_pipeline(
    mut commands: Commands,
    render_device: Res<RenderDevice>,
    fullscreen_shader: Res<FullscreenShader>,
    asset_server: Res<AssetServer>,
) {
    let fragment_shader = load_embedded_asset!(asset_server.as_ref(), "exponential_height_fog.wgsl");
    commands.insert_resource(ExponentialHeightFogPipeline::new(
        &render_device,
        fullscreen_shader.clone(),
        fragment_shader,
    ));
}

#[derive(PartialEq, Eq, Hash, Clone, Copy)]
pub struct ExponentialHeightFogPipelineKey {
    pub hdr: bool,
    pub samples: u32,
}

impl SpecializedRenderPipeline for ExponentialHeightFogPipeline {
    type Key = ExponentialHeightFogPipelineKey;

    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
        let mut shader_defs = vec![];
        let layout = if key.samples > 1 {
            shader_defs.push(ShaderDefVal::from("MULTISAMPLED"));
            vec![self.layout_msaa.clone()]
        } else {
            vec![self.layout.clone()]
        };

        RenderPipelineDescriptor {
            label: Some("exponential_height_fog_pipeline".into()),
            layout,
            vertex: self.fullscreen_shader.to_vertex_state(),
            fragment: Some(FragmentState {
                shader: self.fragment_shader.clone(),
                shader_defs,
                targets: vec![Some(ColorTargetState {
                    format: if key.hdr {
                        ViewTarget::TEXTURE_FORMAT_HDR
                    } else {
                        TextureFormat::bevy_default()
                    },
                    blend: None,
                    write_mask: ColorWrites::ALL,
                })],
                ..default()
            }),
            ..default()
        }
    }
}

#[derive(Component)]
pub struct ExponentialHeightFogPipelineId(pub CachedRenderPipelineId);

pub fn prepare_exponential_height_fog_pipelines(
    mut commands: Commands,
    pipeline_cache: Res<PipelineCache>,
    mut pipelines: ResMut<SpecializedRenderPipelines<ExponentialHeightFogPipeline>>,
    pipeline: Res<ExponentialHeightFogPipeline>,
    views: Query<(Entity, &ExtractedView, &Msaa), With<ExponentialHeightFogUniform>>,
) {
    for (entity, view, msaa) in &views {
        let pipeline_id = pipelines.specialize(
            &pipeline_cache,
            &pipeline,
            ExponentialHeightFogPipelineKey {
                hdr: view.hdr,
                samples: msaa.samples(),
            },
        );
        commands
            .entity(entity)
            .insert(ExponentialHeightFogPipelineId(pipeline_id));
    }
}

pub fn configure_depth_texture_for_exponential_height_fog(
    mut cameras: Query<&mut Camera3d, With<crate::component::ExponentialHeightFog>>,
) {
    for mut camera in &mut cameras {
        let mut depth_texture_usages = bevy::render::render_resource::TextureUsages::from(
            camera.depth_texture_usages,
        );
        depth_texture_usages |= bevy::render::render_resource::TextureUsages::TEXTURE_BINDING;
        camera.depth_texture_usages = depth_texture_usages.into();
    }
}