raychel-rs/src/vec3.rs

use std::{ops::{Add, Sub, AddAssign, SubAssign, Mul, Div, MulAssign, DivAssign, Neg}};

use crate::rtweekend::{random_f64, random_range_f64};

#[derive(Debug, Clone, Copy)]
pub struct Vec3 {
    pub elements: [f64; 3],
}

impl Vec3 {
    pub const fn new(x: f64, y: f64, z: f64) -> Self {
        Vec3 { elements: [x, y, z] }
    }

    pub fn new_empty() -> Self {
        Vec3::new(0.0, 0.0, 0.0)
    }

    pub fn x(&self) -> f64 {
        self.elements[0]
    }

    pub fn y(&self) -> f64 {
        self.elements[1]
    }

    pub fn z(&self) -> f64 {
        self.elements[2]
    }

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

    pub fn length_squared(&self) -> f64 {
        self.dot(self)
    }

    pub fn dot(&self, rhs: &Self) -> f64 {
        self.x() * rhs.x()
        + self.y() * rhs.y()
        + self.z() * rhs.z()
    }

    pub fn cross(&self, rhs: &Self) -> Self {
        Self {
            elements: [
                self.y() * rhs.z() + self.z() * rhs.y(),
                self.z() * rhs.x() + self.x() * rhs.z(),
                self.x() * rhs.y() + self.y() * rhs.x()
            ]
        }
    }

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

    pub fn random() -> Self {
        Vec3::new(random_f64(), random_f64(), random_f64())
    }

    pub fn random_range(min: f64, max: f64) -> Self {
        Vec3::new(random_range_f64(min, max), random_range_f64(min, max), random_range_f64(min, max))
    }

    pub fn near_zero(&self) -> bool {
        let s = 1e-8;
        (self.x().abs() < s) && (self.y().abs() < s) && (self.z().abs() < s)
    }
}

impl Add for Vec3 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self {
        Self {
            elements: [
                self.x() + rhs.x(),
                self.y() + rhs.y(),
                self.z() + rhs.z()
            ]
        }
    }
}

impl Sub for Vec3 {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self {
            elements: [
                self.x() - rhs.x(),
                self.y() - rhs.y(),
                self.z() - rhs.z()
            ]
        }
    }
}

impl Neg for Vec3 {
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self {
            elements: [
                -self.x(),
                -self.y(),
                -self.z()
            ]
        }
    }
}

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

    fn mul(self, rhs: f64) -> Self::Output {
        Self {
            elements: [
                self.x() * rhs,
                self.y() * rhs,
                self.z() * rhs
            ]
        }
    }
}

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

    fn mul(self, rhs: Self) -> Self::Output {
        Self {
            elements: [
                self.x() * rhs.x(),
                self.y() * rhs.y(),
                self.z() * rhs.z()
            ]
        }
    }
}

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

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

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

    fn div(self, rhs: f64) -> Self::Output {
        Self {
            elements: [
                self.x() / rhs,
                self.y() / rhs,
                self.z() / rhs
            ]
        }
    }
}


impl AddAssign for Vec3 {
    fn add_assign(&mut self, rhs: Self) {
        *self = Self {
            elements: [
                self.x() + rhs.x(),
                self.y() + rhs.y(),
                self.z() + rhs.z()
            ]
        }
    }
}

impl SubAssign for Vec3 {
    fn sub_assign(&mut self, rhs: Self) {
        *self = Self {
            elements: [
                self.x() - rhs.x(),
                self.y() - rhs.y(),
                self.z() - rhs.z()
            ]
        }
    }
}

impl MulAssign<f64> for Vec3 {
    fn mul_assign(&mut self, rhs: f64) {
        *self = Self {
            elements: [
                self.x() * rhs,
                self.y() * rhs,
                self.z() * rhs
            ]
        }
    }
}

impl DivAssign<f64> for Vec3 {
    fn div_assign(&mut self, rhs: f64) {
        *self = Self {
            elements: [
                self.x() / rhs,
                self.y() / rhs,
                self.z() / rhs
            ]
        }
    }
}

pub type Color = Vec3;
pub type Point3 = Vec3;

pub fn random_in_unit_sphere() -> Vec3 {
    loop {
        let p = Vec3::random_range(-1.0, 1.0);
        if p.length_squared() >= 1.0 {
            continue;
        }
        return p;
    }
}

pub fn random_unit_vector() -> Vec3 {
    random_in_unit_sphere().unit_vector()
}

pub fn random_in_hemisphere(normal: &Vec3) -> Vec3 {
    let in_unit_sphere = random_in_unit_sphere();
    if in_unit_sphere.dot(&normal) > 0.0 {
        return in_unit_sphere;
    } else {
        return -in_unit_sphere;
    }
}

pub fn reflect(v: &Vec3, n: &Vec3) -> Vec3 {
    *v - 2.0 * v.dot(n) * *n
}

pub fn refract(uv: &Vec3, n: &Vec3, etai_over_etat: f64) -> Vec3 {
    let cos_theta = f64::min(-uv.dot(&n), 1.0);
    let r_out_perp = etai_over_etat * (*uv + cos_theta * *n);
    let r_out_parallel = -f64::sqrt(f64::abs(1.0 - r_out_perp.length_squared())) * *n;
    
    r_out_perp + r_out_parallel
}