Trac3r-rust/src/vkprocessor/compute_image.rs

use vulkano::buffer::{BufferUsage, CpuAccessibleBuffer, DeviceLocalBuffer, ImmutableBuffer, BufferAccess};
use vulkano::command_buffer::{AutoCommandBufferBuilder, DynamicState};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSet, StdDescriptorPoolAlloc};
use vulkano::device::{Device, DeviceExtensions, QueuesIter, Queue};
use vulkano::instance::{Instance, InstanceExtensions, PhysicalDevice, QueueFamily};
use vulkano::pipeline::{ComputePipeline, GraphicsPipeline, GraphicsPipelineAbstract, GraphicsPipelineBuilder};
use vulkano::sync::{GpuFuture, FlushError};
use vulkano::sync;
use std::time::SystemTime;
use std::sync::Arc;
use std::ffi::CStr;
use std::path::PathBuf;
use shade_runner as sr;
use image::{DynamicImage, ImageBuffer};
use image::GenericImageView;
use vulkano::descriptor::pipeline_layout::PipelineLayout;
use image::GenericImage;
use shade_runner::{ComputeLayout, CompileError, FragLayout, FragInput, FragOutput, VertInput, VertOutput, VertLayout, CompiledShaders, Entry};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSetBuf, PersistentDescriptorSetImg, PersistentDescriptorSetSampler};
use shaderc::CompileOptions;
use vulkano::framebuffer::{Subpass, RenderPass, RenderPassAbstract, Framebuffer, FramebufferAbstract};
use vulkano::pipeline::shader::{GraphicsShaderType, ShaderModule, GraphicsEntryPoint, SpecializationConstants, SpecializationMapEntry};
use vulkano::swapchain::{Swapchain, PresentMode, SurfaceTransform, Surface, SwapchainCreationError, AcquireError};
use vulkano::swapchain::acquire_next_image;
use vulkano::image::swapchain::SwapchainImage;
use winit::{EventsLoop, WindowBuilder, Window, Event, WindowEvent};
use vulkano_win::VkSurfaceBuild;
use vulkano::pipeline::vertex::{SingleBufferDefinition, Vertex};
use vulkano::descriptor::PipelineLayoutAbstract;
use std::alloc::Layout;
use vulkano::pipeline::viewport::Viewport;
use image::ImageFormat;
use vulkano::image::immutable::ImmutableImage;
use vulkano::image::attachment::AttachmentImage;
use vulkano::image::{Dimensions, ImageUsage};
use vulkano::format::Format;
use vulkano::sampler::{Sampler, Filter, MipmapMode, SamplerAddressMode};
use image::flat::NormalForm::ColumnMajorPacked;
use image::Rgba;
use crate::vkprocessor::SimpleSpecializationConstants;


#[derive(Default, Debug, Clone)]
struct tVertex { position: [f32; 2] }

#[derive(Clone)]
pub struct ComputeImage {

    device: Arc<Device>,

    compute_graphics_swap_buffer: std::sync::Arc<vulkano::image::attachment::AttachmentImage>,

    image_buffer: Vec<u8>,

    xy: (u32, u32),

    pub rw_buffers: Vec<Arc<CpuAccessibleBuffer<[u8]>>>,
    pub settings_buffer: Arc<CpuAccessibleBuffer<[u32]>>,
}


impl ComputeImage {

    fn load_raw(filename: String) -> (Vec<u8>, (u32,u32)) {

        let project_root =
            std::env::current_dir()
                .expect("failed to get root directory");

        let mut compute_path = project_root.clone();
        compute_path.push(PathBuf::from("resources/images/"));
        compute_path.push(PathBuf::from(filename.clone()));

        let img = image::open(compute_path).expect("Couldn't find image");

        let xy = img.dimensions();

        let data_length = xy.0 * xy.1 * 4;
        let pixel_count = img.raw_pixels().len();

        let mut image_buffer = Vec::new();

        if pixel_count != data_length as usize {
            println!("Creating apha channel...");
            for i in img.raw_pixels().iter() {
                if (image_buffer.len() + 1) % 4 == 0 {
                    image_buffer.push(255);
                }
                image_buffer.push(*i);
            }
            image_buffer.push(255);
        } else {
            image_buffer = img.raw_pixels();
        }

        (image_buffer, xy)
    }

    pub fn new(device: Arc<Device>, image_filename: String) -> ComputeImage {

        let (image_buffer, xy) = ComputeImage::load_raw(image_filename);

        let compute_graphics_swap_buffer = {

            let mut usage = ImageUsage::none();
            usage.transfer_destination = true;
            usage.storage = true;

            AttachmentImage::with_usage(
                device.clone(),
                [xy.0, xy.1],
                Format::R8G8B8A8Uint,
                usage)
        };

        let data_length = xy.0 * xy.1 * 4;

        // Pull out the image data and place it in a buffer for the kernel to write to and for us to read from
        let write_buffer = {
            let mut buff = image_buffer.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
        };


        // Pull out the image data and place it in a buffer for the kernel to read from
        let read_buffer = {
            let mut buff = image_buffer.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(), BufferUsage::all(), data_iter).unwrap()
        };


        // A buffer to hold many i32 values to use as settings
        let settings_buffer = {
            let vec = vec![xy.0, xy.1];
            let mut buff = vec.iter();
            let data_iter =
                (0..2).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(device.clone(),
                                           BufferUsage::all(),
                                           data_iter).unwrap()
        };

        ComputeImage{
            device: device.clone(),
            compute_graphics_swap_buffer: compute_graphics_swap_buffer.unwrap(),
            image_buffer: image_buffer,
            xy: (0, 0),
            rw_buffers: vec![write_buffer, read_buffer],
            settings_buffer: settings_buffer
        }
    }

    pub fn get_swap_buffer(&mut self) -> Arc<AttachmentImage> {
        self.compute_graphics_swap_buffer.clone()
    }

    pub fn read_read_buffer(&mut self) -> ImageBuffer<Rgba<u8>, Vec<u8>>{

        let data_buffer_content = self.rw_buffers.get(0).unwrap().read().unwrap();
        ImageBuffer::from_fn(self.xy.0, self.xy.1, |x, y| {
            let r = data_buffer_content[((self.xy.0 * y + x) * 4 + 0) as usize] as u8;
            let g = data_buffer_content[((self.xy.0 * y + x) * 4 + 1) as usize] as u8;
            let b = data_buffer_content[((self.xy.0 * y + x) * 4 + 2) as usize] as u8;
            let a = data_buffer_content[((self.xy.0 * y + x) * 4 + 3) as usize] as u8;

            image::Rgba([r, g, b, a])
        })
    }

    pub fn get_descriptor_set(&self, compute_pipeline: std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>)
        -> Arc<PersistentDescriptorSet<std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>, ((((),
                PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u8]>>>),
                PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u8]>>>),
                PersistentDescriptorSetBuf<std::sync::Arc<vulkano::buffer::cpu_access::CpuAccessibleBuffer<[u32]>>>)>> {

        Arc::new(PersistentDescriptorSet::start(compute_pipeline.clone(), 0)
            .add_buffer(self.rw_buffers.get(0).unwrap().clone()).unwrap()
            .add_buffer(self.rw_buffers.get(1).unwrap().clone()).unwrap()
            .add_buffer(self.settings_buffer.clone()).unwrap()
            .build().unwrap())
    }
}