rust_vulkan_test/src/renderer/vulkan/vk_swapchain.rs

use super::{SwapchainSupportDetails, VkDevice, VkFramebuffer, VkPhysicalDevice, VkRenderPass, VkSemaphore, VkSurface};
use crate::display::Window;
use ash::prelude::VkResult;
use ash::vk;
use std::sync::Arc;

pub struct VkSwapchain {
    surface: Arc<VkSurface>,
    device: Arc<VkDevice>,

    pub handle: Option<vk::SwapchainKHR>,
    swapchain_support_details: SwapchainSupportDetails,

    pub desired_image_count: u32,
    pub surface_format: vk::SurfaceFormatKHR,
    pub surface_resolution: vk::Extent2D,
    pub new_requested_surface_resolution: Option<vk::Extent2D>,
    pub present_mode: vk::PresentModeKHR,
    pub pre_transform: vk::SurfaceTransformFlagsKHR,

    pub present_images: Option<Vec<vk::Image>>,
    pub present_image_views: Option<Vec<Arc<vk::ImageView>>>,
}

impl VkSwapchain {
    pub fn new(
        window: &Window,
        surface: &Arc<VkSurface>,
        device: &Arc<VkDevice>,
        physical_device: &VkPhysicalDevice,
    ) -> anyhow::Result<Self> {
        log::debug!("Creating swapchain");

        let window_size = window
            .physical_size::<u32>()
            .and_then(|size| {
                Some(vk::Extent2D {
                    width: size.width,
                    height: size.height,
                })
            })
            .ok_or_else(|| anyhow::anyhow!("Failed to get swapchain extent"))?;
        log::debug!("Window size ({}x{})", window_size.width, window_size.height);

        let swapchain_support_details =
            surface.get_physical_device_swapchain_support_details(physical_device)?;
        let SwapchainSupportDetails(surface_formats, surface_capabilities, present_modes) =
            &swapchain_support_details;
        log::debug!("Supported surface formats by physical device: {surface_formats:#?}");
        log::debug!("Surface capabilities: {surface_capabilities:#?}");
        log::debug!("Present modes: {present_modes:#?}");

        let surface_format = Self::choose_surface_format(surface_formats)
            .ok_or_else(|| anyhow::anyhow!("No available surface formats"))?;
        let desired_image_count = Self::choose_desired_image_count(surface_capabilities);
        let swapchain_extent = Self::choose_swapchain_extent(window_size, surface_capabilities);
        let pre_transform = Self::choose_pre_transform(surface_capabilities);
        let present_mode = Self::choose_present_mode(present_modes);

        let mut swapchain = Self {
            surface: surface.clone(),
            device: device.clone(),

            handle: None,
            new_requested_surface_resolution: None,
            swapchain_support_details,
            desired_image_count,
            surface_format,
            surface_resolution: swapchain_extent,
            present_mode,
            pre_transform,
            present_images: None,
            present_image_views: None,
        };

        swapchain.create_swapchain()?;

        Ok(swapchain)
    }

    pub fn create_swapchain(&mut self) -> VkResult<()> {
        if let Some(new_requested_surface_resolution) = self.new_requested_surface_resolution {
            self.surface_resolution = new_requested_surface_resolution;
            self.new_requested_surface_resolution = None;
        }

        let mut swapchain_create_info = self.create_swapchain_info(&self.surface);

        if let Some(old_swapchain) = self.handle {
            swapchain_create_info.old_swapchain = old_swapchain;
        }

        let swapchain = unsafe {
            self.device
                .swapchain_loader
                .create_swapchain(&swapchain_create_info, None)?
        };

        let present_images = unsafe {
            self.device
                .swapchain_loader
                .get_swapchain_images(swapchain)?
        };
        let present_images_view = present_images
            .iter()
            .map(|i| {
                self.create_present_image_view(*i)
                    .expect("Failed to create image view")
            })
            .map(|i| Arc::new(i))
            .collect::<Vec<_>>();

        if log::log_enabled!(log::Level::Debug) {
            let label = match self.handle {
                None => "Swapchain created",
                Some(_) => "Swapchain updated",
            };
            log::debug!("{label} ({swapchain:?}) : {swapchain_create_info:#?}");
        }

        self.handle = Some(swapchain);
        self.present_image_views = Some(present_images_view);
        self.present_images = Some(present_images);

        Ok(())
    }

    pub fn create_framebuffers(
        &self,
        render_pass: &Arc<VkRenderPass>,
    ) -> Option<Vec<Arc<VkFramebuffer>>> {
        let present_image_views = self.present_image_views.as_ref()?;

        Some(
            present_image_views
                .iter()
                .map(|image_view| {
                    VkFramebuffer::from_swapchain_image_view(
                        &self.device,
                        &render_pass,
                        &image_view,
                        &self,
                    )
                        .unwrap()
                })
                .map(|framebuffer| Arc::new(framebuffer))
                .collect::<Vec<_>>(),
        )
    }

    pub fn update_resolution(&mut self, width: u32, height: u32) -> VkResult<()> {
        log::debug!("New resolution requested ({width}x{height})");

        let chosen_extent = Self::choose_swapchain_extent(
            vk::Extent2D { width, height },
            &self.swapchain_support_details.1,
        );
        if chosen_extent.width != self.surface_resolution.width
            || chosen_extent.height != self.surface_resolution.height
        {
            self.new_requested_surface_resolution = Some(chosen_extent);
            log::debug!(
                "New resolution submitted ({}x{})",
                chosen_extent.width,
                chosen_extent.height
            );
        } else {
            log::debug!("New resolution skipped ({width}x{height}) : Same resolution");
        }

        Ok(())
    }

    pub fn acquire_next_image(&self, semaphore: &VkSemaphore) -> VkResult<(u32, bool)> {
        unsafe {
            self.device.swapchain_loader.acquire_next_image(
                self.handle.unwrap(),
                u64::MAX,
                semaphore.handle,
                vk::Fence::null(),
            )
        }
    }

    pub fn is_dirty(&self) -> bool {
        self.new_requested_surface_resolution.is_some()
    }

    fn create_swapchain_info(&self, surface: &VkSurface) -> vk::SwapchainCreateInfoKHR {
        vk::SwapchainCreateInfoKHR::default()
            .surface(surface.handle)
            .min_image_count(self.desired_image_count)
            .image_color_space(self.surface_format.color_space)
            .image_format(self.surface_format.format)
            .image_extent(self.surface_resolution)
            .image_usage(vk::ImageUsageFlags::COLOR_ATTACHMENT)
            .image_sharing_mode(vk::SharingMode::EXCLUSIVE)
            .pre_transform(self.pre_transform)
            .composite_alpha(vk::CompositeAlphaFlagsKHR::OPAQUE)
            .present_mode(self.present_mode)
            .clipped(true)
            .image_array_layers(1)
    }

    fn choose_swapchain_extent(
        window_size: vk::Extent2D,
        surface_capabilities: &vk::SurfaceCapabilitiesKHR,
    ) -> vk::Extent2D {
        vk::Extent2D {
            width: window_size
                .width
                .max(surface_capabilities.min_image_extent.width)
                .min(surface_capabilities.max_image_extent.width),
            height: window_size
                .height
                .max(surface_capabilities.min_image_extent.height)
                .min(surface_capabilities.max_image_extent.height),
        }
    }

    fn choose_surface_format(
        surface_formats: &Vec<vk::SurfaceFormatKHR>,
    ) -> Option<vk::SurfaceFormatKHR> {
        surface_formats.first().and_then(|f| Some(*f))
    }

    fn choose_desired_image_count(surface_capabilities: &vk::SurfaceCapabilitiesKHR) -> u32 {
        let mut desired_image_count = surface_capabilities.min_image_count + 1;
        if surface_capabilities.max_image_count > 0
            && desired_image_count > surface_capabilities.max_image_count
        {
            desired_image_count = surface_capabilities.max_image_count;
        }
        desired_image_count
    }

    fn choose_pre_transform(
        surface_capabilities: &vk::SurfaceCapabilitiesKHR,
    ) -> vk::SurfaceTransformFlagsKHR {
        if surface_capabilities
            .supported_transforms
            .contains(vk::SurfaceTransformFlagsKHR::IDENTITY)
        {
            vk::SurfaceTransformFlagsKHR::IDENTITY
        } else {
            surface_capabilities.current_transform
        }
    }

    fn choose_present_mode(present_modes: &Vec<vk::PresentModeKHR>) -> vk::PresentModeKHR {
        present_modes
            .iter()
            .cloned()
            .find(|&mode| mode == vk::PresentModeKHR::MAILBOX)
            .unwrap_or(vk::PresentModeKHR::FIFO)
    }

    fn create_present_image_view(&self, image: vk::Image) -> VkResult<vk::ImageView> {
        let create_view_info = vk::ImageViewCreateInfo::default()
            .view_type(vk::ImageViewType::TYPE_2D)
            .format(self.surface_format.format)
            .components(vk::ComponentMapping {
                r: vk::ComponentSwizzle::IDENTITY,
                g: vk::ComponentSwizzle::IDENTITY,
                b: vk::ComponentSwizzle::IDENTITY,
                a: vk::ComponentSwizzle::IDENTITY,
            })
            .subresource_range(vk::ImageSubresourceRange {
                aspect_mask: vk::ImageAspectFlags::COLOR,
                base_mip_level: 0,
                level_count: 1,
                base_array_layer: 0,
                layer_count: 1,
            })
            .image(image);

        unsafe {
            self.device
                .handle
                .create_image_view(&create_view_info, None)
        }
    }
}

impl Drop for VkSwapchain {
    fn drop(&mut self) {
        if let Some(swapchain) = self.handle {
            unsafe {
                self.device
                    .swapchain_loader
                    .destroy_swapchain(swapchain, None);
            }
            self.handle = None;
            log::debug!("Swapchain destroyed ({swapchain:?})");
        }
    }
}