Trac3r-rust/src/main.rs

#![allow(dead_code)]
#![allow(unused_variables)]
#![allow(unused_mut)]

extern crate cgmath;
extern crate image;
extern crate nalgebra as na;
extern crate quick_xml;
extern crate rand;
extern crate sfml;
extern crate time;

use image::{DynamicImage, GenericImage, GenericImageView, Pixel, SubImage};
use sfml::graphics::*;
use sfml::graphics::{
    Color, RenderTarget, RenderWindow,
};
use sfml::system::*;
use sfml::system::Vector2 as sfVec2;
use sfml::window::*;
use sfml::window::{Event, Key, Style};

use vulkano::buffer::{BufferUsage, CpuAccessibleBuffer, DeviceLocalBuffer, ImmutableBuffer, BufferAccess};
use vulkano::command_buffer::AutoCommandBufferBuilder;
use vulkano::descriptor::descriptor_set::PersistentDescriptorSet;
use vulkano::device::{Device, DeviceExtensions};
use vulkano::instance::{Instance, InstanceExtensions, PhysicalDevice};
use vulkano::pipeline::ComputePipeline;
use vulkano::sync::GpuFuture;
use vulkano::sync;
use std::sync::Arc;
use std::{fs, mem, iter, ptr};
use std::path::PathBuf;
use std::result;

use crate::input::Input;
use crate::slider::Slider;
use crate::timer::Timer;
use na::DimAdd;
use std::time::{SystemTime, Duration};
use shade_runner as sr;
use std::ffi::CStr;
use std::ptr::write;

mod slider;
mod timer;
mod input;
mod util;

fn main() {

    // Load up the input image, determine some details
    let mut img = image::open("resources/images/test2.png").unwrap();
    let xy = img.dimensions();
    let data_length = xy.0 * xy.1 * 4;
    let mut image_buffer = Vec::new();

    {
        // Create the vulkan instance, device, and device queue
        let instance = Instance::new(None, &InstanceExtensions::none(), None).unwrap();
        let physical = PhysicalDevice::enumerate(&instance).next().unwrap();
        let queue_family = physical.queue_families().find(|&q| q.supports_compute()).unwrap();
        let (device, mut queues) = Device::new(physical,
                                               physical.supported_features(),
                                               &DeviceExtensions::none(),
                                               [(queue_family, 0.5)].iter().cloned()).unwrap();
        let queue = queues.next().unwrap();

        println!("Device initialized");

        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/shaders/add.compute"));

        let shader = sr::load_compute(compute_path).expect("Failed to compile");
        let vulkano_entry = sr::parse_compute(&shader).expect("failed to parse");

        let x = unsafe {
            vulkano::pipeline::shader::ShaderModule::from_words(device.clone(), &shader.compute)
        }.unwrap();

        // Compile the shader and add it to a pipeline
        let pipeline = Arc::new({
            unsafe {
                ComputePipeline::new(device.clone(), &x.compute_entry_point(
                    CStr::from_bytes_with_nul_unchecked(b"main\0"),
                    vulkano_entry.compute_layout), &()
                ).unwrap()
            }
        });

        let pixel_count = img.raw_pixels().len();
        println!("Pixel count {}", pixel_count);

        if pixel_count != data_length as usize {
            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();
        }

        println!("Buffer length {}", image_buffer.len());
        println!("Size {:?}", 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 = 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()
            };

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

            let mut set = Arc::new(set.build().unwrap());

            println!("Running Kernel...");
            // The command buffer I think pretty much serves to define what runs where for how many times
            let command_buffer = AutoCommandBufferBuilder::primary_one_time_submit(device.clone(), queue.family()).unwrap()
                .dispatch([xy.0, xy.1, 1], pipeline.clone(), set.clone(), ()).unwrap()
                .build().unwrap();

            // Create a future for running the command buffer and then just fence it
            let future = sync::now(device.clone())
                .then_execute(queue.clone(), command_buffer).unwrap()
                .then_signal_fence_and_flush().unwrap();

            // I think this is redundant and returns immediately
            future.wait(None).unwrap();

            println!("Done running kernel");

            // The buffer is sync'd so we can just read straight from the handle
            let mut data_buffer_content = write_buffer.read().unwrap();

            println!("Reading output");

            for y in 0..xy.1 {
                for x in 0..xy.0 {
                    let r = data_buffer_content[((xy.0 * y + x) * 4 + 0) as usize] as u8;
                    let g = data_buffer_content[((xy.0 * y + x) * 4 + 1) as usize] as u8;
                    let b = data_buffer_content[((xy.0 * y + x) * 4 + 2) as usize] as u8;
                    let a = data_buffer_content[((xy.0 * y + x) * 4 + 3) as usize] as u8;

                    *image_buffer.get_mut(((xy.0 * y + x) * 4 + 0) as usize).unwrap() = r;
                    *image_buffer.get_mut(((xy.0 * y + x) * 4 + 1) as usize).unwrap() = g;
                    *image_buffer.get_mut(((xy.0 * y + x) * 4 + 2) as usize).unwrap() = b;
                    *image_buffer.get_mut(((xy.0 * y + x) * 4 + 3) as usize).unwrap() = a;

                    img.put_pixel(x, y, image::Rgba([r, g, b, a]))
                }
            }
        }

        // Currently bringing all this start shit outta scope to see if it stops my gpu from screaming
        println!("Saving output");
        img.save(format!("output/{}.png", SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs()));
    }

    let mut window = RenderWindow::new(
        (900, 900),
        "Custom drawable",
        Style::CLOSE,
        &Default::default(),
    );

    let mut timer = Timer::new();
    let mut input = Input::new();

    let font = Font::from_file("resources/fonts/sansation.ttf").unwrap();

    let mut bg_texture = Texture::new(xy.0, xy.1).unwrap();
    bg_texture.update_from_pixels(image_buffer.as_slice(), xy.0, xy.1, 0, 0);

    let mut background_sprite = Sprite::with_texture(&bg_texture);
    background_sprite.set_position((0., 0.));

    let mut slider = Slider::new(40.0, None);

    let step_size:            f32 = 0.005;
    let mut elapsed_time:     f32;
    let mut delta_time:       f32;
    let mut accumulator_time: f32 = 0.0;
    let mut current_time:     f32 = timer.elap_time();

    while window.is_open() {

        while let Some(event) = window.poll_event() {
            match event {
                Event::Closed => return,
                Event::KeyPressed { code, .. } => {
                    if code == Key::Escape {
                        return;
                    }
                }
                _ => {}
            }
            input.ingest(&event)
        }

        if input.is_held(Key::W) {
        }
        if input.is_held(Key::A) {
        }
        if input.is_held(Key::S) {
        }
        if input.is_held(Key::D) {
        }

        elapsed_time = timer.elap_time();
        delta_time = elapsed_time - current_time;
        current_time = elapsed_time;
        if delta_time > 0.02 {
            delta_time = 0.02;
        }
        accumulator_time += delta_time;

        while (accumulator_time - step_size) >= step_size {
            accumulator_time -= step_size;
        }

        window.clear(&Color::BLACK);

        window.draw(&background_sprite);
        window.draw(&slider);


        window.display();

    }
}