imgui-ash 0.1.0

ash/VMA renderer for imgui-rs
Documentation
use ash::vk::{self, Handle};

use crate::{
    DeviceObjects, MIN_DESCRIPTOR_POOL_SIZE, Renderer, RendererCreateError, RendererCreateInfo,
    util::RaiiWrapper,
};

#[repr(align(4))]
pub struct ShaderCode<const N: usize>(pub [u8; N]);
macro_rules! load_shadercode {
    ($bytes:expr) => {{
        static SHADER_CODE: &[u8] = &ShaderCode($bytes).0;
        debug_assert!(SHADER_CODE.len() % 4 == 0);
        unsafe { SHADER_CODE.align_to::<u32>().1 }
    }};
}

impl<A> Renderer<A>
where
    A: vk_mem::Alloc,
{
    pub(crate) unsafe fn create_device_objects(
        _instance: &ash::Instance,
        device: &ash::Device,
        create_info: &RendererCreateInfo,
    ) -> Result<RaiiWrapper<DeviceObjects, impl FnMut(DeviceObjects)>, RendererCreateError> {
        let sampler = {
            let create_info = vk::SamplerCreateInfo::default()
                .mag_filter(vk::Filter::LINEAR)
                .min_filter(vk::Filter::LINEAR)
                .mipmap_mode(vk::SamplerMipmapMode::LINEAR)
                .address_mode_u(vk::SamplerAddressMode::CLAMP_TO_EDGE)
                .address_mode_v(vk::SamplerAddressMode::CLAMP_TO_EDGE)
                .address_mode_w(vk::SamplerAddressMode::CLAMP_TO_EDGE)
                .min_lod(-1000.0)
                .max_lod(1000.0)
                .max_anisotropy(1.0);
            RaiiWrapper::new(
                unsafe { device.create_sampler(&create_info, None) }
                    .map_err(RendererCreateError::SamplerCreateError)?,
                |s| unsafe { device.destroy_sampler(s, None) },
            )
        };

        let descriptor_layout = {
            let descriptor_set_bindings = [vk::DescriptorSetLayoutBinding::default()
                .descriptor_type(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .descriptor_count(1)
                .stage_flags(vk::ShaderStageFlags::FRAGMENT)];
            let create_info =
                vk::DescriptorSetLayoutCreateInfo::default().bindings(&descriptor_set_bindings);

            RaiiWrapper::new(
                unsafe { device.create_descriptor_set_layout(&create_info, None) }
                    .map_err(RendererCreateError::DescriptorSetCreateError)?,
                |l| unsafe {
                    device.destroy_descriptor_set_layout(l, None);
                },
            )
        };

        let descriptor_pool = {
            let pool_size = [vk::DescriptorPoolSize::default()
                .ty(vk::DescriptorType::COMBINED_IMAGE_SAMPLER)
                .descriptor_count(
                    create_info
                        .descriptor_pool_size
                        .unwrap_or(MIN_DESCRIPTOR_POOL_SIZE),
                )];
            let create_info = vk::DescriptorPoolCreateInfo::default()
                .pool_sizes(&pool_size)
                .max_sets(
                    create_info
                        .descriptor_pool_size
                        .unwrap_or(MIN_DESCRIPTOR_POOL_SIZE),
                )
                .flags(vk::DescriptorPoolCreateFlags::FREE_DESCRIPTOR_SET);

            RaiiWrapper::new(
                unsafe { device.create_descriptor_pool(&create_info, None) }
                    .map_err(RendererCreateError::DescriptorPoolCreateError)?,
                |p| unsafe { device.destroy_descriptor_pool(p, None) },
            )
        };

        let mut pc_size = core::mem::size_of::<f32>() * 4 + core::mem::size_of::<bool>();
        let pc_off = 4 - pc_size % 4;
        pc_size += pc_off;
        // SAFETY: references layout but drop order in scope is reverse to declaration so this is fine
        let pipeline_layout = {
            let push_constants = [vk::PushConstantRange::default()
                .stage_flags(vk::ShaderStageFlags::VERTEX)
                .offset(0)
                .size(pc_size as _)];

            let layout = [*descriptor_layout];
            let create_info = vk::PipelineLayoutCreateInfo::default()
                .set_layouts(&layout)
                .push_constant_ranges(&push_constants);

            RaiiWrapper::new(
                unsafe { device.create_pipeline_layout(&create_info, None) }
                    .map_err(RendererCreateError::PipelineLayoutCreateError)?,
                |p| unsafe { device.destroy_pipeline_layout(p, None) },
            )
        };

        let pipeline_and_shaders = RaiiWrapper::new(
            unsafe { Self::create_pipeline(device, create_info, *pipeline_layout) }?,
            |(p, vs, fs)| unsafe {
                device.destroy_pipeline(p, None);
                device.destroy_shader_module(vs, None);
                device.destroy_shader_module(fs, None);
            },
        );

        let tex_command_pool = {
            let create_info =
                vk::CommandPoolCreateInfo::default().queue_family_index(create_info.queue_family);
            RaiiWrapper::new(
                unsafe { device.create_command_pool(&create_info, None) }
                    .map_err(RendererCreateError::TexCommandPoolCreateError)?,
                |p| unsafe { device.destroy_command_pool(p, None) },
            )
        };

        let tex_command_buffer = {
            let create_info = vk::CommandBufferAllocateInfo::default()
                .command_pool(*tex_command_pool)
                .command_buffer_count(1);
            unsafe { device.allocate_command_buffers(&create_info) }
                .map_err(RendererCreateError::TexCommandBufferAllocError)?[0]
        };

        let fence = {
            let create_info = vk::FenceCreateInfo::default().flags(vk::FenceCreateFlags::empty());
            RaiiWrapper::new(
                unsafe { device.create_fence(&create_info, None) }
                    .map_err(RendererCreateError::SyncobjCreateError)?,
                |f| unsafe {
                    device.destroy_fence(f, None);
                },
            )
        };

        let (pipeline, vert_shader, frag_shader) = pipeline_and_shaders.finalise();
        Ok(RaiiWrapper::new(
            DeviceObjects {
                tex_sampler: sampler.finalise(),
                descriptor_layout: descriptor_layout.finalise(),
                descriptor_pool: descriptor_pool.finalise(),
                pipeline_layout: pipeline_layout.finalise(),
                wait_fence: fence.finalise(),
                pipeline,
                vert_shader,
                frag_shader,
                tex_command_pool: tex_command_pool.finalise(),
                tex_command_buffer,
            },
            |device_o| unsafe { Self::destroy_device_objects(device, &device_o) },
        ))
    }

    unsafe fn create_pipeline(
        device: &ash::Device,
        create_info: &RendererCreateInfo,
        layout: vk::PipelineLayout,
    ) -> Result<(vk::Pipeline, vk::ShaderModule, vk::ShaderModule), RendererCreateError> {
        debug_assert!(!layout.is_null(), "pipeline layout must be valid (is null)");

        let vert_shader =
            load_shadercode!(*include_bytes!(concat!(env!("OUT_DIR"), "/main.vert.spv")));
        let frag_shader =
            load_shadercode!(*include_bytes!(concat!(env!("OUT_DIR"), "/main.frag.spv")));

        let vert_shader = RaiiWrapper::new(
            unsafe { create_shader_module(device, vert_shader) }?,
            |sm| unsafe { device.destroy_shader_module(sm, None) },
        );

        let frag_shader = RaiiWrapper::new(
            unsafe { create_shader_module(device, frag_shader) }?,
            |sm| unsafe {
                device.destroy_shader_module(sm, None);
            },
        );

        let stages = [
            vk::PipelineShaderStageCreateInfo::default()
                .stage(vk::ShaderStageFlags::VERTEX)
                .module(*vert_shader)
                .name(c"main"),
            vk::PipelineShaderStageCreateInfo::default()
                .stage(vk::ShaderStageFlags::FRAGMENT)
                .module(*frag_shader)
                .name(c"main"),
        ];

        let binding_desc = [vk::VertexInputBindingDescription::default()
            .stride(core::mem::size_of::<imgui::DrawVert>() as u32)
            .input_rate(vk::VertexInputRate::VERTEX)
            .binding(0)];

        let attribute_description = [
            vk::VertexInputAttributeDescription::default()
                .location(0)
                .binding(binding_desc[0].binding)
                .format(vk::Format::R32G32_SFLOAT)
                .offset(core::mem::offset_of!(imgui::DrawVert, pos) as u32),
            vk::VertexInputAttributeDescription::default()
                .location(1)
                .binding(binding_desc[0].binding)
                .format(vk::Format::R32G32_SFLOAT)
                .offset(core::mem::offset_of!(imgui::DrawVert, uv) as u32),
            vk::VertexInputAttributeDescription::default()
                .location(2)
                .binding(binding_desc[0].binding)
                .format(vk::Format::R8G8B8A8_UNORM)
                .offset(core::mem::offset_of!(imgui::DrawVert, col) as u32),
        ];

        let dynamic_states = [vk::DynamicState::VIEWPORT, vk::DynamicState::SCISSOR];
        let dynamic_state =
            vk::PipelineDynamicStateCreateInfo::default().dynamic_states(&dynamic_states);

        let vertex_info = vk::PipelineVertexInputStateCreateInfo::default()
            .vertex_binding_descriptions(&binding_desc)
            .vertex_attribute_descriptions(&attribute_description);

        let input_assembler = vk::PipelineInputAssemblyStateCreateInfo::default()
            .topology(vk::PrimitiveTopology::TRIANGLE_LIST);

        let viewport_info = vk::PipelineViewportStateCreateInfo::default()
            .viewport_count(1)
            .scissor_count(1);

        let raster_info = vk::PipelineRasterizationStateCreateInfo::default()
            .polygon_mode(vk::PolygonMode::FILL)
            .cull_mode(vk::CullModeFlags::NONE)
            .front_face(vk::FrontFace::COUNTER_CLOCKWISE)
            .line_width(1.0);

        let ms_info = vk::PipelineMultisampleStateCreateInfo::default().rasterization_samples(
            create_info
                .msaa_samples
                .unwrap_or(vk::SampleCountFlags::TYPE_1),
        );

        let color_attachment = [vk::PipelineColorBlendAttachmentState::default()
            .blend_enable(true)
            .src_color_blend_factor(vk::BlendFactor::SRC_ALPHA)
            .dst_color_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)
            .color_blend_op(vk::BlendOp::ADD)
            .src_alpha_blend_factor(vk::BlendFactor::ONE)
            .dst_color_blend_factor(vk::BlendFactor::ONE_MINUS_SRC_ALPHA)
            .alpha_blend_op(vk::BlendOp::ADD)
            .color_write_mask(vk::ColorComponentFlags::RGBA)];

        let depth_info = vk::PipelineDepthStencilStateCreateInfo::default();

        let blend_info =
            vk::PipelineColorBlendStateCreateInfo::default().attachments(&color_attachment);

        let pipeline_info = vk::GraphicsPipelineCreateInfo::default()
            .stages(&stages)
            .vertex_input_state(&vertex_info)
            .input_assembly_state(&input_assembler)
            .viewport_state(&viewport_info)
            .rasterization_state(&raster_info)
            .multisample_state(&ms_info)
            .color_blend_state(&blend_info)
            .depth_stencil_state(&depth_info)
            .dynamic_state(&dynamic_state)
            .layout(layout)
            .render_pass(create_info.render_pass)
            .subpass(create_info.subpass);

        let pipeline = unsafe {
            device.create_graphics_pipelines(create_info.pipeline_cache, &[pipeline_info], None)
        }
        .map_err(|(_, err)| RendererCreateError::PipelineCreateError(err))?;
        Ok((pipeline[0], vert_shader.finalise(), frag_shader.finalise()))
    }
}

unsafe fn create_shader_module(
    device: &ash::Device,
    code: &[u32],
) -> Result<vk::ShaderModule, RendererCreateError> {
    let create_info = vk::ShaderModuleCreateInfo::default().code(code);
    unsafe { device.create_shader_module(&create_info, None) }
        .map_err(RendererCreateError::ShaderModuleCreateError)
}