AOBench written in Rust

aobench.rs

extern crate rand;

use std::fs;
use std::io::{BufWriter, Write};
use std::ops::{Add, AddAssign, Sub, Mul, Div};

const WIDTH: usize = 256;
const HEIGHT: usize = 256;
const NSUBSAMPLES: usize = 2;
const NAO_SAMPLES: usize = 8;

const NOHIT: f64 = 1.0e+17;
const M_PI: f64 = 3.1415926535;

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

impl Add for Vec3 {
    type Output = Vec3;
    fn add(self, other: Vec3) -> Vec3 {
        Vec3{x: self.x + other.x, y: self.y + other.y, z: self.z + other.z}
    }
}

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

impl Sub for Vec3 {
    type Output = Vec3;
    fn sub(self, other: Vec3) -> Vec3 {
        Vec3{x: self.x - other.x, y: self.y - other.y, z: self.z - other.z}
    }
}

impl Mul<f64> for Vec3 {
    type Output = Vec3;
    fn mul(self, other: f64) -> Vec3 {
        Vec3{x: self.x * other, y: self.y * other, z: self.z * other}
    }
}

impl Mul<Vec3> for f64 {
    type Output = Vec3;
    fn mul(self, other: Vec3) -> Vec3 {
        other * self
    }
}

impl Div<f64> for Vec3 {
    type Output = Vec3;
    fn div(self, other: f64) -> Vec3 {
        self * (1.0 / other)
    }
}

impl Clone for Vec3 {
    fn clone(&self) -> Vec3{*self}
}

impl Vec3 {
    fn normalize(&self) -> Vec3 {
        let length = self.dot(&self).sqrt();
        if length.abs() > 1.0e-17 {
            *self / length
        } else {
            *self
        }
    }

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

    fn cross(&self, v1: &Vec3) -> Vec3 {
        Vec3{x: self.y * v1.z - self.z * v1.y,
             y: self.z * v1.x - self.x * v1.z,
             z: self.x * v1.y - self.y * v1.x}
    }
}

struct Isect {
    t: f64,
    p: Vec3,
    n: Vec3,
}

struct Sphere {
    center: Vec3,
    radius: f64,
}

struct Plane {
    p: Vec3,
    n: Vec3,
}

struct Ray {
    org: Vec3,
    dir: Vec3,
}

const SPHERES: [Sphere; 3] = [
    Sphere{center: Vec3{x: -2.0, y: 0.0, z: -3.5},
           radius: 0.5},
    Sphere{center: Vec3{x: -0.5, y: 0.0, z: -3.0},
           radius: 0.5},
    Sphere{center: Vec3{x: 1.0, y: 0.0, z: -2.2},
           radius: 0.5},
];

const PLANE: Plane = Plane{p: Vec3{x: 0.0, y: -0.5, z: 0.0},
                           n: Vec3{x: 0.0, y: 1.0, z: 0.0}};

fn ray_sphere_intersect(isect: Isect, ray: &Ray, sphere: &Sphere) -> Isect {
    let rs = ray.org - sphere.center;
    let b = rs.dot(&ray.dir);
    let c = rs.dot(&rs) - sphere.radius * sphere.radius;
    let d = b * b - c;

    if d > 0.0 {
        let t = -b - d.sqrt();
        if t > 0.0 && t < isect.t {
            let p = ray.org + ray.dir * t;
            return Isect{t: t,
                         p: p,
                         n: (p - sphere.center).normalize()};
        }
    }
    return Isect{..isect};
}

fn ray_plane_intersect(isect: Isect, ray: &Ray, plane: &Plane) -> Isect {
    let d = -plane.p.dot(&plane.n);
    let v = ray.dir.dot(&plane.n);

    if v.abs() > 1.0e-17 {
        let t = -(ray.org.dot(&plane.n) + d) / v;

        if t > 0.0 && t < isect.t {
            return Isect{t: t,
                         p: ray.org + ray.dir * t,
                         n: Vec3{..plane.n}};
        }
    }
    return Isect{..isect};
}

fn ortho_basic(n: &Vec3) -> (Vec3, Vec3, Vec3) {
    let bz = Vec3{..*n};
    let by_: Vec3;

    if n.x < 0.6 && n.x > -0.6 {
        by_ = Vec3{x: 1.0, y: 0.0, z: 0.0};
    } else if n.y < 0.6 && n.y > -0.6 {
        by_ = Vec3{x: 0.0, y: 1.0, z: 0.0};
    } else if n.z < 0.6 && n.z > -0.6 {
        by_ = Vec3{x: 0.0, y: 0.0, z: 1.0};
    } else {
        by_ = Vec3{x: 1.0, y: 0.0, z: 0.0};
    }

    let bx = by_.cross(&bz).normalize();
    let by = bz.cross(&bx).normalize();
    return (bx, by, bz);
}

fn ambient_occlustion(isect: &Isect) -> Vec3 {
    let ntheta = NAO_SAMPLES;
    let nphi   = NAO_SAMPLES;
    let eps = 0.0001;
    let p = isect.p + isect.n * eps;

    let (bx, by, bz) = ortho_basic(&isect.n);

    let mut occlusion = 0.0;

    for _j in 0..ntheta {
        for _i in 0..nphi {
            let theta = rand::random::<f64>().sqrt();
            let phi = 2.0 * M_PI * rand::random::<f64>();

            let x = phi.cos() * theta;
            let y = phi.sin() * theta;
            let z = (1.0 - theta * theta).sqrt();

            // local -> global
            let rx = x * bx.x + y * by.x + z * bz.x;
            let ry = x * bx.y + y * by.y + z * bz.y;
            let rz = x * bx.z + y * by.z + z * bz.z;

            let ray = Ray{org: Vec3{..p},
                          dir: Vec3{x: rx, y: ry, z: rz}};
            let mut occ_isect = Isect{t: NOHIT,
                                      p: Vec3{x: 0.0, y: 0.0, z: 0.0},
                                      n: Vec3{x: 0.0, y: 0.0, z: 0.0}};

            occ_isect = ray_sphere_intersect(occ_isect, &ray, &SPHERES[0]);
            occ_isect = ray_sphere_intersect(occ_isect, &ray, &SPHERES[1]);
            occ_isect = ray_sphere_intersect(occ_isect, &ray, &SPHERES[2]);
            occ_isect = ray_plane_intersect(occ_isect, &ray, &PLANE);

            if occ_isect.t < NOHIT {
                occlusion += 1.0;
            }
        }
    }

    occlusion = ((ntheta as f64) * (nphi as f64) - occlusion) / ((ntheta * nphi) as f64);

    return Vec3{x: occlusion,
                y: occlusion,
                z: occlusion};
}

fn render(w: usize, h: usize, nsubsamples: usize) -> Vec<Vec3> {
    let mut img = vec![Vec3{x: 0.0, y: 0.0, z: 0.0}; w * h];
    for y in 0..h {
        for x in 0..w {
            let mut sum = Vec3{x: 0.0, y: 0.0, z: 0.0};
            for v in 0..nsubsamples {
                for u in 0..nsubsamples {
                    let px = ((x as f64) + ((u as f64) / (nsubsamples as f64)) - ((w as f64) / 2.0)) / ((w as f64) / 2.0);
                    let py = -((y as f64) + ((v as f64) / (nsubsamples as f64)) - ((h as f64) / 2.0)) / ((h as f64) / 2.0);

                    let ray = Ray{org: Vec3{x: 0.0, y: 0.0, z: 0.0},
                                  dir: Vec3{x: px, y: py, z: -1.0}.normalize()};
                    let mut isect = Isect{t: NOHIT,
                                          n: Vec3{x: 0.0, y: 0.0, z: 0.0},
                                          p: Vec3{x: 0.0, y: 0.0, z: 0.0},};

                    isect = ray_sphere_intersect(isect, &ray, &SPHERES[0]);
                    isect = ray_sphere_intersect(isect, &ray, &SPHERES[1]);
                    isect = ray_sphere_intersect(isect, &ray, &SPHERES[2]);
                    isect = ray_plane_intersect(isect, &ray, &PLANE);

                    if isect.t < NOHIT {
                        let col = ambient_occlustion(&isect);
                        sum += col;
                    }
                }
            }
            img[y * w + x] = sum / ((nsubsamples * nsubsamples) as f64);
        }
    }
    return img;
}

fn clamp(f: f64) -> u8 {
    let i = f * 255.5;
    if i < 0.0 {
        return 0;
    } else if i > 255.0 {
        return 255;
    } else {
        return i as u8;
    }
}

fn saveppm(fname: String, w: usize, h: usize, img: &Vec<Vec3>) {
    let mut f = BufWriter::new(fs::File::create(fname).unwrap());
    f.write(format!("P3\n{} {}\n255\n\n", w, h).as_bytes()).unwrap();
    for i in 0..h {
        for j in 0..w {
            if j != 0 {
                f.write(b" ").unwrap();
            }
            let c = &img[i * w + j];
            let r: u8 = clamp(c.x);
            let g: u8 = clamp(c.y);
            let b: u8 = clamp(c.z);
            f.write(format!("{} {} {}", r, g, b).as_bytes()).unwrap();
        }
        f.write(b"\n").unwrap();
    }
}

fn main() {
    let img = render(WIDTH, HEIGHT, NSUBSAMPLES);
    saveppm("ao.ppm".to_string(), WIDTH, HEIGHT, &img);
}