rust_vulkan_test/src/renderer/vulkan/vk_render_context.rs

use super::{
    VkCommandPool, VkDevice, VkFence, VkFramebuffer, VkGraphicsPipeline, VkInstance, VkPhysicalDevice,
    VkRenderPass, VkSemaphore, VkSurface, VkSwapchain,
};
use ash::vk;
use std::sync::Arc;
use crate::renderer::Renderable;

pub struct VkRenderContext {
    instance: Arc<VkInstance>,
    surface: Arc<VkSurface>,
    device: Arc<VkDevice>,

    swapchain: Arc<VkSwapchain>,
    render_pass: Arc<VkRenderPass>,
    framebuffers: Vec<Arc<VkFramebuffer>>,

    command_pool: VkCommandPool,
    command_buffers: Vec<vk::CommandBuffer>,
    image_available_semaphore: VkSemaphore,
    render_finished_semaphore: VkSemaphore,
    in_flight_fence: VkFence,
}

impl VkRenderContext {
    pub fn init(window: &crate::display::Window) -> anyhow::Result<Self> {
        let required_extensions = window.required_extensions()?;

        let instance = Arc::new(VkInstance::new(&required_extensions));
        let surface = Arc::new(VkSurface::new(&window, instance.clone())?);

        let mut physical_devices = instance.get_physical_devices();
        physical_devices.sort_by(|a, b| b.priority().cmp(&a.priority()));

        let (physical_device, queue_family_index, properties) =
            VkPhysicalDevice::pick_physical_device_and_queue_by(
                &physical_devices,
                Some(vk::QueueFlags::GRAPHICS),
                Some(&surface),
            )
                .ok_or_else(|| anyhow::anyhow!("Unable to find physical device"))?;
        log::debug!(
            "Selected queue {properties:#?} for physical device {:?}",
            physical_device.properties.device_name_as_c_str()
        );

        let device = Arc::new(VkDevice::new_graphics_device(
            &instance,
            &physical_device,
            queue_family_index,
        )?);

        let swapchain = Arc::new(VkSwapchain::new(
            &window,
            &surface,
            &device,
            &physical_device,
        )?);

        let render_pass = Arc::new(VkRenderPass::new(&device, &swapchain)?);

        let framebuffers = swapchain
            .create_framebuffers(&render_pass)
            .ok_or_else(|| anyhow::anyhow!("Failed to get framebuffers"))?;

        let command_pool = VkCommandPool::new(&device)?;

        // Destroyed with command pool
        let command_buffers = command_pool
            .allocate_command_buffers_for_framebuffers(framebuffers.len() as u32)?;

        let image_available_semaphore = VkSemaphore::new(&device)?;
        let render_finished_semaphore = VkSemaphore::new(&device)?;
        let in_flight_fence = VkFence::new(&device)?;

        Ok(Self {
            instance,
            surface,
            device,

            swapchain,
            render_pass,
            framebuffers,

            command_pool,
            command_buffers,

            image_available_semaphore,
            render_finished_semaphore,
            in_flight_fence,
        })
    }

    pub fn render(&mut self, scene: Option<&Box<dyn Renderable>>) -> anyhow::Result<()> {
        unsafe { self.device.handle.wait_for_fences(&[self.in_flight_fence.handle], true, u64::MAX)? };
        unsafe { self.device.handle.reset_fences(&[self.in_flight_fence.handle])? };

        let (index, _) = self
            .swapchain
            .acquire_next_image(&self.image_available_semaphore)?;

        // if self.swapchain.is_dirty() {
        //     self.update_swapchain()?
        // }

        let command_buffer = self.command_buffers[index as usize];
        unsafe { self.device.handle.reset_command_buffer(command_buffer, vk::CommandBufferResetFlags::default())? };

        let render_area = vk::Rect2D::default().extent(self.swapchain.surface_resolution);
        let clear_value = vk::ClearValue::default();
        let command_buffer_begin_info = vk::CommandBufferBeginInfo::default();
        unsafe {
            self.device
                .handle
                .begin_command_buffer(command_buffer, &command_buffer_begin_info)?
        };

        let clear_values = [clear_value];
        let framebuffer = self.framebuffers[index as usize].as_ref();
        let render_pass_begin_info = vk::RenderPassBeginInfo::default()
            .render_pass(self.render_pass.handle)
            .framebuffer(framebuffer.handle)
            .render_area(render_area)
            .clear_values(&clear_values);

        unsafe {
            self.device.handle.cmd_begin_render_pass(
                command_buffer,
                &render_pass_begin_info,
                vk::SubpassContents::INLINE,
            );
        };

        if let Some(scene) = scene {
            scene.render(&self.device, &self.swapchain, &command_buffer)?;
        }

        unsafe { self.device.handle.cmd_end_render_pass(command_buffer) };

        unsafe { self.device.handle.end_command_buffer(command_buffer)? };

        let wait_semaphores = [self.image_available_semaphore.handle];
        let signal_semaphores = [self.render_finished_semaphore.handle];
        let wait_stages = [vk::PipelineStageFlags::COLOR_ATTACHMENT_OUTPUT];
        let command_buffers_to_submit = [command_buffer];
        let submit_info = vk::SubmitInfo::default()
            .wait_semaphores(&wait_semaphores)
            .wait_dst_stage_mask(&wait_stages)
            .command_buffers(&command_buffers_to_submit)
            .signal_semaphores(&signal_semaphores);

        let queue = self
            .device
            .get_device_queue(0)
            .ok_or_else(|| anyhow::anyhow!("Failed to get a queue"))?;

        unsafe {
            self.device
                .handle
                .queue_submit(*queue, &[submit_info], self.in_flight_fence.handle)?
        };

        let swapchains = [self.swapchain.handle.unwrap()];
        let indices = [index];
        let present_info = vk::PresentInfoKHR::default()
            .wait_semaphores(&signal_semaphores)
            .swapchains(&swapchains)
            .image_indices(&indices);

        unsafe { self.device.swapchain_loader.queue_present(*queue, &present_info)? };

        Ok(())
    }

    pub fn update_resolution(&mut self, width: u32, height: u32) -> anyhow::Result<()> {
        match Arc::get_mut(&mut self.swapchain) {
            Some(swapchain) => swapchain.update_resolution(width, height)?,
            None => log::warn!("Impossible to get mutable swapchain"),
        }

        Ok(())
    }

    pub fn exit(&self) {
        unsafe { self.device.handle.device_wait_idle().unwrap() }
    }

    pub fn init_scene(&self, scene: &mut Box<dyn Renderable>) -> anyhow::Result<()> {
        scene.init(&self.device, &self.render_pass)
    }

    fn update_swapchain(&mut self) -> anyhow::Result<()> {
        match Arc::get_mut(&mut self.swapchain) {
            Some(swapchain) => {
                swapchain.create_swapchain()?;

                self.framebuffers = self.swapchain
                    .create_framebuffers(&self.render_pass)
                    .ok_or_else(|| anyhow::anyhow!("Failed to get framebuffers"))?;
            }
            None => log::warn!("Impossible to get mutable swapchain"),
        }

        Ok(())
    }
}