vulkan_lib/vulkan-rs/src/shadermodule.rs

use crate::prelude::*;

use anyhow::{Context, Result};

use std::fs::File;
use std::io::Read;
use std::sync::Arc;

#[allow(clippy::cast_ptr_alignment)]
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ShaderType {
    None,
    Vertex,
    Fragment,
    Geometry,
    TesselationControl,
    TesselationEvaluation,
    Compute,
    RayGeneration,
    ClosestHit,
    Miss,
    AnyHit,
    Intersection,
}

impl From<u8> for ShaderType {
    fn from(value: u8) -> Self {
        match value {
            0 => Self::None,
            1 => Self::Vertex,
            2 => Self::Fragment,
            3 => Self::Geometry,
            4 => Self::TesselationControl,
            5 => Self::TesselationEvaluation,
            6 => Self::Compute,
            7 => Self::RayGeneration,
            8 => Self::ClosestHit,
            9 => Self::Miss,
            10 => Self::AnyHit,
            11 => Self::Intersection,

            _ => panic!("can't convert ShaderType from {}", value),
        }
    }
}

impl Default for ShaderType {
    fn default() -> Self {
        ShaderType::None
    }
}

pub trait VertexInputDescription: ReprC {
    fn bindings() -> Vec<VkVertexInputBindingDescription>;
    fn attributes() -> Vec<VkVertexInputAttributeDescription>;
}

pub trait PipelineStageInfo {
    fn pipeline_stage_info(&self) -> VkPipelineShaderStageCreateInfo;
}

macro_rules! impl_pipeline_stage_info {
    ($func:ident, $type:ident) => {
        impl PipelineStageInfo for ShaderModule<{ ShaderType::$type as u8 }> {
            fn pipeline_stage_info(&self) -> VkPipelineShaderStageCreateInfo {
                VkPipelineShaderStageCreateInfo::$func(self.shader_module)
            }
        }
    };
}

#[derive(Debug)]
pub struct ShaderModule<const TYPE: u8> {
    device: Arc<Device>,
    shader_module: VkShaderModule,
}

impl<const TYPE: u8> ShaderModule<TYPE> {
    pub fn new(device: Arc<Device>, path: &str) -> Result<Arc<ShaderModule<TYPE>>> {
        let code = Self::shader_code(path)?;

        Self::from_slice(device, code.as_slice())
    }

    pub fn from_slice(device: Arc<Device>, code: &[u8]) -> Result<Arc<ShaderModule<TYPE>>> {
        let shader_module_ci =
            VkShaderModuleCreateInfo::new(VK_SHADER_MODULE_CREATE_NULL_BIT, code);

        let shader_module = device.create_shader_module(&shader_module_ci)?;

        Ok(Arc::new(ShaderModule {
            device,
            shader_module,
        }))
    }

    fn shader_code(path: &str) -> Result<Vec<u8>> {
        let mut file = File::open(path).with_context({
            let path = path.to_string();
            || path
        })?;

        let mut code: Vec<u8> = Vec::new();

        file.read_to_end(&mut code)?;

        Ok(code)
    }
}

impl_pipeline_stage_info!(vertex, Vertex);
impl_pipeline_stage_info!(geometry, Geometry);
impl_pipeline_stage_info!(tesselation_control, TesselationControl);
impl_pipeline_stage_info!(tesselation_evaluation, TesselationEvaluation);
impl_pipeline_stage_info!(fragment, Fragment);
impl_pipeline_stage_info!(compute, Compute);
impl_pipeline_stage_info!(any_hit, AnyHit);
impl_pipeline_stage_info!(intersection, Intersection);
impl_pipeline_stage_info!(closest_hit, ClosestHit);
impl_pipeline_stage_info!(ray_generation, RayGeneration);
impl_pipeline_stage_info!(miss, Miss);

impl<const TYPE: u8> VulkanDevice for ShaderModule<TYPE> {
    fn device(&self) -> &Arc<Device> {
        &self.device
    }
}

impl_vk_handle!(ShaderModule<const TYPE: u8,>, VkShaderModule, shader_module);

impl<const TYPE: u8> Drop for ShaderModule<TYPE> {
    fn drop(&mut self) {
        self.device.destroy_shader_module(self.shader_module);
    }
}

pub trait AddSpecializationConstant<T> {
    fn add(&mut self, value: T, id: u32);
}

pub struct SpecializationConstants {
    // store data as raw bytes
    data: Vec<u8>,
    entries: Vec<VkSpecializationMapEntry>,

    info: VkSpecializationInfo,
}

impl SpecializationConstants {
    pub fn new() -> Self {
        let mut me = SpecializationConstants {
            data: Vec::new(),
            entries: Vec::new(),

            info: VkSpecializationInfo::empty(),
        };

        me.info.set_data(&me.data);
        me.info.set_map_entries(&me.entries);

        me
    }

    pub fn vk_handle(&self) -> &VkSpecializationInfo {
        &self.info
    }
}

macro_rules! impl_add_specialization_constant {
    ($($type: ty),+) => {
        $(
            impl AddSpecializationConstant<$type> for SpecializationConstants {
                fn add(&mut self, value: $type, id: u32) {
                    let bytes = value.to_ne_bytes();

                    self.entries.push(VkSpecializationMapEntry {
                        constantID: id,
                        offset: self.data.len() as u32,
                        size: bytes.len(),
                    });

                    self.data.extend(&bytes);
                }
            }
        )+
    };
}

impl_add_specialization_constant!(f32, f64, u64, i64, u32, i32, u16, i16, u8, i8, usize, isize);