rusting_engine 2.7.0

//! A shader pack's `sky_lut` program: the procedural sky drawn every frame
//! into a small mipmapped equirectangular image with the pack's hook, then
//! bound as the environment map in place of the [`SkyBake`](super::sky).

use std::sync::Arc;

use vulkano::buffer::Subbuffer;
use vulkano::command_buffer::{
    AutoCommandBufferBuilder, PrimaryAutoCommandBuffer, RenderPassBeginInfo,
    SubpassBeginInfo, SubpassContents,
};
use vulkano::descriptor_set::allocator::StandardDescriptorSetAllocator;
use vulkano::descriptor_set::{DescriptorSet, WriteDescriptorSet};
use vulkano::device::Queue;
use vulkano::format::Format;
use vulkano::image::sampler::{
    Filter, Sampler, SamplerAddressMode, SamplerCreateInfo, SamplerMipmapMode,
    LOD_CLAMP_NONE,
};
use vulkano::image::view::{ImageView, ImageViewCreateInfo};
use vulkano::image::{
    Image, ImageCreateInfo, ImageSubresourceRange, ImageUsage,
};
use vulkano::memory::allocator::{
    AllocationCreateInfo, MemoryTypeFilter, StandardMemoryAllocator,
};
use vulkano::pipeline::graphics::color_blend::{
    ColorBlendAttachmentState, ColorBlendState,
};
use vulkano::pipeline::graphics::input_assembly::InputAssemblyState;
use vulkano::pipeline::graphics::multisample::MultisampleState;
use vulkano::pipeline::graphics::rasterization::RasterizationState;
use vulkano::pipeline::graphics::subpass::PipelineSubpassType;
use vulkano::pipeline::graphics::viewport::{Viewport, ViewportState};
use vulkano::pipeline::graphics::GraphicsPipelineCreateInfo;
use vulkano::pipeline::layout::PipelineDescriptorSetLayoutCreateInfo;
use vulkano::pipeline::{
    GraphicsPipeline, Pipeline, PipelineBindPoint, PipelineLayout,
    PipelineShaderStageCreateInfo,
};
use vulkano::render_pass::{Framebuffer, FramebufferCreateInfo, Subpass};

use super::composite::{FrameBlock, PackHooks, SCENE_FRAME_BINDING};
use super::scene_renderer::{
    compile_glsl, hooked_glsl, name_object, record_mip_chain,
    tonemap_vertex_shader, SceneRenderError,
};
use super::sky::SkyParams;

fn error(error: impl std::fmt::Display) -> SceneRenderError {
    SceneRenderError(error.to_string())
}

/// Size of the LUT; `sky_lut.glsl` divides by it too.
const LUT_EXTENT: [u32; 2] = [256, 128];

/// The `sky_lut` program's fragment shader with `source`, a
/// `#rusting sky_lut 1` shader, as its hook.
pub(super) fn sky_lut_glsl(source: &str) -> Result<String, String> {
    let frag = include_str!("../shaders/sky_lut.glsl").replacen(
        "#include \"frame.glsl\"",
        include_str!("../shaders/frame.glsl"),
        1,
    );
    hooked_glsl(&frag, &[("sky_lut", source, ("sky_lut_hook.glsl", ""))])
}

/// The pack's sky LUT: its pipeline and image.
pub(super) struct SkyLut {
    pipeline: Arc<GraphicsPipeline>,
    framebuffer: Arc<Framebuffer>,
    /// Every mip, for lighting.
    pub(super) view: Arc<ImageView>,
    pub(super) sampler: Arc<Sampler>,
    /// Recorded at least once, so its image can be kept while the sky is
    /// hidden.
    pub(super) drawn: bool,
}

/// The LUT of one pack revision; `None` inside when it failed to build, so
/// it is not tried again and the sky bake lights the scene.
pub(super) type SkyLutCache = Option<(u64, Option<SkyLut>)>;

/// The LUT for `hooks`, built on first use; `None` when the pack has no
/// `sky_lut` program or it does not build.
pub(super) fn sky_lut<'a>(
    cache: &'a mut SkyLutCache,
    queue: &Arc<Queue>,
    memory: &Arc<StandardMemoryAllocator>,
    hooks: &PackHooks,
) -> Option<&'a mut SkyLut> {
    if hooks.sky_lut.is_empty() {
        return None;
    }
    if cache
        .as_ref()
        .is_none_or(|(revision, _)| *revision != hooks.revision)
    {
        let lut = SkyLut::new(queue, memory, &hooks.sky_lut)
            .inspect_err(|e| {
                eprintln!("sky_lut skipped, using the sky bake:\n{}", e.0)
            })
            .ok();
        *cache = Some((hooks.revision, lut));
    }
    cache.as_mut().and_then(|(_, lut)| lut.as_mut())
}

impl SkyLut {
    fn new(
        queue: &Arc<Queue>,
        memory: &Arc<StandardMemoryAllocator>,
        source: &str,
    ) -> Result<Self, SceneRenderError> {
        let device = queue.device().clone();
        let fragment = compile_glsl(
            &device,
            "fragment",
            &sky_lut_glsl(source).map_err(error)?,
        )
        .map_err(error)?
        .entry_point("main")
        .ok_or_else(|| error("sky_lut entry point is missing"))?;
        let vertex = tonemap_vertex_shader::load(device.clone())
            .map_err(error)?
            .entry_point("main")
            .ok_or_else(|| error("sky_lut vertex entry point is missing"))?;
        let stages = [
            PipelineShaderStageCreateInfo::new(vertex),
            PipelineShaderStageCreateInfo::new(fragment),
        ];
        let layout = PipelineLayout::new(
            device.clone(),
            PipelineDescriptorSetLayoutCreateInfo::from_stages(&stages)
                .into_pipeline_layout_create_info(device.clone())
                .map_err(error)?,
        )
        .map_err(error)?;
        let render_pass = vulkano::single_pass_renderpass!(
            device.clone(),
            attachments: {
                color: {
                    format: Format::R16G16B16A16_SFLOAT,
                    samples: 1,
                    load_op: DontCare,
                    store_op: Store,
                }
            },
            pass: {
                color: [color],
                depth_stencil: {}
            }
        )
        .map_err(error)?;
        let subpass = Subpass::from(render_pass.clone(), 0)
            .ok_or_else(|| error("sky_lut subpass is missing"))?;
        let pipeline = GraphicsPipeline::new(
            device.clone(),
            None,
            GraphicsPipelineCreateInfo {
                stages: stages.into_iter().collect(),
                vertex_input_state: Some(Default::default()),
                input_assembly_state: Some(InputAssemblyState::default()),
                viewport_state: Some(ViewportState {
                    viewports: [Viewport {
                        offset: [0.0, 0.0],
                        extent: LUT_EXTENT.map(|side| side as f32),
                        depth_range: 0.0..=1.0,
                    }]
                    .into_iter()
                    .collect(),
                    ..Default::default()
                }),
                rasterization_state: Some(RasterizationState::default()),
                multisample_state: Some(MultisampleState::default()),
                color_blend_state: Some(
                    ColorBlendState::with_attachment_states(
                        1,
                        ColorBlendAttachmentState::default(),
                    ),
                ),
                subpass: Some(PipelineSubpassType::BeginRenderPass(subpass)),
                ..GraphicsPipelineCreateInfo::layout(layout)
            },
        )
        .map_err(error)?;
        let [width, height] = LUT_EXTENT;
        let image = Image::new(
            memory.clone(),
            ImageCreateInfo {
                format: Format::R16G16B16A16_SFLOAT,
                extent: [width, height, 1],
                mip_levels: width.max(height).ilog2() + 1,
                // Transfer source and destination for the mip chain.
                usage: ImageUsage::COLOR_ATTACHMENT
                    | ImageUsage::SAMPLED
                    | ImageUsage::TRANSFER_SRC
                    | ImageUsage::TRANSFER_DST,
                ..Default::default()
            },
            AllocationCreateInfo {
                memory_type_filter: MemoryTypeFilter::PREFER_DEVICE,
                ..Default::default()
            },
        )
        .map_err(error)?;
        name_object(&*image, "Pack sky LUT");
        let top = ImageView::new(
            image.clone(),
            ImageViewCreateInfo {
                subresource_range: ImageSubresourceRange {
                    mip_levels: 0..1,
                    ..image.subresource_range()
                },
                ..ImageViewCreateInfo::from_image(&image)
            },
        )
        .map_err(error)?;
        let framebuffer = Framebuffer::new(
            render_pass,
            FramebufferCreateInfo {
                attachments: vec![top],
                ..Default::default()
            },
        )
        .map_err(error)?;
        // The same sampling as the sky bake: around wraps, up and down clamp.
        let sampler = Sampler::new(
            device,
            SamplerCreateInfo {
                mag_filter: Filter::Linear,
                min_filter: Filter::Linear,
                mipmap_mode: SamplerMipmapMode::Linear,
                address_mode: [
                    SamplerAddressMode::Repeat,
                    SamplerAddressMode::ClampToEdge,
                    SamplerAddressMode::ClampToEdge,
                ],
                lod: 0.0..=LOD_CLAMP_NONE,
                ..Default::default()
            },
        )
        .map_err(error)?;
        Ok(Self {
            pipeline,
            framebuffer,
            view: ImageView::new_default(image).map_err(error)?,
            sampler,
            drawn: false,
        })
    }

    /// Records the LUT for `params` with `frame` bound, and its mip chain.
    pub(super) fn record(
        &mut self,
        commands: &mut AutoCommandBufferBuilder<PrimaryAutoCommandBuffer>,
        descriptors: &Arc<StandardDescriptorSetAllocator>,
        params: SkyParams,
        frame: Subbuffer<FrameBlock>,
    ) -> Result<(), SceneRenderError> {
        let layout = self.pipeline.layout().clone();
        let set = DescriptorSet::new(
            descriptors.clone(),
            layout.set_layouts()[0].clone(),
            [WriteDescriptorSet::buffer(SCENE_FRAME_BINDING, frame)],
            [],
        )
        .map_err(error)?;
        commands
            .begin_render_pass(
                RenderPassBeginInfo {
                    clear_values: vec![None],
                    ..RenderPassBeginInfo::framebuffer(self.framebuffer.clone())
                },
                SubpassBeginInfo {
                    contents: SubpassContents::Inline,
                    ..Default::default()
                },
            )
            .map_err(error)?
            .bind_pipeline_graphics(self.pipeline.clone())
            .map_err(error)?
            .bind_descriptor_sets(
                PipelineBindPoint::Graphics,
                layout.clone(),
                0,
                set,
            )
            .map_err(error)?
            .push_constants(layout, 0, params)
            .map_err(error)?;
        // Safety: one full-screen triangle, no vertex buffers.
        unsafe { commands.draw(3, 1, 0, 0) }.map_err(error)?;
        commands
            .end_render_pass(Default::default())
            .map_err(error)?;
        self.drawn = true;
        record_mip_chain(commands, self.view.image())
    }
}