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(),
        }
    }
}

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<(u32, TypeId)>,
}

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(|(_, type_id)| {
                let scene: &mut dyn TScene = unsafe {
                    std::mem::transmute(destructure_traitobject::data_mut(
                        world.resources.get_mut_by_type_id_untyped(*type_id),
                    ))
                };

                scene.process(
                    &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(|(_, type_id)| {
                let scene: &mut dyn TScene = unsafe {
                    std::mem::transmute(destructure_traitobject::data_mut(
                        world.resources.get_mut_by_type_id_untyped(*type_id),
                    ))
                };

                scene.resize(width, height, &images)
            })?;

        Ok(())
    }

    /// lower priority means it is more important
    pub fn add_render_routine<T: TScene>(&mut self, priority: u32) {
        self.render_routines.push((priority, TypeId::of::<T>()));
        self.render_routines.sort_by_key(|(p, _)| *p);
    }

    pub fn remove_render_routine<T: TScene>(&mut self) {
        if let Some(index) = self
            .render_routines
            .iter()
            .find(|(_, type_id)| *type_id == TypeId::of::<T>())
        {
            self.render_routines.remove(index)
        }
    }

    // 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);
    }
}