Trac3r-rust/src/vkprocessor.rs

use vulkano::buffer::{BufferUsage, CpuAccessibleBuffer, DeviceLocalBuffer, ImmutableBuffer, BufferAccess};
use vulkano::command_buffer::{AutoCommandBufferBuilder, DynamicState};
use vulkano::descriptor::descriptor_set::{PersistentDescriptorSet, StdDescriptorPoolAlloc, PersistentDescriptorSetBuilder, FixedSizeDescriptorSetsPool, StdDescriptorPool};
use vulkano::descriptor::descriptor_set::collection::DescriptorSetsCollection;
use vulkano::device::{Device, DeviceExtensions, QueuesIter, Queue};
use vulkano::instance::{Instance, InstanceExtensions, PhysicalDevice, QueueFamily};
use vulkano::pipeline::{ComputePipeline, GraphicsPipeline, GraphicsPipelineAbstract};
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, DescriptorSet};
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, ImageAccess, ImageDimensions};
use vulkano::format::Format;
use vulkano::sampler::{Sampler, Filter, MipmapMode, SamplerAddressMode};
use image::flat::NormalForm::ColumnMajorPacked;

mod compute_kernel;
use crate::vkprocessor::compute_kernel::ComputeKernel;

mod shader_kernels;
use crate::vkprocessor::shader_kernels::ShaderKernels;
use vulkano::descriptor::descriptor::DescriptorDesc;

//
//#[derive(Clone)]
//struct ImageBuffers {
//    pub image_buffers : Vec<Box<ImageAccess + Send + Sync>>,
//}
//
//impl ImageBuffers {
//
//    pub fn new() -> ImageBuffers {
//        ImageBuffers {
//            image_buffers: vec![]
//        }
//    }
//    pub fn add_image(self) -> Self {
//
//        self
//    }
//}
//
//unsafe impl DescriptorSetsCollection for ImageBuffers {
//    fn into_vec(self) -> Vec<Box<DescriptorSet>> {
//        unimplemented!()
//    }
//
//    fn num_bindings_in_set(&self, set: usize) -> Option<usize> {
//        unimplemented!()
//    }
//
//    fn descriptor(&self, set: usize, binding: usize) -> Option<DescriptorDesc> {
//        unimplemented!()
//    }
//}















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

/// This method is called once during initialization, then again whenever the window is resized
fn window_size_dependent_setup(
    images: &[Arc<SwapchainImage<Window>>],
    render_pass: Arc<dyn RenderPassAbstract + Send + Sync>,
    dynamic_state: &mut DynamicState,
) -> Vec<Arc<dyn FramebufferAbstract + Send + Sync>> {
    let dimensions = images[0].dimensions();

    let viewport = Viewport {
        origin: [0.0, 0.0],
        dimensions: [dimensions.width() as f32, dimensions.height() as f32],
        depth_range: 0.0..1.0,
    };
    dynamic_state.viewports = Some(vec!(viewport));

    images.iter().map(|image| {
        Arc::new(
            Framebuffer::start(render_pass.clone())
                .add(image.clone()).unwrap()
                .build().unwrap()
        ) as Arc<dyn FramebufferAbstract + Send + Sync>
    }).collect::<Vec<_>>()
}

#[repr(C)]
#[derive(Default, Debug, Clone)]
struct SimpleSpecializationConstants {
    first_constant: i32,
    second_constant: u32,
    third_constant: f32,
}

unsafe impl SpecializationConstants for SimpleSpecializationConstants {
    fn descriptors() -> &'static [SpecializationMapEntry] {
        static DESCRIPTORS: [SpecializationMapEntry; 3] = [
            SpecializationMapEntry {
                constant_id: 0,
                offset: 0,
                size: 4,
            },
            SpecializationMapEntry {
                constant_id: 1,
                offset: 4,
                size: 4,
            },
            SpecializationMapEntry {
                constant_id: 2,
                offset: 8,
                size: 4,
            },
        ];

        &DESCRIPTORS
    }
}

pub struct VkProcessor<'a> {
    pub shader_kernels: Option<ShaderKernels>,
    pub compute_kernel: Option<ComputeKernel>,
    pub vertex_shader_path: PathBuf,
    pub fragment_shader_path: PathBuf,
    pub instance: Arc<Instance>,
    pub physical: PhysicalDevice<'a>,
    pub graphics_pipeline: Option<Arc<GraphicsPipelineAbstract + Sync + Send>>,
    pub compute_pipeline: Option<std::sync::Arc<ComputePipeline<PipelineLayout<shade_runner::layouts::ComputeLayout>>>>,
    pub device: Arc<Device>,
    pub queues: QueuesIter,
    pub queue: Arc<Queue>,
    pub compute_set: Option<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]>>>)>>>,
    pub img_set: Option<Arc<PersistentDescriptorSet<Arc<dyn GraphicsPipelineAbstract + Send + Sync>, ((((), PersistentDescriptorSetImg<Arc<ImmutableImage<Format>>>), PersistentDescriptorSetSampler), PersistentDescriptorSetImg<Arc<AttachmentImage>>)>>>,
    pub graphics_image_buffer: Option<Arc<ImmutableImage<Format>>>,
    pub image_buffer: Vec<u8>,
    pub compute_image_buffers: Vec<Arc<CpuAccessibleBuffer<[u8]>>>,
    pub settings_buffer: Option<Arc<CpuAccessibleBuffer<[u32]>>>,
//    pub swapchain: Option<Arc<Swapchain<Window>>>,
//    pub images: Option<Vec<Arc<SwapchainImage<Window>>>>,
    pub xy: (u32, u32),
    pub render_pass: Option<Arc<RenderPassAbstract + Send + Sync>>,
    pub vertex_buffer: Option<Arc<(dyn BufferAccess + std::marker::Send + std::marker::Sync + 'static)>>,
    pub dynamic_state: DynamicState,
    pub graphics_image_swap_buffer: Option<std::sync::Arc<vulkano::image::attachment::AttachmentImage>>,


    pub image_buffer_store : Vec<Box<ImageAccess + Send + Sync>>,

 //   pub image_buffers_obj : ImageBuffers,
}


impl<'a> VkProcessor<'a> {

    pub fn new(instance: &'a Arc<Instance>, surface: &'a Arc<Surface<Window>>) -> VkProcessor<'a> {
        let physical = PhysicalDevice::enumerate(instance).next().unwrap();

        let queue_family = physical.queue_families().find(|&q| {
            // We take the first queue that supports drawing to our window.
            q.supports_graphics() &&
                surface.is_supported(q).unwrap_or(false) &&
                q.supports_compute()
        }).unwrap();

        let device_ext = DeviceExtensions { khr_swapchain: true, ..DeviceExtensions::none() };

        let (device, mut queues) = Device::new(physical,
                                               physical.supported_features(),
                                               &device_ext,
                                               [(queue_family, 0.5)].iter().cloned()).unwrap();
        let queue = queues.next().unwrap();

        VkProcessor {
            shader_kernels: Option::None,
            compute_kernel: Option::None,
            vertex_shader_path: Default::default(),
            fragment_shader_path: Default::default(),
            instance: instance.clone(),
            physical: physical.clone(),
            graphics_pipeline: Option::None,
            compute_pipeline: Option::None,
            device: device.clone(),
            queue: queue,
            queues: queues,
            compute_set: Option::None,
            img_set: Option::None,
            graphics_image_buffer: None,
            image_buffer: Vec::new(),
            compute_image_buffers: Vec::new(),
            settings_buffer: Option::None,
            xy: (0, 0),
            render_pass: Option::None,
            vertex_buffer: Option::None,
            dynamic_state: DynamicState { line_width: None, viewports: None, scissors: None },
            graphics_image_swap_buffer: None,

            image_buffer_store: vec![],
            //image_buffers_obj: ImageBuffers::new(),
        }
    }


    pub fn compile_kernel(&mut self, filename: String) {
        self.compute_kernel   = Some(ComputeKernel::new(filename, self.device.clone()));
        self.compute_pipeline = Some(self.compute_kernel.clone().unwrap().get_pipeline());
    }

    pub fn compile_shaders(&mut self, filename: String, surface: &'a Arc<Surface<Window>>) {
        self.shader_kernels = Some(
            ShaderKernels::new(filename.clone(),
                               surface, self.queue.clone(),
                               self.physical,
                               self.device.clone())
        );
    }

    // On resizes we have to recreate the swapchain
    pub fn recreate_swapchain(&mut self, surface: &'a Arc<Surface<Window>>) {
        self.shader_kernels = Some(self.shader_kernels.take().unwrap().recreate_swapchain(surface));
    }

    pub fn load_buffers(&mut self, image_filename: String)
    {
        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(image_filename));

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

        self.xy = img.dimensions();

        let data_length = self.xy.0 * self.xy.1 * 4;
        let pixel_count = img.raw_pixels().len();
        println!("Pixel count {}", pixel_count);

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

        println!("Buffer length {}", self.image_buffer.len());
        println!("Size {:?}", self.xy);

        println!("Allocating Buffers...");

        // 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 = self.image_buffer.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(self.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 = self.image_buffer.iter();
            let data_iter = (0..data_length).map(|n| *(buff.next().unwrap()));
            CpuAccessibleBuffer::from_iter(self.device.clone(), BufferUsage::all(), data_iter).unwrap()
        };


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

        println!("Done");

        // Create the data descriptor set for our previously created shader pipeline
        let mut set =
            PersistentDescriptorSet::start(self.compute_pipeline.clone().unwrap().clone(), 0)
                .add_buffer(write_buffer.clone()).unwrap()
                .add_buffer(read_buffer.clone()).unwrap()
                .add_buffer(settings_buffer.clone()).unwrap();

        self.compute_set = Some(Arc::new(set.build().unwrap()));

        self.compute_image_buffers.push(write_buffer);
        self.compute_image_buffers.push(read_buffer);
        self.settings_buffer = Some(settings_buffer);


        let vertex_buffer = {
            vulkano::impl_vertex!(tVertex, position);

            CpuAccessibleBuffer::from_iter(self.device.clone(), BufferUsage::all(), [
                tVertex { position: [-1.0, -1.0 ] },
                tVertex { position: [-1.0,  1.0 ] },
                tVertex { position: [ 1.0,  1.0 ] },
                tVertex { position: [ 1.0, -1.0 ] },
            ].iter().cloned()).unwrap()
        };

        self.vertex_buffer = Some(vertex_buffer);

        let (texture, tex_future) = {

            ImmutableImage::from_iter(
                self.image_buffer.iter().cloned(),
                Dimensions::Dim2d { width: self.xy.0, height: self.xy.1 },
                Format::R8G8B8A8Srgb,
                self.queue.clone()
            ).unwrap()
        };

        let compute_transfer_image = {

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

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

        self.image_buffer_store.push(Box::new(texture.clone()));

        self.graphics_image_buffer = Some(texture.clone());
        self.graphics_image_swap_buffer = Some(compute_transfer_image.clone().unwrap());
    }




    // The image set is the containing object for all texture and image hooks.
    // todo, make this pull from the image_buffer_store
    fn get_image_set(&mut self) -> Box<DescriptorSet + Send + Sync> {

        let sampler = Sampler::new(self.device.clone(), Filter::Linear, Filter::Linear,
                                   MipmapMode::Nearest, SamplerAddressMode::Repeat, SamplerAddressMode::Repeat,
                                   SamplerAddressMode::Repeat, 0.0, 1.0, 0.0, 0.0).unwrap();

        let mut descriptor_sets = PersistentDescriptorSet::start(
            self.shader_kernels.clone().unwrap().graphics_pipeline.clone().unwrap().clone(), 0
        );

        let descriptor_sets = descriptor_sets.add_sampled_image(self.graphics_image_buffer.clone().unwrap().clone(), sampler.clone()).unwrap();

        let o : Box<DescriptorSet + Send + Sync> = Box::new(
        PersistentDescriptorSet::start(
            self.shader_kernels.clone().unwrap().graphics_pipeline.clone().unwrap().clone(), 0
        )
            .add_sampled_image(self.graphics_image_buffer.clone().unwrap().clone(), sampler.clone()).unwrap()
            .add_image(self.graphics_image_swap_buffer.clone().unwrap().clone()).unwrap()
            .build().unwrap());
        o
    }

    pub fn run(&mut self, surface: &'a Arc<Surface<Window>>, mut frame_future: Box<dyn GpuFuture>) -> Box<dyn GpuFuture> {

        let mut framebuffers = window_size_dependent_setup(&self.shader_kernels.clone().unwrap().swapchain_images.clone(),
                                                           self.shader_kernels.clone().unwrap().render_pass.clone(),
                                                           &mut self.dynamic_state);

        let mut recreate_swapchain = false;

        // The docs said to call this on each loop.
        frame_future.cleanup_finished();

        // Whenever the window resizes we need to recreate everything dependent on the window size.
        // In this example that includes the swapchain, the framebuffers and the dynamic state viewport.
        if recreate_swapchain {
            self.shader_kernels = Some(self.shader_kernels.clone().unwrap().recreate_swapchain(surface));
            framebuffers = window_size_dependent_setup(&self.shader_kernels.clone().unwrap().swapchain_images.clone(),
                                                       self.render_pass.clone().unwrap().clone(),
                                                       &mut self.dynamic_state);
            recreate_swapchain = false;
        }


        // This function can block if no image is available. The parameter is an optional timeout
        // after which the function call will return an error.
        let (image_num, acquire_future) = match vulkano::swapchain::acquire_next_image(self.shader_kernels.clone().unwrap().swapchain.clone(), None) {
            Ok(r) => r,
            Err(AcquireError::OutOfDate) => {
                recreate_swapchain = true;
                //continue;
                panic!("Weird thing");
            }
            Err(err) => panic!("{:?}", err)
        };

        // Specify the color to clear the framebuffer with i.e. blue
        let clear_values = vec!([0.0, 0.0, 1.0, 1.0].into());


        {
            // In order to draw, we have to build a *command buffer*. The command buffer object holds
            // the list of commands that are going to be executed.
            //
            // Building a command buffer is an expensive operation (usually a few hundred
            // microseconds), but it is known to be a hot path in the driver and is expected to be
            // optimized.
            //
            // Note that we have to pass a queue family when we create the command buffer. The command
            // buffer will only be executable on that given queue family.
            let mut v = Vec::new();
            v.push(self.vertex_buffer.clone().unwrap().clone());

            let command_buffer =
                AutoCommandBufferBuilder::primary_one_time_submit(self.device.clone(), self.queue.family())
                    .unwrap()

                    .dispatch([self.xy.0, self.xy.1, 1],
                              self.compute_pipeline.clone().unwrap().clone(),
                              self.compute_set.clone().unwrap().clone(), ()).unwrap()

                    .copy_buffer_to_image(self.compute_image_buffers.get(0).unwrap().clone(),
                                          self.graphics_image_swap_buffer.clone().unwrap()).unwrap()
                    .begin_render_pass(framebuffers[image_num].clone(), false, clear_values)
                    .unwrap()

                    .draw(self.shader_kernels.clone().unwrap().graphics_pipeline.clone().unwrap().clone(),
                          &self.dynamic_state.clone(), v,
                          vec![self.get_image_set()], ())
                    .unwrap()

                    .end_render_pass()
                    .unwrap()

                    .build().unwrap();

            let mut data_buffer_content = self.compute_image_buffers.get(0).unwrap().read().unwrap();
            let img = 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])
            });

            // Wait on the previous frame, then execute the command buffer and present the image
            let future = frame_future.join(acquire_future)
                .then_execute(self.queue.clone(), command_buffer).unwrap()
                .then_swapchain_present(self.queue.clone(), self.shader_kernels.clone().unwrap().swapchain.clone(), image_num)
                .then_signal_fence_and_flush();

            match future {
                Ok(future) => {
                    (Box::new(future) as Box<_>)
                }
                Err(FlushError::OutOfDate) => {
                    recreate_swapchain = true;
                    (Box::new(sync::now(self.device.clone())) as Box<_>)
                }
                Err(e) => {
                    println!("{:?}", e);
                    (Box::new(sync::now(self.device.clone())) as Box<_>)
                }
            }
        }
    }

//    pub fn read_image(&self) -> Vec<u8> {
//
//        // The buffer is sync'd so we can just read straight from the handle
//        let mut data_buffer_content = self.img_buffers.get(0).unwrap().read().unwrap();
//
//        println!("Reading output");
//
//        let mut image_buffer = Vec::new();
//
//        for y in 0..self.xy.1 {
//            for x in 0..self.xy.0 {
//
//                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_buffer.push(r);
//                image_buffer.push(g);
//                image_buffer.push(b);
//                image_buffer.push(a);
//            }
//        }
//
//        image_buffer
//    }

//    pub fn save_image(&self) {
//        println!("Saving output");
//
//        let img_data = self.read_image();
//
//        let img = ImageBuffer::from_fn(self.xy.0, self.xy.1, |x, y| {
//
//            let r = img_data[((self.xy.0 * y + x) * 4 + 0) as usize] as u8;
//            let g = img_data[((self.xy.0 * y + x) * 4 + 1) as usize] as u8;
//            let b = img_data[((self.xy.0 * y + x) * 4 + 2) as usize] as u8;
//            let a = img_data[((self.xy.0 * y + x) * 4 + 3) as usize] as u8;
//
//            image::Rgba([r, g, b, a])
//        });
//
//        img.save(format!("output/{}.png", SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs()));
//    }
}