raytracing/src/vec3.rs

use core::f32::math::sqrt;
use is_close::default;
use rand::RngExt;
use std::{
    fmt::Display,
    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign},
};

#[derive(Copy, Clone)]
pub struct Vec3 {
    r: f32,
    g: f32,
    b: f32,
}

pub type Colour = Vec3;

impl Vec3 {
    pub fn new(r: f32, g: f32, b: f32) -> Self {
        Self { r: r, g: g, b: b }
    }

    pub fn nil() -> Self {
        Self {
            r: 0.,
            g: 0.,
            b: 0.,
        }
    }

    pub fn random() -> Self {
        let mut rng = rand::rng();
        Self {
            r: rng.random_range(-1.0..1.),
            g: rng.random_range(-1.0..1.),
            b: rng.random_range(-1.0..1.),
        }
    }

    pub fn random_unit() -> Self {
        loop {
            let v = Self::random();
            if v.length_squared() <= 1. {
                return v / v.length();
            }
        }
    }

    pub fn random_unit_hemisphere(n: &Self) -> Self {
        let v = Self::random_unit();
        if n.dot(&v) > 0.0 {
            v
        } else {
            -v
        }
    }

    pub fn x(&self) -> &f32 {
        &self.r
    }

    pub fn y(&self) -> &f32 {
        &self.g
    }

    pub fn z(&self) -> &f32 {
        &self.b
    }

    pub fn length(&self) -> f32 {
        sqrt(self.length_squared())
    }

    pub fn length_squared(&self) -> f32 {
        (self.r * self.r + self.g * self.g + self.b * self.b) as f32
    }

    pub fn dot(&self, other: &Self) -> f32 {
        self.r * other.r + self.g * other.g + self.b * other.b
    }

    pub fn cross(&self, other: &Self) -> Self {
        Self {
            r: self.g * other.b - self.b * other.g,
            g: self.b * other.r - self.r * other.b,
            b: self.r * other.g - self.g * other.r,
        }
    }

    pub fn make_unit(mut self) {
        self /= self.length();
    }

    pub fn get_unit(self) -> Self {
        self / self.length()
    }

    pub fn near_zero(&self) -> bool {
        default().is_close(self.r, 0.)
            && default().is_close(self.g, 0.)
            && default().is_close(self.b, 0.)
    }

    pub fn reflect(&self, n: &Self) -> Self {
        *self - 2. * self.dot(n) * n
    }

    pub fn refract(&self, n: &Self, etai_over_etat: f32) -> Self {
        let mut cos_theta = -self.dot(n);
        cos_theta = if cos_theta > 1.0 { 1.0 } else { cos_theta };

        let r_out_perp = etai_over_etat * (*self + cos_theta * n);
        let r_out_parr = -sqrt((1. - r_out_perp.length_squared()).abs()) * n;

        return r_out_perp + r_out_parr;
    }

    pub fn output(self) -> image::Rgb<u8> {
        // gamma correction
        let r = if self.r > 0. {
            sqrt(self.r).clamp(0., 1.)
        } else {
            0.
        };
        let g = if self.g > 0. {
            sqrt(self.g).clamp(0., 1.)
        } else {
            0.
        };
        let b = if self.b > 0. {
            sqrt(self.b).clamp(0.1, 1.)
        } else {
            0.
        };

        let ir = (255.599 * r) as u8;
        let ig = (255.599 * g) as u8;
        let ib = (255.599 * b) as u8;

        image::Rgb([ir, ig, ib])
    }

    pub fn clone(&self) -> Self {
        Self {
            r: self.r,
            g: self.g,
            b: self.b,
        }
    }
}

impl Neg for &Vec3 {
    type Output = Vec3;

    fn neg(self) -> Self::Output {
        -*self
    }
}

impl Neg for Vec3 {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self {
            r: -self.r,
            g: -self.g,
            b: -self.b,
        }
    }
}

impl Add<Vec3> for Vec3 {
    type Output = Self;

    fn add(self, _rhs: Self) -> Self {
        Self {
            r: self.r + _rhs.r,
            g: self.g + _rhs.g,
            b: self.b + _rhs.b,
        }
    }
}

impl Add<f32> for Vec3 {
    type Output = Self;

    fn add(self, f: f32) -> Self {
        Self {
            r: self.r + f,
            g: self.g + f,
            b: self.b + f,
        }
    }
}

impl AddAssign<Vec3> for Vec3 {
    fn add_assign(&mut self, other: Self) {
        *self = Self {
            r: self.r + other.r,
            g: self.g + other.g,
            b: self.b + other.b,
        };
    }
}

impl AddAssign<f32> for Vec3 {
    fn add_assign(&mut self, f: f32) {
        *self = Self {
            r: self.r + f,
            g: self.g + f,
            b: self.b + f,
        };
    }
}
impl Sub<&Vec3> for Vec3 {
    type Output = Self;

    fn sub(self, rhs: &Self) -> Self {
        Self {
            r: self.r - rhs.r,
            g: self.g - rhs.g,
            b: self.b - rhs.b,
        }
    }
}

impl Sub<Vec3> for Vec3 {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self {
        Self {
            r: self.r - rhs.r,
            g: self.g - rhs.g,
            b: self.b - rhs.b,
        }
    }
}

impl Sub<f32> for Vec3 {
    type Output = Self;

    fn sub(self, f: f32) -> Self {
        Self {
            r: self.r - f,
            g: self.g - f,
            b: self.b - f,
        }
    }
}

impl SubAssign<Vec3> for Vec3 {
    fn sub_assign(&mut self, rhs: Self) {
        *self = Self {
            r: self.r - rhs.r,
            g: self.g - rhs.g,
            b: self.b - rhs.b,
        };
    }
}

impl SubAssign<f32> for Vec3 {
    fn sub_assign(&mut self, f: f32) {
        *self = Self {
            r: self.r - f,
            g: self.g - f,
            b: self.b - f,
        };
    }
}

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

    fn mul(self, rhs: Self) -> Self::Output {
        Self {
            r: self.r * rhs.r,
            g: self.g * rhs.g,
            b: self.b * rhs.b,
        }
    }
}

impl Mul<Vec3> for u32 {
    type Output = Vec3;

    fn mul(self, rhs: Vec3) -> Self::Output {
        let f = self as f32;
        Vec3 {
            r: rhs.r * f,
            g: rhs.g * f,
            b: rhs.b * f,
        }
    }
}

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

    fn mul(self, f: f32) -> Self::Output {
        Self {
            r: self.r * f,
            g: self.g * f,
            b: self.b * f,
        }
    }
}

impl Mul<Vec3> for f32 {
    type Output = Vec3;

    fn mul(self, rhs: Vec3) -> Self::Output {
        Vec3 {
            r: rhs.r * self,
            g: rhs.g * self,
            b: rhs.b * self,
        }
    }
}

impl Mul<&Vec3> for f32 {
    type Output = Vec3;

    fn mul(self, rhs: &Vec3) -> Self::Output {
        Vec3 {
            r: rhs.r * self,
            g: rhs.g * self,
            b: rhs.b * self,
        }
    }
}

impl MulAssign<Vec3> for Vec3 {
    fn mul_assign(&mut self, rhs: Self) {
        *self = Self {
            r: self.r * rhs.r,
            g: self.g * rhs.g,
            b: self.b * rhs.b,
        };
    }
}

impl MulAssign<f32> for Vec3 {
    fn mul_assign(&mut self, f: f32) {
        *self = Self {
            r: self.r * f,
            g: self.g * f,
            b: self.b * f,
        };
    }
}

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

    fn div(self, rhs: Self) -> Self::Output {
        Self {
            r: self.r / rhs.r,
            g: self.g / rhs.g,
            b: self.b / rhs.b,
        }
    }
}
impl Div<i32> for Vec3 {
    type Output = Self;

    fn div(self, i: i32) -> Self::Output {
        let f: f32 = i as f32;
        Self {
            r: self.r / f,
            g: self.g / f,
            b: self.b / f,
        }
    }
}

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

    fn div(self, f: f32) -> Self::Output {
        Self {
            r: self.r / f,
            g: self.g / f,
            b: self.b / f,
        }
    }
}

impl DivAssign<Vec3> for Vec3 {
    fn div_assign(&mut self, rhs: Self) {
        *self = Self {
            r: self.r / rhs.r,
            g: self.g / rhs.g,
            b: self.b / rhs.b,
        };
    }
}

impl DivAssign<f32> for Vec3 {
    fn div_assign(&mut self, f: f32) {
        *self = Self {
            r: self.r / f,
            g: self.g / f,
            b: self.b / f,
        };
    }
}

impl PartialEq for Vec3 {
    fn eq(&self, other: &Self) -> bool {
        self.r == other.r && self.g == other.g && self.b == other.b
    }
}

impl Display for Vec3 {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "({}, {}, {})", self.r, self.g, self.b)
    }
}