rustic-zen 0.3.0

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

//! VulkanImage contains the full colour depth framebuffer and a GPU rendering Context, and provides functions to export it.
use std::sync::Arc;
use std::thread::{self, JoinHandle};

use crossbeam::channel::{bounded, Receiver, Sender};

use vulkano::buffer::{Buffer, BufferContents, BufferCreateInfo, BufferUsage, Subbuffer};
use vulkano::command_buffer::{
    CopyImageToBufferInfo, RenderPassBeginInfo, SubpassBeginInfo, SubpassContents,
};
use vulkano::format::Format;
use vulkano::image::view::ImageView;
use vulkano::image::Image;
use vulkano::memory::allocator::{AllocationCreateInfo, MemoryTypeFilter};
use vulkano::pipeline::graphics::color_blend::{
    AttachmentBlend, ColorBlendAttachmentState, ColorBlendState,
};
use vulkano::pipeline::graphics::input_assembly::{InputAssemblyState, PrimitiveTopology};
use vulkano::pipeline::graphics::multisample::MultisampleState;
use vulkano::pipeline::graphics::rasterization::RasterizationState;
use vulkano::pipeline::graphics::vertex_input::{Vertex, VertexDefinition};
use vulkano::pipeline::graphics::viewport::{Viewport, ViewportState};
use vulkano::pipeline::graphics::GraphicsPipelineCreateInfo;
use vulkano::pipeline::layout::PipelineDescriptorSetLayoutCreateInfo;
use vulkano::pipeline::{GraphicsPipeline, PipelineLayout, PipelineShaderStageCreateInfo};
use vulkano::render_pass::{Framebuffer, FramebufferCreateInfo, Subpass};

use super::vulkan_context::VulkanContext;
use super::{ExportImage, RenderImage};
use crate::ray::RayResult;

/// Provides a context for rendering an image on the GPU
pub struct VulkanImage {
    image_buffer: Subbuffer<[f32]>,
    width: usize,
    height: usize,
    lightpower: f32,
    sender: Sender<RayMsg>,
    render_thread: Option<JoinHandle<()>>,
}

#[derive(BufferContents, Vertex)]
#[repr(C)]
struct RayVertex {
    #[format(R32G32_SFLOAT)]
    position: [f32; 2],
    #[format(R32G32B32A32_SFLOAT)]
    color: [f32; 4],
}

#[derive(BufferContents)]
#[repr(C)]
struct ImageBounds {
    bounds: [f32; 2],
}

enum RayMsg {
    Ray(RayResult),
    EndOfFrame(Sender<()>),
    Shutdown,
}

impl VulkanImage {
    pub(crate) const BATCH_SIZE: usize = 1_000_000;

    /// Create new Image on the GPU for GPU rendering, with the given size.
    pub fn new(width: usize, height: usize) -> Self {
        let ctx = VulkanContext::new();

        let image_buffer = Buffer::new_slice::<f32>(
            ctx.memory_allocator(),
            BufferCreateInfo {
                usage: BufferUsage::TRANSFER_DST,
                ..Default::default()
            },
            AllocationCreateInfo {
                memory_type_filter: MemoryTypeFilter::PREFER_HOST
                    | MemoryTypeFilter::HOST_RANDOM_ACCESS,
                ..Default::default()
            },
            (width * height * 4) as u64,
        )
        .expect("failed to create buffer");

        let (tx, rx) = bounded(Self::BATCH_SIZE);

        let mut this = Self {
            image_buffer: image_buffer.clone(),
            width,
            height,
            lightpower: 0.0,
            sender: tx,
            render_thread: None,
        };

        let w = width as u32;
        let h = height as u32;
        let render_thread = thread::spawn(move || Self::draw(ctx, rx, w, h, image_buffer));

        this.render_thread = Some(render_thread);

        this
    }

    fn draw(
        mut ctx: VulkanContext,
        recv: Receiver<RayMsg>,
        width: u32,
        height: u32,
        dest: Subbuffer<[f32]>,
    ) {
        let image = ctx.new_framebuffer(width, height);
        let mut result_buffer = Vec::with_capacity(Self::BATCH_SIZE * 2);
        loop {
            match recv.recv().unwrap() {
                RayMsg::Ray(ray) => {
                    let (r, g, b) = ray.color::<f32>();

                    let start = RayVertex {
                        position: [ray.origin.x as f32, ray.origin.y as f32],
                        color: [r, g, b, 1.0],
                    };

                    let end = RayVertex {
                        position: [ray.termination.x as f32, ray.termination.y as f32],
                        color: [r, g, b, 1.0],
                    };

                    result_buffer.push(start);
                    result_buffer.push(end);

                    if result_buffer.len() == Self::BATCH_SIZE * 2 {
                        Self::draw_batch(
                            &mut ctx,
                            width,
                            height,
                            result_buffer.drain(..),
                            image.clone(),
                            dest.clone(),
                        );
                    }
                }
                RayMsg::EndOfFrame(ack) => {
                    Self::draw_batch(
                        &mut ctx,
                        width,
                        height,
                        result_buffer.drain(..),
                        image.clone(),
                        dest.clone(),
                    );
                    ctx.wait_gpu();
                    ack.send(()).unwrap();
                }
                RayMsg::Shutdown => {
                    break;
                }
            }
        }
    }

    fn draw_batch(
        ctx: &mut VulkanContext,
        width: u32,
        height: u32,
        rays: impl ExactSizeIterator<Item = RayVertex>,
        image: Arc<Image>,
        dest: Subbuffer<[f32]>,
    ) {
        let rays = rays.into_iter();
        let vertex_buffer = Buffer::from_iter(
            ctx.memory_allocator(),
            BufferCreateInfo {
                usage: BufferUsage::VERTEX_BUFFER,
                ..Default::default()
            },
            AllocationCreateInfo {
                memory_type_filter: MemoryTypeFilter::PREFER_DEVICE
                    | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
                ..Default::default()
            },
            rays,
        )
        .unwrap();

        let render_pass = vulkano::single_pass_renderpass!(ctx.device(),
            attachments: {
                color: {
                    format: Format::R32G32B32A32_SFLOAT,
                    samples: 1,
                    load_op: DontCare,
                    store_op: Store,
                },
            },
            pass: {
                color: [color],
                depth_stencil: {},
            },
        )
        .unwrap();

        let view = ImageView::new_default(image.clone()).unwrap();
        let frame_buffer = Framebuffer::new(
            render_pass.clone(),
            FramebufferCreateInfo {
                attachments: vec![view],
                ..Default::default()
            },
        )
        .unwrap();

        let vs = shaders::vertex::load(ctx.device()).expect("failed to create shader module");
        let fs = shaders::fragment::load(ctx.device()).expect("failed to create shader module");

        let viewport = Viewport {
            offset: [0.0, 0.0],
            extent: [width as f32, height as f32],
            depth_range: 0.0..=1.0,
        };

        let vs = vs.entry_point("main").unwrap();
        let fs = fs.entry_point("main").unwrap();

        let vertex_input_state = RayVertex::per_vertex()
            .definition(&vs.info().input_interface)
            .unwrap();

        let stages = [
            PipelineShaderStageCreateInfo::new(vs),
            PipelineShaderStageCreateInfo::new(fs),
        ];

        let layout = PipelineLayout::new(
            ctx.device(),
            PipelineDescriptorSetLayoutCreateInfo::from_stages(&stages)
                .into_pipeline_layout_create_info(ctx.device())
                .unwrap(),
        )
        .unwrap();

        let subpass = Subpass::from(render_pass.clone(), 0).unwrap();

        let pipeline = GraphicsPipeline::new(
            ctx.device(),
            None,
            GraphicsPipelineCreateInfo {
                stages: stages.into_iter().collect(),
                vertex_input_state: Some(vertex_input_state),
                input_assembly_state: Some(InputAssemblyState {
                    topology: PrimitiveTopology::LineList,
                    primitive_restart_enable: false,
                    ..Default::default()
                }),
                viewport_state: Some(ViewportState {
                    viewports: [viewport].into_iter().collect(),
                    ..Default::default()
                }),
                rasterization_state: Some(RasterizationState::default()),
                multisample_state: Some(MultisampleState::default()),
                color_blend_state: Some(ColorBlendState::with_attachment_states(
                    subpass.num_color_attachments(),
                    ColorBlendAttachmentState {
                        blend: Some(AttachmentBlend::additive()),
                        ..Default::default()
                    },
                )),
                subpass: Some(subpass.into()),
                ..GraphicsPipelineCreateInfo::layout(layout.clone())
            },
        )
        .unwrap();

        let num_vertices = vertex_buffer.len();

        let mut builder = ctx.command_builder();

        builder
            .begin_render_pass(
                RenderPassBeginInfo {
                    clear_values: vec![None],
                    ..RenderPassBeginInfo::framebuffer(frame_buffer)
                },
                SubpassBeginInfo {
                    contents: SubpassContents::Inline,
                    ..Default::default()
                },
            )
            .unwrap()
            .bind_pipeline_graphics(pipeline)
            .unwrap()
            .bind_vertex_buffers(0, vertex_buffer)
            .unwrap()
            .push_constants(
                layout,
                0,
                ImageBounds {
                    bounds: [width as f32, height as f32],
                },
            )
            .unwrap()
            .draw(
                num_vertices as u32,
                1,
                0,
                0, // 3 is the number of vertices, 1 is the number of instances
            )
            .unwrap()
            .end_render_pass(Default::default())
            .unwrap()
            .copy_image_to_buffer(CopyImageToBufferInfo::image_buffer(image, dest))
            .unwrap();

        let command_buffer = builder.build().unwrap();

        ctx.run_command_buffer(command_buffer);
    }
}

impl RenderImage for VulkanImage {
    fn draw_line(&self, ray: RayResult) {
        self.sender.send(RayMsg::Ray(ray)).unwrap();
    }

    fn prepare_render(&mut self, lightpower: f32) {
        self.lightpower = lightpower;
    }

    fn finish_render(&mut self) {
        let (tx, rx) = bounded(1);
        self.sender.send(RayMsg::EndOfFrame(tx)).unwrap();
        rx.recv().unwrap();
    }
}

impl ExportImage for VulkanImage {
    fn get_size(&self) -> (usize, usize) {
        (self.width, self.height)
    }

    fn get_lightpower(&self) -> f32 {
        self.lightpower
    }

    fn to_rgbaf32(&self) -> Vec<f32> {
        self.image_buffer
            .read()
            .expect("failed to read frambuffer")
            .to_vec()
    }
}

impl Drop for VulkanImage {
    fn drop(&mut self) {
        self.sender.send(RayMsg::Shutdown).unwrap();
    }
}

mod shaders {
    pub mod vertex {
        vulkano_shaders::shader! {
            ty: "vertex",
            src: "
                #version 460               
                layout(location = 0) in vec2 position;
                layout(location = 1) in vec4 color;
                
                layout(location=0) out vec4 outcolor;

                layout( push_constant ) uniform Constants
                {
                    vec2 image_bounds;
                } constants;

                void main() {
                    gl_Position = vec4((2.0 * (position.x / constants.image_bounds.x)) - 1.0, (2.0 * (position.y / constants.image_bounds.y)) - 1.0, 0.0, 1.0);
                    outcolor = color;
                }
            ",
        }
    }

    pub mod fragment {
        vulkano_shaders::shader! {
            ty: "fragment",
            src: r"
                #version 460
                layout(location=0) in vec4 incolor;

                layout(location = 0) out vec4 f_color;

                void main() {
                    f_color = incolor;
                }
            ",
        }
    }
}

#[cfg(test)]
mod tests {
    use super::VulkanImage;
    use crate::image::{ExportImage, RenderImage};
    use crate::ray::RayResult;

    use itertools::Itertools as _;

    #[test]
    fn traced_ray_is_not_black() {
        let mut i = VulkanImage::new(100, 100);
        i.prepare_render(0.0);
        i.draw_line(RayResult::new((10.0, 10.0), (90.0, 90.0), 620.0)); //red
        i.draw_line(RayResult::new((20.0, 10.0), (90.0, 80.0), 520.0)); //green
        i.draw_line(RayResult::new((10.0, 20.0), (80.0, 90.0), 470.0)); //blue
        i.finish_render();
        let mut r_count = 0.0;
        let mut g_count = 0.0;
        let mut b_count = 0.0;
        for (r, g, b, _) in i.to_rgbaf32().iter().tuples() {
            r_count += r;
            g_count += g;
            b_count += b;
        }
        assert_ne!(r_count, 0.0);
        assert_ne!(g_count, 0.0);
        assert_ne!(b_count, 0.0);
    }

    #[test]
    fn empty_image_is_black() {
        let i = VulkanImage::new(1920, 1080);
        let v = i.to_rgbaf32();
        for i in v.iter() {
            assert_eq!(*i, 0.0);
        }
    }

    #[test]
    fn output_len_u8() {
        let i = VulkanImage::new(1920, 1080);
        let v = i.to_rgba8(0, 1.0, 1.0);
        assert_eq!(v.len(), 1920 * 1080 * 4);
    }

    #[test]
    fn output_len_f32() {
        let i = VulkanImage::new(1920, 1080);
        let v = i.to_rgbaf32();
        assert_eq!(v.len(), 1920 * 1080 * 4);
    }
}