Rust port of https://medium.com/swlh/ray-tracing-from-scratch-in-python-41670e6a96f9

rayt.rs

extern crate image;
extern crate itertools_num;

use std::ops::{Add, AddAssign, Div, Mul, Sub};

use image::{ImageBuffer, Rgb, RgbImage};
use itertools_num::linspace;

static WIDTH: u32 = 1280;
static HEIGHT: u32 = 720;

static MAX_DEPTH: usize = 3;

#[derive(Clone, Copy, Debug, PartialEq)]
struct Vec3 {
    x: f64,
    y: f64,
    z: f64,
}

impl Vec3 {
    fn new(x: f64, y: f64, z: f64) -> Self {
        Vec3 { x: x, y: y, z: z }
    }

    fn norm(self) -> f64 {
        self.dot(self).sqrt()
    }

    fn dot(self, other: Vec3) -> f64 {
        self.x * other.x + self.y * other.y + self.z * other.z
    }

    fn normalize(self) -> Vec3 {
        self / self.norm()
    }

    fn clamp(self, min: f64, max: f64) -> Vec3 {
        let x = self.x.clamp(min, max);
        let y = self.y.clamp(min, max);
        let z = self.z.clamp(min, max);
        Vec3::new(x, y, z)
    }
}

impl Add for Vec3 {
    type Output = Self;

    fn add(self, other: Self) -> Self::Output {
        let x = self.x + other.x;
        let y = self.y + other.y;
        let z = self.z + other.z;
        Vec3::new(x, y, z)
    }
}

impl AddAssign for Vec3 {
    fn add_assign(&mut self, other: Self) {
        let x = self.x + other.x;
        let y = self.y + other.y;
        let z = self.z + other.z;
        *self = Vec3::new(x, y, z)
    }
}

impl Sub for Vec3 {
    type Output = Self;

    fn sub(self, other: Self) -> Self::Output {
        let x = self.x - other.x;
        let y = self.y - other.y;
        let z = self.z - other.z;
        Vec3::new(x, y, z)
    }
}

impl Mul for Vec3 {
    type Output = Self;

    fn mul(self, other: Self) -> Self::Output {
        let x = self.x * other.x;
        let y = self.y * other.y;
        let z = self.z * other.z;
        Vec3::new(x, y, z)
    }
}

impl Mul<f64> for Vec3 {
    type Output = Self;

    fn mul(self, other: f64) -> Self::Output {
        let x = self.x * other;
        let y = self.y * other;
        let z = self.z * other;
        Vec3::new(x, y, z)
    }
}

impl Div<f64> for Vec3 {
    type Output = Self;

    fn div(self, other: f64) -> Self::Output {
        let x = self.x / other;
        let y = self.y / other;
        let z = self.z / other;
        Vec3::new(x, y, z)
    }
}

impl From<Vec3> for Rgb<u8> {
    fn from(v: Vec3) -> Self {
        let v = v.clamp(0., 1.);
        let r = (v.x * 256.) as u8;
        let g = (v.y * 256.) as u8;
        let b = (v.z * 256.) as u8;
        Rgb([r, g, b])
    }
}

fn reflected(v: Vec3, axis: Vec3) -> Vec3 {
    v - axis * 2. * v.dot(axis)
}

#[derive(Debug)]
struct Sphere {
    center: Vec3,
    radius: f64,
    ambient: Vec3,
    diffuse: Vec3,
    specular: Vec3,
    shineness: f64,
    reflection: f64,
}

impl Sphere {
    fn new(
        center: Vec3,
        radius: f64,
        ambient: Vec3,
        diffuse: Vec3,
        specular: Vec3,
        shineness: f64,
        reflection: f64,
    ) -> Self {
        Sphere {
            center: center,
            radius: radius,
            ambient: ambient,
            diffuse: diffuse,
            specular: specular,
            shineness: shineness,
            reflection: reflection,
        }
    }

    fn intersect(&self, ray_origin: Vec3, ray_direction: Vec3) -> Option<f64> {
        let b = 2. * ray_direction.dot(ray_origin - self.center);
        let c = (ray_origin - self.center).norm().powi(2) - self.radius.powi(2);
        let delta = b.powi(2) - 4. * c;
        if delta > 0. {
            let t1 = (-b + delta.sqrt()) / 2.;
            let t2 = (-b - delta.sqrt()) / 2.;
            if t1 > 0. && t2 > 0. {
                return Some(t1.min(t2));
            }
        }
        None
    }
}

fn nearest_intersected_object(
    objects: &[Sphere],
    ray_origin: Vec3,
    ray_direction: Vec3,
) -> (Option<&Sphere>, f64) {
    let distances = objects
        .iter()
        .map(|object| object.intersect(ray_origin, ray_direction));
    let mut nearest_object: Option<&Sphere> = None;
    let mut min_distance = f64::INFINITY;
    for (index, distance) in distances.enumerate() {
        match distance {
            Some(d) => {
                if d < min_distance {
                    nearest_object = Some(&objects[index]);
                    min_distance = d;
                }
            }
            None => (),
        }
    }
    (nearest_object, min_distance)
}

#[derive(Debug)]
struct Light {
    position: Vec3,
    ambient: Vec3,
    diffuse: Vec3,
    specular: Vec3,
}

impl Light {
    fn new(position: Vec3, ambient: Vec3, diffuse: Vec3, specular: Vec3) -> Self {
        Light {
            position: position,
            ambient: ambient,
            diffuse: diffuse,
            specular: specular,
        }
    }
}

fn main() {
    let camera = Vec3::new(0., 0., 1.);
    let ratio = (WIDTH as f64) / (HEIGHT as f64);
    let screen = (-1., 1. / ratio, 1., -1. / ratio); // left, top, right, botom

    let light = Light::new(
        Vec3::new(5., 5., 5.),
        Vec3::new(1., 1., 1.),
        Vec3::new(1., 1., 1.),
        Vec3::new(1., 1., 1.),
    );

    let objects = vec![
        Sphere::new(
            Vec3::new(-0.2, 0., -1.),
            0.7,
            Vec3::new(0.1, 0., 0.),
            Vec3::new(0.7, 0., 0.),
            Vec3::new(1., 1., 1.),
            100.,
            0.5,
        ),
        Sphere::new(
            Vec3::new(0.1, -0.3, 0.),
            0.1,
            Vec3::new(0.1, 0., 0.1),
            Vec3::new(0.7, 0., 0.7),
            Vec3::new(1., 1., 1.),
            100.,
            0.5,
        ),
        Sphere::new(
            Vec3::new(-0.3, 0., 0.),
            0.15,
            Vec3::new(0., 0.1, 0.),
            Vec3::new(0., 0.6, 0.),
            Vec3::new(1., 1., 1.),
            100.,
            0.5,
        ),
        Sphere::new(
            Vec3::new(0., -9000., 0.),
            9000. - 0.7,
            Vec3::new(0.1, 0.1, 0.1),
            Vec3::new(0.6, 0.6, 0.6),
            Vec3::new(1., 1., 1.),
            100.,
            0.5,
        ),
    ];

    let mut image: RgbImage = ImageBuffer::new(WIDTH, HEIGHT);

    for (i, y) in linspace(screen.1, screen.3, HEIGHT as usize).enumerate() {
        for (j, x) in linspace(screen.0, screen.2, WIDTH as usize).enumerate() {
            let pixel = Vec3::new(x, y, 0.);
            let mut origin = camera;
            let mut direction = (pixel - origin).normalize();

            let mut color = Vec3::new(0., 0., 0.);
            let mut reflection = 1.;

            for _ in 0..MAX_DEPTH {
                let (nearest_object, min_distance) =
                    nearest_intersected_object(&objects, origin, direction);
                let nearest_object = match nearest_object {
                    Some(object) => object,
                    None => break,
                };
                let intersection = origin + direction * min_distance;

                let normal_to_surface = (intersection - nearest_object.center).normalize();
                let shifted_point = intersection + normal_to_surface * 1e-5;
                let intersection_to_light = (light.position - shifted_point).normalize();

                let (_, min_distance_to_light) =
                    nearest_intersected_object(&objects, shifted_point, intersection_to_light);
                let intersection_to_light_distance = (light.position - intersection).norm();
                let is_shadowed = min_distance_to_light < intersection_to_light_distance;

                if is_shadowed {
                    break;
                }

                let mut illumination = Vec3::new(0., 0., 0.);
                // ambient
                illumination += nearest_object.ambient * light.ambient;
                // diffuse
                illumination += nearest_object.diffuse
                    * light.diffuse
                    * intersection_to_light.dot(normal_to_surface);
                // specular
                let intersection_to_camera = (camera - intersection).normalize();
                let h = (intersection_to_light + intersection_to_camera).normalize();
                illumination += nearest_object.specular
                    * light.specular
                    * normal_to_surface.dot(h).powf(nearest_object.shineness / 4.);

                // reflection
                color += illumination * reflection;
                reflection *= nearest_object.reflection;

                origin = shifted_point;
                direction = reflected(direction, normal_to_surface);
            }

            image.put_pixel(j as u32, i as u32, Rgb::from(color));
        }

        print!("\r{}/{}", i + 1, HEIGHT);
    }

    println!("\nsave a image");
    image.save("image.png").unwrap();
}