ray-tracing2/ray-tracing-egui/src/render.rs

use std::sync::{
    mpsc::{Receiver, Sender},
    Arc,
};

use rand::{rngs::SmallRng, SeedableRng};
use ray_tracing_core::{
    camera::{BasicCamera, Camera},
    prelude::*,
    renderer::ClassicalRenderer,
    scene::Scene,
};
use ray_tracing_renderer::{
    depth_renderer::DepthRenderer, mis::MIS, next_event_estimation::NextEventEstimation,
    path_tracer::PathTracer, path_tracer_importance::PathTracerImportance,
};
use ray_tracing_scene::examples;
use rayon::{
    iter::{IndexedParallelIterator, ParallelIterator},
    slice::ParallelSliceMut,
};
use vulkano::{
    buffer::{Buffer, BufferCreateInfo, BufferUsage, Subbuffer},
    memory::allocator::{AllocationCreateInfo, MemoryAllocator, MemoryTypeFilter},
};

#[derive(Debug, Clone)]
pub struct RenderSettings {
    pub width: u32,
    pub height: u32,
    pub scene: &'static str,
    pub renderer_id: usize,
    pub camera_pos: Pos3,
    pub camera_look_at: Pos3,
    pub camera_up: Dir3,
    pub camera_horizontal_fov: Float,
}

pub struct Data {
    pub width: u32,
    pub height: u32,
    pub samples: u32,
    pub buffer: Subbuffer<[f32]>,
}

fn render_pass<
    R: ClassicalRenderer<SmallRng, S, BasicCamera> + std::marker::Sync,
    S: Scene<SmallRng> + std::marker::Sync,
>(
    settings: &RenderSettings,
    renderer: &R,
    scene: &S,
    camera: &BasicCamera,
    samples: u32,
    buffer: &mut [Float],
) {
    buffer.par_chunks_mut(3).enumerate().for_each(|(i, c)| {
        let x = (i % settings.width as usize) as u32;
        let y = (i / settings.width as usize) as u32;

        let mut rng = SmallRng::seed_from_u64(
            (x + y * settings.width) as u64
                + (settings.width as u64 * settings.height as u64 * samples as u64),
        );
        let r = renderer.render_pixel(&scene, &camera, x, y, &mut rng);
        c[0] += r.r();
        c[1] += r.g();
        c[2] += r.b();
    });
}

#[allow(clippy::type_complexity)]
pub const RENDERER: [&str; 5] = [
    "Depth",
    "Path tracer",
    "Path tracer importance",
    "Next event estimation",
    "MIS",
];

fn render_pass_renderer<S: Scene<SmallRng> + Sync>(
    settings: &RenderSettings,
    scene: &S,
    camera: &BasicCamera,
    samples: u32,
    buffer: &mut [Float],
) {
    match settings.renderer_id {
        0 => {
            let r = DepthRenderer::new(settings.width, settings.height);
            render_pass(settings, &r, scene, camera, samples, buffer);
        }
        1 => {
            let r = PathTracer::new(settings.width, settings.height);
            render_pass(settings, &r, scene, camera, samples, buffer);
        }
        2 => {
            let r = PathTracerImportance::new(settings.width, settings.height);
            render_pass(settings, &r, scene, camera, samples, buffer);
        }
        3 => {
            let r = NextEventEstimation::new(settings.width, settings.height);
            render_pass(settings, &r, scene, camera, samples, buffer);
        }
        4 => {
            let r = MIS::new(settings.width, settings.height);
            render_pass(settings, &r, scene, camera, samples, buffer);
        }

        _ => panic!("Unknown renderer"),
    }
}

pub fn render_thread(
    s: RenderSettings,
    rx: Receiver<RenderSettings>,
    tx: Sender<Data>,
    allocator: Arc<impl MemoryAllocator>,
) {
    let mut settings = s;

    let mut buffer = vec![0.0; settings.width as usize * settings.height as usize * 3];

    let example_scenes = examples::example_scenes::<SmallRng>();

    let mut scene = example_scenes[settings.scene].get_scene();

    let mut camera = BasicCamera::from_look_at(
        settings.width,
        settings.height,
        settings.camera_pos,
        settings.camera_look_at,
        settings.camera_up,
        settings.camera_horizontal_fov,
    );

    let mut samples = 0;
    loop {
        while let Ok(s) = rx.try_recv() {
            println!("Settings changed.");
            settings = s;
            scene = example_scenes[settings.scene].get_scene();
            camera = BasicCamera::from_look_at(
                settings.width,
                settings.height,
                settings.camera_pos,
                settings.camera_look_at,
                settings.camera_up,
                settings.camera_horizontal_fov,
            );
            buffer = vec![0.0; settings.width as usize * settings.height as usize * 3];
            samples = 0;
        }

        match &scene {
            examples::ExampleSceneEnum::AccelerationStructureScene(s) => {
                render_pass_renderer(&settings, s, &camera, samples, &mut buffer)
            }
            examples::ExampleSceneEnum::TriangleBVH(s) => {
                render_pass_renderer(&settings, s, &camera, samples, &mut buffer)
            }
        }
        samples += 1;

        let data = Data {
            width: settings.width,
            height: settings.height,
            samples,
            buffer: Buffer::from_iter(
                allocator.clone(),
                BufferCreateInfo {
                    usage: BufferUsage::STORAGE_BUFFER,
                    ..Default::default()
                },
                AllocationCreateInfo {
                    memory_type_filter: MemoryTypeFilter::PREFER_DEVICE
                        | MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
                    ..Default::default()
                },
                buffer.iter().copied(),
            )
            .unwrap(),
        };

        let _ = tx.send(data);
    }
}