renderer: extract swapchain to its own file, refactor types

crates/renderer/src/swapchain.rs

@@ -0,0 +1,763 @@
+use std::{
+    marker::PhantomData,
+    sync::{
+        atomic::{AtomicU32, AtomicU64},
+        Arc,
+    },
+};
+
+use ash::{
+    prelude::VkResult,
+    vk::{self, Handle},
+};
+use parking_lot::{RawMutex, RwLock};
+use raw_window_handle::{RawDisplayHandle, RawWindowHandle};
+
+use crate::{
+    define_device_owned_handle,
+    device::{Device, DeviceOwned},
+    images, sync,
+    util::RawMutexGuard,
+    Instance, Result,
+};
+
+define_device_owned_handle! {
+    #[derive(Debug)]
+    pub Surface(vk::SurfaceKHR) {
+    } => |this| unsafe {
+        this.device().instance().surface.destroy_surface(this.handle(), None);
+    }
+}
+
+impl Surface {
+    #[allow(dead_code)]
+    pub fn headless(device: Device) -> Result<Self> {
+        unsafe {
+            let instance = device.instance();
+            let headless_instance =
+                ash::ext::headless_surface::Instance::new(&instance.entry, &instance.instance);
+
+            let surface = headless_instance
+                .create_headless_surface(&vk::HeadlessSurfaceCreateInfoEXT::default(), None)?;
+
+            Ok(Self::construct(
+                device,
+                surface,
+                Some("headless-surface".into()),
+            )?)
+        }
+    }
+
+    pub fn create(
+        device: Device,
+        display_handle: RawDisplayHandle,
+        window_handle: raw_window_handle::RawWindowHandle,
+    ) -> Result<Self> {
+        let instance = device.instance();
+        let surface = unsafe {
+            ash_window::create_surface(
+                &instance.entry,
+                &instance.instance,
+                display_handle,
+                window_handle,
+                None,
+            )?
+        };
+
+        Ok(Self::construct(
+            device,
+            surface,
+            Some(format!("{:?}_surface", window_handle).into()),
+        )?)
+    }
+}
+
+define_device_owned_handle! {
+    pub Swapchain(vk::SwapchainKHR) {
+        mutex: RawMutex,
+        surface: Arc<Surface>,
+
+        #[allow(unused)]
+        present_mode: vk::PresentModeKHR,
+        #[allow(unused)]
+        color_space: vk::ColorSpaceKHR,
+        format: vk::Format,
+        images: Vec<Arc<images::Image>>,
+        image_views: Vec<vk::ImageView>,
+        extent: vk::Extent2D,
+        min_image_count: u32,
+
+        // sync objects:
+        // we need two semaphores per each image, one acquire-semaphore and one release-semaphore.
+        // semaphores must be unique to each frame and cannot be reused per swapchain.
+        acquire_semaphores: Vec<vk::Semaphore>,
+        release_semaphores: Vec<vk::Semaphore>,
+
+        // one fence per in-flight frame, to synchronize image acquisition
+        fences: Vec<Arc<sync::Fence>>,
+
+        current_frame: AtomicU32,
+
+        // for khr_present_id/khr_present_wait
+        #[allow(unused)]
+        present_id: AtomicU64,
+    } => |this| unsafe {
+        _ = this.device().wait_queue_idle(this.device().present_queue());
+        tracing::debug!("dropping swapchain {:?}", this.handle());
+        for view in &this.image_views {
+            this.device().dev().destroy_image_view(*view, None);
+        }
+
+        this.with_locked(|swapchain| {
+            this.device().swapchain().destroy_swapchain(swapchain, None)
+        });
+
+        for &semaphore in this
+            .acquire_semaphores
+            .iter()
+            .chain(&this.release_semaphores)
+        {
+            this.device().dev().destroy_semaphore(semaphore, None);
+        }
+    }
+}
+
+impl core::fmt::Debug for Swapchain {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Swapchain")
+            .field("inner", &self.inner)
+            .field("surface", &self.surface)
+            .field("present_mode", &self.present_mode)
+            .field("color_space", &self.color_space)
+            .field("format", &self.format)
+            .field("images", &self.images)
+            .field("image_views", &self.image_views)
+            .field("extent", &self.extent)
+            .field("min_image_count", &self.min_image_count)
+            .field("acquire_semaphores", &self.acquire_semaphores)
+            .field("release_semaphores", &self.release_semaphores)
+            .field("fences", &self.fences)
+            .field("current_frame", &self.current_frame)
+            .field("present_id", &self.present_id)
+            .finish()
+    }
+}
+
+impl Swapchain {
+    const PREFERRED_IMAGES_IN_FLIGHT: u32 = 3;
+
+    fn get_swapchain_params_from_surface(
+        instance: &Arc<Instance>,
+        surface: vk::SurfaceKHR,
+        pdev: vk::PhysicalDevice,
+        requested_extent: Option<vk::Extent2D>,
+    ) -> Result<SwapchainParams> {
+        let caps = unsafe {
+            instance
+                .surface
+                .get_physical_device_surface_capabilities(pdev, surface)?
+        };
+        let formats = unsafe {
+            instance
+                .surface
+                .get_physical_device_surface_formats(pdev, surface)?
+        };
+        let present_modes = unsafe {
+            instance
+                .surface
+                .get_physical_device_surface_present_modes(pdev, surface)?
+        };
+
+        let present_mode = present_modes
+            .iter()
+            .find(|&mode| mode == &vk::PresentModeKHR::MAILBOX)
+            .cloned()
+            .unwrap_or(vk::PresentModeKHR::FIFO);
+
+        let format = formats
+            .iter()
+            .max_by_key(|&&format| {
+                let is_rgba_unorm = format.format == vk::Format::R8G8B8A8_UNORM
+                    || format.format == vk::Format::B8G8R8A8_UNORM;
+                let is_srgb = format.color_space == vk::ColorSpaceKHR::SRGB_NONLINEAR;
+                is_rgba_unorm as u8 * 10 + is_srgb as u8
+            })
+            .or(formats.first())
+            .cloned()
+            .expect("no surface format available!");
+
+        // 0 here means no limit
+        let max_image_count = core::num::NonZero::new(caps.max_image_count)
+            .map(|n| n.get())
+            .unwrap_or(u32::MAX);
+
+        // we want PREFERRED_IMAGES_IN_FLIGHT images acquired at the same time,
+        let image_count =
+            (caps.min_image_count + Self::PREFERRED_IMAGES_IN_FLIGHT).min(max_image_count);
+
+        let extent = current_extent_or_clamped(
+            &caps,
+            requested_extent.unwrap_or(vk::Extent2D::default().width(1).height(1)),
+        );
+
+        Ok(SwapchainParams {
+            present_mode,
+            format: format.format,
+            color_space: format.color_space,
+            image_count,
+            extent,
+            min_image_count: caps.min_image_count,
+        })
+    }
+
+    pub fn new(
+        device: Device,
+        surface: Arc<Surface>,
+        pdev: vk::PhysicalDevice,
+        extent: vk::Extent2D,
+    ) -> Result<Self> {
+        Self::create(device, surface, pdev, Some(extent), None)
+    }
+
+    fn create(
+        device: Device,
+        surface: Arc<Surface>,
+        pdev: vk::PhysicalDevice,
+        extent: Option<vk::Extent2D>,
+        old_swapchain: Option<&Self>,
+    ) -> Result<Self> {
+        let SwapchainParams {
+            present_mode,
+            format,
+            color_space,
+            image_count,
+            min_image_count,
+            extent,
+        } = Self::get_swapchain_params_from_surface(
+            device.instance(),
+            surface.handle(),
+            pdev,
+            extent,
+        )?;
+
+        let (swapchain, images) = {
+            let lock = old_swapchain.as_ref().map(|handle| handle.lock());
+
+            Self::create_vkswapchainkhr(
+                &device,
+                surface.handle(),
+                &device.queue_families().swapchain_family_indices(),
+                extent,
+                lock.as_ref().map(|lock| **lock),
+                present_mode,
+                format,
+                color_space,
+                image_count,
+            )
+        }?;
+
+        let images = images
+            .iter()
+            .enumerate()
+            .map(|(i, image)| unsafe {
+                images::Image::from_swapchain_image(
+                    device.clone(),
+                    *image,
+                    Some(format!("swapchain-{:x}-image-{i}", swapchain.as_raw()).into()),
+                    vk::Extent3D {
+                        width: extent.width,
+                        height: extent.height,
+                        depth: 1,
+                    },
+                    format,
+                )
+                .inspect(|img| {
+                    _ = img.get_view(images::ImageViewDesc {
+                        // TODO: make this a function that uses debug name/surface handle
+                        name: Some(
+                            format!("swapchain-{:x}-image-{i}-view", swapchain.as_raw()).into(),
+                        ),
+                        kind: vk::ImageViewType::TYPE_2D,
+                        format,
+                        aspect: vk::ImageAspectFlags::COLOR,
+                        ..Default::default()
+                    });
+                })
+                .map(|img| Arc::new(img))
+            })
+            .collect::<VkResult<Vec<_>>>()?;
+
+        let image_views = images
+            .iter()
+            .enumerate()
+            .map(|(i, image)| {
+                image.get_view(images::ImageViewDesc {
+                    name: Some(format!("swapchain-{:x}-image-{i}-view", swapchain.as_raw()).into()),
+                    kind: vk::ImageViewType::TYPE_2D,
+                    format,
+                    aspect: vk::ImageAspectFlags::COLOR,
+                    ..Default::default()
+                })
+            })
+            .collect::<VkResult<Vec<_>>>()?;
+
+        let num_images = images.len() as u32;
+        let inflight_frames = num_images - min_image_count;
+
+        let acquire_semaphores = {
+            (0..inflight_frames)
+                .map(|i| unsafe {
+                    device
+                        .dev()
+                        .create_semaphore(&vk::SemaphoreCreateInfo::default(), None)
+                        .inspect(|r| {
+                            #[cfg(debug_assertions)]
+                            {
+                                device
+                                    .debug_name_object(
+                                        *r,
+                                        &format!("semaphore-{:x}_{i}-acquire", swapchain.as_raw()),
+                                    )
+                                    .unwrap();
+                            }
+                        })
+                })
+                .collect::<VkResult<Vec<_>>>()?
+        };
+
+        let release_semaphores = {
+            (0..inflight_frames)
+                .map(|i| unsafe {
+                    device
+                        .dev()
+                        .create_semaphore(&vk::SemaphoreCreateInfo::default(), None)
+                        .inspect(|r| {
+                            #[cfg(debug_assertions)]
+                            {
+                                device
+                                    .debug_name_object(
+                                        *r,
+                                        &format!("semaphore-{:x}_{i}-release", swapchain.as_raw()),
+                                    )
+                                    .unwrap();
+                            }
+                        })
+                })
+                .collect::<VkResult<Vec<_>>>()?
+        };
+
+        let fences = {
+            (0..inflight_frames)
+                .map(|i| {
+                    Ok(Arc::new(sync::Fence::create(device.clone()).inspect(
+                        |r| {
+                            #[cfg(debug_assertions)]
+                            {
+                                device
+                                    .debug_name_object(
+                                        r.fence(),
+                                        &format!("fence-{:x}_{i}", swapchain.as_raw()),
+                                    )
+                                    .unwrap();
+                            }
+                        },
+                    )?))
+                })
+                .collect::<VkResult<Vec<_>>>()?
+        };
+
+        tracing::trace!("fences: {fences:?}");
+
+        Ok(Self::construct(
+            device,
+            swapchain,
+            Some(
+                format!(
+                    "swapchain-{}_{}",
+                    surface.handle().as_raw(),
+                    SWAPCHAIN_COUNT.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
+                )
+                .into(),
+            ),
+            <parking_lot::RawMutex as parking_lot::lock_api::RawMutex>::INIT,
+            surface,
+            present_mode,
+            color_space,
+            format,
+            images,
+            image_views,
+            extent,
+            min_image_count,
+            acquire_semaphores,
+            release_semaphores,
+            fences,
+            AtomicU32::new(0),
+            AtomicU64::new(1),
+        )?)
+    }
+
+    pub fn max_in_flight_images(&self) -> u32 {
+        self.num_images() - self.min_image_count
+    }
+
+    pub fn num_images(&self) -> u32 {
+        self.images.len() as u32
+    }
+
+    fn recreate(&self, extent: Option<vk::Extent2D>) -> Result<Self> {
+        Self::create(
+            self.device().clone(),
+            self.surface.clone(),
+            self.device().phy(),
+            extent,
+            Some(&self),
+        )
+    }
+
+    /// returns a future yielding the frame, and true if the swapchain is
+    /// suboptimal and should be recreated.
+    fn acquire_image(
+        self: Arc<Self>,
+    ) -> impl std::future::Future<Output = VkResult<(SwapchainFrame, bool)>> {
+        let frame = self
+            .current_frame
+            .fetch_update(
+                std::sync::atomic::Ordering::Release,
+                std::sync::atomic::Ordering::Relaxed,
+                |i| Some((i + 1) % self.max_in_flight_images()),
+            )
+            .unwrap() as usize;
+
+        tracing::trace!(frame, "acquiring image for frame {frame}");
+
+        async move {
+            let fence = self.fences[frame].clone();
+            let acquire = self.acquire_semaphores[frame];
+            let release = self.release_semaphores[frame];
+
+            // spawn on threadpool because it might block.
+            let (idx, suboptimal) = smol::unblock({
+                let this = self.clone();
+                let fence = fence.clone();
+                move || unsafe {
+                    this.with_locked(|swapchain| {
+                        this.device().swapchain().acquire_next_image(
+                            swapchain,
+                            u64::MAX,
+                            acquire,
+                            fence.fence(),
+                        )
+                    })
+                }
+            })
+            .await?;
+
+            // wait for image to become available.
+            sync::FenceFuture::new(fence.clone()).await;
+
+            let idx = idx as usize;
+            let image = self.images[idx].clone();
+            let view = self.image_views[idx];
+
+            Ok((
+                SwapchainFrame {
+                    index: idx as u32,
+                    swapchain: self.clone(),
+                    format: self.format,
+                    image,
+                    view,
+                    acquire,
+                    release,
+                },
+                suboptimal,
+            ))
+        }
+    }
+
+    fn present(&self, frame: SwapchainFrame, wait: Option<vk::Semaphore>) -> Result<()> {
+        let swpchain = self.lock();
+        let queue = self.device().present_queue().lock();
+
+        let wait_semaphores = wait
+            .as_ref()
+            .map(|sema| core::slice::from_ref(sema))
+            .unwrap_or_default();
+
+        // TODO: make this optional for devices with no support for present_wait/present_id
+        // let present_id = self
+        //     .present_id
+        //     .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        // let mut present_id =
+        //     vk::PresentIdKHR::default().present_ids(core::slice::from_ref(&present_id));
+
+        let present_info = vk::PresentInfoKHR::default()
+            .image_indices(core::slice::from_ref(&frame.index))
+            .swapchains(core::slice::from_ref(&swpchain))
+            .wait_semaphores(wait_semaphores);
+        //.push_next(&mut present_id)
+
+        // call winits pre_present_notify here
+
+        unsafe {
+            self.device()
+                .swapchain()
+                .queue_present(*queue, &present_info)?;
+        }
+        Ok(())
+    }
+
+    fn create_vkswapchainkhr(
+        device: &Device,
+        surface: vk::SurfaceKHR,
+        queue_families: &[u32],
+        image_extent: vk::Extent2D,
+        old_swapchain: Option<vk::SwapchainKHR>,
+        present_mode: vk::PresentModeKHR,
+        image_format: vk::Format,
+        image_color_space: vk::ColorSpaceKHR,
+        image_count: u32,
+    ) -> Result<(vk::SwapchainKHR, Vec<vk::Image>)> {
+        let create_info = vk::SwapchainCreateInfoKHR::default()
+            .surface(surface)
+            .present_mode(present_mode)
+            .image_color_space(image_color_space)
+            .image_format(image_format)
+            .min_image_count(image_count)
+            .image_usage(vk::ImageUsageFlags::TRANSFER_DST | vk::ImageUsageFlags::COLOR_ATTACHMENT)
+            .image_array_layers(1)
+            .image_extent(image_extent)
+            .image_sharing_mode(if queue_families.len() <= 1 {
+                vk::SharingMode::EXCLUSIVE
+            } else {
+                vk::SharingMode::CONCURRENT
+            })
+            .queue_family_indices(queue_families)
+            .pre_transform(vk::SurfaceTransformFlagsKHR::IDENTITY)
+            .composite_alpha(vk::CompositeAlphaFlagsKHR::OPAQUE)
+            .old_swapchain(old_swapchain.unwrap_or(vk::SwapchainKHR::null()))
+            .clipped(true);
+
+        let (swapchain, images) = unsafe {
+            let swapchain = device.swapchain().create_swapchain(&create_info, None)?;
+
+            let images = device.swapchain().get_swapchain_images(swapchain)?;
+
+            (swapchain, images)
+        };
+
+        Ok((swapchain, images))
+    }
+}
+static SWAPCHAIN_COUNT: AtomicU64 = AtomicU64::new(0);
+
+#[derive(Debug)]
+#[must_use = "This struct represents an acquired image from the swapchain and
+ must be presented in order to free resources on the device."]
+pub struct SwapchainFrame {
+    pub swapchain: Arc<Swapchain>,
+    pub index: u32,
+    pub image: Arc<images::Image>,
+    pub format: vk::Format,
+    pub view: vk::ImageView,
+    pub acquire: vk::Semaphore,
+    pub release: vk::Semaphore,
+}
+
+impl Eq for SwapchainFrame {}
+impl PartialEq for SwapchainFrame {
+    fn eq(&self, other: &Self) -> bool {
+        self.index == other.index && self.image == other.image
+    }
+}
+
+impl SwapchainFrame {
+    pub fn present(self, wait: Option<vk::Semaphore>) -> crate::Result<()> {
+        self.swapchain.clone().present(self, wait)
+    }
+}
+
+fn current_extent_or_clamped(
+    caps: &vk::SurfaceCapabilitiesKHR,
+    fallback: vk::Extent2D,
+) -> vk::Extent2D {
+    if caps.current_extent.width == u32::MAX {
+        vk::Extent2D {
+            width: fallback
+                .width
+                .clamp(caps.min_image_extent.width, caps.max_image_extent.width),
+            height: fallback
+                .height
+                .clamp(caps.min_image_extent.height, caps.max_image_extent.height),
+        }
+    } else {
+        caps.current_extent
+    }
+}
+
+struct SwapchainParams {
+    present_mode: vk::PresentModeKHR,
+    format: vk::Format,
+    color_space: vk::ColorSpaceKHR,
+    /// the number of images to request from the device
+    image_count: u32,
+    /// the minimum number of images the surface permits
+    min_image_count: u32,
+    extent: vk::Extent2D,
+}
+
+impl Swapchain {
+    pub fn lock(&self) -> RawMutexGuard<'_, vk::SwapchainKHR> {
+        use parking_lot::lock_api::RawMutex;
+        self.mutex.lock();
+        RawMutexGuard {
+            mutex: &self.mutex,
+            value: &self.inner.object,
+            _pd: PhantomData,
+        }
+    }
+
+    pub fn with_locked<T, F: FnOnce(vk::SwapchainKHR) -> T>(&self, f: F) -> T {
+        let lock = self.lock();
+        f(*lock)
+    }
+}
+
+#[derive(Debug)]
+pub struct WindowSurface {
+    // window_handle: RawWindowHandle,
+    pub surface: Arc<Surface>,
+    // this mutex is for guarding the swapchain against being replaced
+    // underneath WindowContext's functions
+    current_swapchain: RwLock<Arc<Swapchain>>,
+}
+
+impl WindowSurface {
+    pub fn new(
+        device: Device,
+        requested_extent: vk::Extent2D,
+        window: RawWindowHandle,
+        display: RawDisplayHandle,
+    ) -> Result<Self> {
+        let surface = Arc::new(Surface::create(device.clone(), display, window)?);
+        let swapchain = RwLock::new(Arc::new(Swapchain::new(
+            device.clone(),
+            surface.clone(),
+            device.phy(),
+            requested_extent,
+        )?));
+
+        Ok(Self {
+            surface,
+            // window_handle: window,
+            current_swapchain: swapchain,
+        })
+    }
+
+    /// spawns a task that continuously requests images from the current
+    /// swapchain, sending them to a channel.  returns the receiver of the
+    /// channel, and a handle to the task, allowing for cancellation.
+    pub fn images(
+        self: Arc<Self>,
+    ) -> (
+        smol::channel::Receiver<SwapchainFrame>,
+        smol::Task<std::result::Result<(), crate::Error>>,
+    ) {
+        let (tx, rx) = smol::channel::bounded(8);
+        let task = smol::spawn(async move {
+            loop {
+                let frame = self.acquire_image().await?;
+                tx.send(frame)
+                    .await
+                    .expect("channel closed on swapchain acquiring frame");
+            }
+        });
+
+        (rx, task)
+    }
+
+    pub async fn acquire_image(&self) -> Result<SwapchainFrame> {
+        // clone swapchain to keep it alive
+        let swapchain = self.current_swapchain.read().clone();
+        let (frame, suboptimal) = swapchain.clone().acquire_image().await?;
+        if suboptimal {
+            let mut lock = self.current_swapchain.write();
+            // only recreate our swapchain if it is still same, or else it might have already been recreated.
+            if Arc::ptr_eq(&swapchain, &lock) {
+                *lock = Arc::new(lock.recreate(None)?);
+            }
+        }
+
+        Ok(frame)
+    }
+
+    pub fn recreate_with(&self, extent: Option<vk::Extent2D>) -> Result<()> {
+        let mut swapchain = self.current_swapchain.write();
+        *swapchain = Arc::new(swapchain.recreate(extent)?);
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::device;
+
+    use super::*;
+
+    fn create_headless_vk() -> Result<(Device, WindowSurface)> {
+        let device = Device::new_from_default_desc(
+            None,
+            &[
+                device::Extension {
+                    name: "VK_EXT_headless_surface",
+                    version: ash::ext::headless_surface::SPEC_VERSION,
+                },
+                device::Extension {
+                    name: "VK_KHR_surface",
+                    version: ash::khr::surface::SPEC_VERSION,
+                },
+            ],
+        )?;
+
+        let surface = Arc::new(Surface::headless(device.clone())?);
+
+        let swapchain = Arc::new(Swapchain::new(
+            device.clone(),
+            surface.clone(),
+            device.phy(),
+            vk::Extent2D::default().width(1).height(1),
+        )?);
+
+        let window_ctx = WindowSurface {
+            // window_handle: RawWindowHandle::Web(raw_window_handle::WebWindowHandle::new(0)),
+            surface,
+            current_swapchain: RwLock::new(swapchain),
+        };
+
+        Ok((device, window_ctx))
+    }
+
+    #[tracing_test::traced_test]
+    #[test]
+    fn async_swapchain_acquiring() {
+        let (_dev, ctx) = create_headless_vk().expect("init");
+        let ctx = Arc::new(ctx);
+        let (rx, handle) = ctx.clone().images();
+
+        eprintln!("hello world!");
+
+        let mut count = 0;
+        loop {
+            let now = std::time::Instant::now();
+            let frame = rx.recv_blocking().expect("recv");
+
+            _ = frame.present(None);
+            tracing::info!("mspf: {:.3}ms", now.elapsed().as_micros() as f32 / 1e3);
+            count += 1;
+            if count > 1000 {
+                smol::block_on(handle.cancel());
+                break;
+            }
+        }
+    }
+}