engine/presentation/src/renderbackend.rs

use crate::Result;
use crate::prelude::*;

use cgmath::{Matrix4, SquareMatrix};
use ecs::World;
use ui::prelude::*;
use vulkan_rs::prelude::*;

use std::any::TypeId;
use std::ops::Deref;
use std::sync::{
    Arc, Mutex, RwLock,
    atomic::{AtomicUsize, Ordering::SeqCst},
};

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Eye {
    Left,
    Right,
}

#[derive(Debug, Clone, Copy)]
pub struct VRTransformations {
    pub proj: Matrix4<f32>,
    pub view: Matrix4<f32>,
}

impl VRTransformations {
    pub fn invert(self) -> Result<Self> {
        Ok(VRTransformations {
            proj: self
                .proj
                .invert()
                .ok_or(anyhow::Error::msg("Could not invert Matrix"))?,
            view: self
                .view
                .invert()
                .ok_or(anyhow::Error::msg("Could not invert Matrix"))?,
        })
    }
}

impl Default for VRTransformations {
    fn default() -> Self {
        VRTransformations {
            proj: Matrix4::identity(),
            view: Matrix4::identity(),
        }
    }
}

struct SceneHandle {
    type_id: TypeId,
    priority: u32,

    render: Box<
        dyn FnMut(
                &mut CommandBufferRecorder<'_>,
                &TargetMode<Vec<Arc<Image>>>,
                &TargetMode<usize>,
                &mut World,
            ) -> Result<()>
            + Send
            + Sync
            + 'static,
    >,

    resize: Box<
        dyn FnMut(f32, f32, &TargetMode<Vec<Arc<Image>>>, &mut World) -> Result<()>
            + Send
            + Sync
            + 'static,
    >,
}

impl SceneHandle {
    fn new<T: TScene>(priority: u32) -> Self {
        Self {
            type_id: TypeId::of::<T>(),
            priority,
            render: Box::new(|recorder, images, indices, world| {
                world
                    .resources
                    .get_mut_unchecked::<T>()
                    .process(recorder, images, indices, world)
            }),

            resize: Box::new(|width, height, images, world| {
                world
                    .resources
                    .get_mut_unchecked::<T>()
                    .resize(width, height, images, world)
            }),
        }
    }
}

pub struct RenderBackend {
    device: Arc<Device>,
    queue: Arc<Mutex<Queue>>,

    // driver provided images
    swapchain_images: Mutex<TargetMode<Vec<Arc<Image>>>>,
    image_count: AtomicUsize,

    clear_color: RwLock<VkClearColorValue>,

    command_buffer: Arc<CommandBuffer>,

    render_routines: Vec<SceneHandle>,
}

impl RenderBackend {
    pub fn new(
        device: &Arc<Device>,
        queue: &Arc<Mutex<Queue>>,
        images: TargetMode<Vec<Arc<Image>>>,
    ) -> Result<Self> {
        let image_count = match &images {
            TargetMode::Mono(images) => images.len(),
            TargetMode::Stereo(left_images, right_images) => {
                debug_assert!(left_images.len() == right_images.len());
                left_images.len()
            }
        };

        // create a new command buffer
        let command_buffer = CommandBuffer::new_primary().build(device.clone(), queue.clone())?;

        Ok(RenderBackend {
            device: device.clone(),
            queue: queue.clone(),

            swapchain_images: Mutex::new(images),
            image_count: AtomicUsize::new(image_count),

            clear_color: RwLock::new(VkClearColorValue::float32([0.0, 0.0, 0.0, 1.0])),

            command_buffer,

            render_routines: Vec::new(),
        })
    }

    pub(crate) fn required_image_usage() -> VkImageUsageFlagBits {
        VK_IMAGE_USAGE_TRANSFER_DST_BIT.into()
    }
}

impl RenderBackend {
    pub fn device(&self) -> &Arc<Device> {
        &self.device
    }

    pub fn queue(&self) -> &Arc<Mutex<Queue>> {
        &self.queue
    }

    pub fn set_clear_color(&self, clear_color: [f32; 4]) {
        *self.clear_color.write().unwrap() = VkClearColorValue::float32(clear_color);
    }

    pub fn render(
        &mut self,
        world: &mut World,
        image_indices: TargetMode<usize>,
    ) -> Result<&Arc<CommandBuffer>> {
        // begin main command buffer
        let mut buffer_recorder = self.command_buffer.begin(VkCommandBufferBeginInfo::new(
            VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
        ))?;

        // clear the current swapchain image
        {
            let swapchain_images = self.swapchain_images.lock().unwrap();
            let target_layout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
            let clear_color = self.clear_color.read().unwrap().clone();

            match (&image_indices, swapchain_images.deref()) {
                (TargetMode::Mono(image_index), TargetMode::Mono(images)) => {
                    let swapchain_image = &images[*image_index];

                    Self::clear_image(
                        &mut buffer_recorder,
                        swapchain_image,
                        clear_color,
                        target_layout,
                    );
                }
                (
                    TargetMode::Stereo(left_image_index, right_image_index),
                    TargetMode::Stereo(left_images, right_images),
                ) => {
                    let left_image = &left_images[*left_image_index];
                    let right_image = &right_images[*right_image_index];

                    Self::clear_image(
                        &mut buffer_recorder,
                        left_image,
                        clear_color.clone(),
                        target_layout,
                    );

                    Self::clear_image(
                        &mut buffer_recorder,
                        right_image,
                        clear_color,
                        target_layout,
                    );
                }
                _ => panic!("Expected another target mode"),
            }
        }

        // make a call to the connected scenes
        self.render_routines
            .iter_mut()
            .try_for_each(|scene_handle| {
                (scene_handle.render)(
                    &mut buffer_recorder,
                    &*self.swapchain_images.lock().unwrap(),
                    &image_indices,
                    world,
                )
            })?;

        Ok(&self.command_buffer)
    }

    pub fn resize(
        &mut self,
        world: &mut World,
        images: TargetMode<Vec<Arc<Image>>>,
        width: u32,
        height: u32,
    ) -> Result<()> {
        self.image_count.store(
            match &images {
                TargetMode::Mono(images) => images.len(),
                TargetMode::Stereo(left_images, right_images) => {
                    debug_assert!(left_images.len() == right_images.len());
                    left_images.len()
                }
            },
            SeqCst,
        );

        self.render_routines
            .iter_mut()
            .try_for_each(|scene_handle| {
                (scene_handle.resize)(width as f32, height as f32, &images, world)
            })?;

        Ok(())
    }

    /// lower priority means it is more important
    pub fn add_render_routine<T: TScene>(&mut self, priority: u32) {
        self.render_routines.push(SceneHandle::new::<T>(priority));
        self.render_routines
            .sort_by_key(|scene_handle| scene_handle.priority);
    }

    pub fn remove_render_routine<T: TScene>(&mut self) {
        if let Some(index) = self
            .render_routines
            .iter()
            .position(|scene_handle| scene_handle.type_id == TypeId::of::<T>())
        {
            self.render_routines.remove(index as usize);
        }
    }

    // getter
    pub fn image_count(&self) -> usize {
        self.image_count.load(SeqCst)
    }

    pub fn images(&self) -> TargetMode<Vec<Arc<Image>>> {
        self.swapchain_images.lock().unwrap().clone()
    }
}

impl RenderBackend {
    #[inline]
    fn clear_image(
        buffer_recorder: &mut CommandBufferRecorder<'_>,
        image: &Arc<Image>,
        clear_color: VkClearColorValue,
        target_layout: VkImageLayout,
    ) {
        buffer_recorder.set_full_image_layout(image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
        buffer_recorder.clear_color_image(image, clear_color);
        buffer_recorder.set_full_image_layout(image, target_layout);
    }
}