bs-trace 0.3.0

Free RayTracing software
Documentation
use crate::image::image::{Colour, Image};
use crate::trace::camera::Camera;
use crate::trace::prelude::{Vec3, World};
use crate::trace::sampling::SamplingStrategy;
use rand::RngCore;
use std::collections::HashMap;
use std::sync::mpsc::Receiver;
use std::sync::{mpsc, Arc, Mutex};
use std::thread;
use std::thread::JoinHandle;

pub struct SceneConfig<const W: usize, const H: usize> {
    camera: Arc<Camera<W, H>>,
    sampling_strategy: Arc<dyn SamplingStrategy + Send + Sync>,
    world: Arc<World>,
    num_threads: usize,
}

impl<const W: usize, const H: usize> SceneConfig<W, H> {
    pub fn new(
        camera: Camera<W, H>,
        sampling_strategy: Arc<dyn SamplingStrategy>,
        world: World,
        num_threads: usize,
    ) -> Self {
        let camera = Arc::new(camera);
        let world = Arc::new(world);
        SceneConfig {
            camera,
            sampling_strategy,
            world,
            num_threads,
        }
    }

    pub fn trace(&self) -> Image<W, H> {
        let (join_handles, rx) = self.spawn_threads(self.num_threads);
        let mut line_map: HashMap<usize, Vec<Colour>> = HashMap::new();
        let mut num_lines = 0usize;
        while let Ok((y, line_colours)) = rx.recv() {
            line_map.insert(y, line_colours);
            num_lines += 1;
            println!("{}/{}", num_lines, H);
        }
        Image::gen(|x, y| {
            let line = line_map.get(&y).unwrap();
            line[x]
        })
    }

    fn spawn_threads(
        &self,
        num_threads: usize,
    ) -> (Vec<JoinHandle<()>>, Receiver<(usize, Vec<Colour>)>) {
        let (tx, rx) = mpsc::channel::<(usize, Vec<Colour>)>();
        let line_numbers = Arc::new(Mutex::new(0..H));
        let join_handles: Vec<JoinHandle<()>> = (0..num_threads)
            .map(|thread_i| {
                let tx = tx.clone();
                let line_numbers = line_numbers.clone();
                let camera = self.camera.clone();
                let world = self.world.clone();
                let sampling_strategy = self.sampling_strategy.clone();

                thread::spawn(move || {
                    let mut rng: Box<dyn RngCore> = Box::new(rand::thread_rng());
                    loop {
                        let y = {
                            let mut line_numbers = line_numbers.lock().unwrap(); // TODO don't unwrap
                            line_numbers.next()
                        };

                        if let Some(y) = y {
                            let mut line_colours: Vec<Colour> = Vec::with_capacity(W);
                            for x in 0..W {
                                let ss = sampling_strategy.clone(); // TODO don't fucking clone this every time
                                let sample_rays =
                                    camera.sampling_rays_for_pixel(&mut rng, ss, x, y);
                                let num_samples = sample_rays.len();
                                let sample_sum: Vec3 = sample_rays
                                    .iter()
                                    .map(|ray| world.trace(&mut rng, ray).0)
                                    .sum();
                                let avg_colour: Colour = (sample_sum / (num_samples as f64))
                                    .map(|c| c.powf(1.0 / camera.gamma))
                                    .into();
                                line_colours.push(avg_colour);
                            }
                            tx.send((y, line_colours)).unwrap(); // TODO don't unwrap
                        } else {
                            break;
                        }
                    }
                })
            })
            .collect();

        (join_handles, rx)
    }
}