Data Oriented Raytracing In a Weekend in Rust

Question

I've been learning about Data Oriented Programming, and when trying to implement this simple Raytracer, I decided to do it in this manner using Rust. This is also in the context of an optimisation course where I wanted to apply the principles we learned (for example: avoiding cache misses).

My question is mostly about the HittableList. (You can see the rest of the code on the repo : https://github.com/St0wy/raytracing)

src/ray.rs

use crate::geometry::sphere::Sphere;
use crate::geometry::xy_rectangle::XyRectangle;
use crate::geometry::xz_rectangle::XzRectangle;
use crate::geometry::yz_rectangle::YzRectangle;
use crate::material::Material;
use crate::{math::vec3::*, ray::Ray};
use rand::Rng;
use std::cmp::Ordering;
use tracy::zone;

#[derive(Debug)]
pub struct HitRecord<'a> {
    point: Vec3,
    normal: Vec3,
    t: f32,
    u: f32,
    v: f32,
    front_face: bool,
    material: &'a Material,
}

impl<'a> HitRecord<'a> {
    pub fn new(
        point: Vec3,
        t: f32,
        u: f32,
        v: f32,
        outward_normal: Vec3,
        ray_direction: &Vec3,
        material: &'a Material,
    ) -> Self {
        let front_face = ray_direction.dot(&outward_normal) < 0.0;
        let normal = if front_face {
            outward_normal
        } else {
            -outward_normal
        };

        Self {
            point,
            normal,
            t,
            u,
            v,
            front_face,
            material,
        }
    }

    pub fn normal(&self) -> &Vec3 {
        &self.normal
    }

    pub fn point(&self) -> &Vec3 {
        &self.point
    }

    pub fn material(&self) -> &Material {
        self.material
    }

    pub fn front_face(&self) -> bool {
        self.front_face
    }

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

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

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

pub trait Hittable {
    fn hit(&self, ray: &Ray, t_min: f32, t_max: f32) -> Option<HitRecord>;
    fn bounding_box(&self, time0: f32, time1: f32) -> Option<Aabb>;
}

#[derive(Copy, Clone, Debug)]
pub enum HittableObjectType {
    Sphere,
    MovingSphere,
    XyRectangle,
    XzRectangle,
    YzRectangle,
    AabbBox,
    BvhNode,
}

#[derive(Copy, Clone, Debug)]
pub struct HittableObjectIndex {
    pub object_type: HittableObjectType,
    pub index: usize,
}

impl HittableObjectIndex {
    pub fn new(object_type: HittableObjectType, index: usize) -> Self {
        HittableObjectIndex { object_type, index }
    }
}

pub struct HittableList {
    spheres: Vec<Sphere>,
    moving_spheres: Vec<MovingSphere>,
    xy_rectangles: Vec<XyRectangle>,
    xz_rectangles: Vec<XzRectangle>,
    yz_rectangles: Vec<YzRectangle>,
    aabb_boxes: Vec<AabbBox>,
    bvh_nodes: Vec<BvhNode>,
    first_node_index: usize,
}

impl HittableList {
    pub fn new() -> Self {
        Self {
            spheres: Vec::new(),
            moving_spheres: Vec::new(),
            xy_rectangles: Vec::new(),
            xz_rectangles: Vec::new(),
            yz_rectangles: Vec::new(),
            aabb_boxes: Vec::new(),
            bvh_nodes: Vec::new(),
            first_node_index: 0,
        }
    }

    pub fn add_sphere(&mut self, sphere: Sphere) {
        self.spheres.push(sphere);
    }

    pub fn add_moving_sphere(&mut self, moving_sphere: MovingSphere) {
        self.moving_spheres.push(moving_sphere);
    }

    pub fn add_xy_rectangle(&mut self, rectangle: XyRectangle) {
        self.xy_rectangles.push(rectangle);
    }

    pub fn add_xz_rectangle(&mut self, rectangle: XzRectangle) {
        self.xz_rectangles.push(rectangle);
    }

    pub fn add_yz_rectangle(&mut self, rectangle: YzRectangle) {
        self.yz_rectangles.push(rectangle);
    }

    pub fn add_aabb_box(&mut self, aabb_box: AabbBox) {
        self.aabb_boxes.push(aabb_box);
    }

    pub fn len(&self) -> usize {
        self.spheres.len()
            + self.moving_spheres.len()
            + self.xy_rectangles.len()
            + self.xz_rectangles.len()
            + self.yz_rectangles.len()
            + self.aabb_boxes.len()
    }

    pub fn hit_no_limit(&self, ray: &Ray) -> Option<HitRecord> {
        self.hit(ray, 0.001, f32::INFINITY)
    }

    pub fn hit_at(
        &self,
        hittable_object_index: &HittableObjectIndex,
        ray: &Ray,
        t_min: f32,
        t_max: f32,
    ) -> Option<HitRecord> {
        zone!();
        match hittable_object_index.object_type {
            HittableObjectType::BvhNode => self.hit_node(
                &self.bvh_nodes[hittable_object_index.index],
                ray,
                t_min,
                t_max,
            ),
            HittableObjectType::Sphere => {
                self.spheres[hittable_object_index.index].hit(ray, t_min, t_max)
            }
            HittableObjectType::MovingSphere => {
                self.moving_spheres[hittable_object_index.index].hit(ray, t_min, t_max)
            }
            HittableObjectType::XyRectangle => {
                self.xy_rectangles[hittable_object_index.index].hit(ray, t_min, t_max)
            }
            HittableObjectType::XzRectangle => {
                self.xz_rectangles[hittable_object_index.index].hit(ray, t_min, t_max)
            }
            HittableObjectType::YzRectangle => {
                self.yz_rectangles[hittable_object_index.index].hit(ray, t_min, t_max)
            }
            HittableObjectType::AabbBox => {
                self.aabb_boxes[hittable_object_index.index].hit(ray, t_min, t_max)
            }
        }
    }

    pub fn get_aabb(
        &self,
        hittable_object_index: HittableObjectIndex,
        time0: f32,
        time1: f32,
    ) -> Option<Aabb> {
        match hittable_object_index.object_type {
            HittableObjectType::Sphere => {
                self.spheres[hittable_object_index.index].bounding_box(time0, time1)
            }
            HittableObjectType::MovingSphere => {
                self.moving_spheres[hittable_object_index.index].bounding_box(time0, time1)
            }
            HittableObjectType::XyRectangle => {
                self.xy_rectangles[hittable_object_index.index].bounding_box(time0, time1)
            }
            HittableObjectType::XzRectangle => {
                self.xz_rectangles[hittable_object_index.index].bounding_box(time0, time1)
            }
            HittableObjectType::YzRectangle => {
                self.yz_rectangles[hittable_object_index.index].bounding_box(time0, time1)
            }
            HittableObjectType::BvhNode => {
                Some(self.bvh_nodes[hittable_object_index.index].aabb().clone())
            }
            HittableObjectType::AabbBox => {
                self.aabb_boxes[hittable_object_index.index].bounding_box(time0, time1)
            }
        }
    }

    pub fn is_empty(&self) -> bool {
        self.spheres.is_empty()
            && self.moving_spheres.is_empty()
            && self.xy_rectangles.is_empty()
            && self.xz_rectangles.is_empty()
            && self.yz_rectangles.is_empty()
            && self.aabb_boxes.is_empty()
    }

    fn hit_node(&self, node: &BvhNode, ray: &Ray, t_min: f32, t_max: f32) -> Option<HitRecord> {
        zone!();
        if !node.aabb().hit(ray, t_min, t_max) {
            return None;
        }

        let record_left_option = self.hit_at(node.left(), ray, t_min, t_max);
        let mut left_distance = t_max;
        let mut record = None;
        if let Some(record_left) = record_left_option {
            left_distance = record_left.t;
            record = Some(record_left);
        }

        let record_right = self.hit_at(node.right(), ray, t_min, left_distance);
        if let Some(record_right) = record_right {
            if left_distance < record_right.t {
                record
            } else {
                Some(record_right)
            }
        } else {
            record
        }
    }

    fn box_compare(
        &self,
        time0: f32,
        time1: f32,
        axis: usize,
    ) -> impl FnMut(&HittableObjectIndex, &HittableObjectIndex) -> Ordering + '_ {
        move |a, b| {
            let a_bbox = self.get_aabb(*a, time0, time1);
            let b_bbox = self.get_aabb(*b, time0, time1);
            if a_bbox.is_none() || b_bbox.is_none() {
                panic!("no bounding box in bvh node")
            }

            if a_bbox.unwrap().min()[axis] - b_bbox.unwrap().min()[axis] < 0.0 {
                Ordering::Less
            } else {
                Ordering::Greater
            }
        }
    }

    fn create_node(
        &mut self,
        hittables: &mut [HittableObjectIndex],
        time0: f32,
        time1: f32,
    ) -> HittableObjectIndex {
        let axis = rand::thread_rng().gen_range(0..3) as usize;

        let len = hittables.len();
        let (left, right) = match len {
            0 => panic!("0 Hittables provided to node creation"),
            1 => (hittables[0].clone(), hittables[0].clone()),
            2 => (hittables[0].clone(), hittables[1].clone()),
            _ => {
                hittables.sort_unstable_by(self.box_compare(time0, time1, axis));
                let mid = len / 2;
                (
                    self.create_node(&mut hittables[0..mid], time0, time1),
                    self.create_node(&mut hittables[mid..len], time0, time1),
                )
            }
        };

        let left_box = self.get_aabb(left, time0, time1);
        let right_box = self.get_aabb(right, time0, time1);

        if left_box.is_none() || right_box.is_none() {
            panic!("No bounding box in Bvh Node");
        }

        let aabb = Aabb::surrounding_box(left_box.unwrap(), right_box.unwrap());

        let node = BvhNode::new(left, right, aabb);
        self.bvh_nodes.push(node);

        HittableObjectIndex::new(HittableObjectType::BvhNode, self.bvh_nodes.len() - 1)
    }

    pub fn init_bvh_nodes(&mut self) {
        let mut hittables = Vec::new();

        for i in 0..self.spheres.len() {
            hittables.push(HittableObjectIndex::new(HittableObjectType::Sphere, i))
        }

        for i in 0..self.moving_spheres.len() {
            hittables.push(HittableObjectIndex::new(
                HittableObjectType::MovingSphere,
                i,
            ))
        }

        for i in 0..self.xy_rectangles.len() {
            hittables.push(HittableObjectIndex::new(HittableObjectType::XyRectangle, i));
        }

        for i in 0..self.xz_rectangles.len() {
            hittables.push(HittableObjectIndex::new(HittableObjectType::XzRectangle, i));
        }

        for i in 0..self.yz_rectangles.len() {
            hittables.push(HittableObjectIndex::new(HittableObjectType::YzRectangle, i));
        }

        for i in 0..self.aabb_boxes.len() {
            hittables.push(HittableObjectIndex::new(HittableObjectType::AabbBox, i));
        }

        let node = self.create_node(&mut hittables[..], 0.0, 1.0);
        self.first_node_index = node.index;
    }
}

fn get_objects_bounding_box<T: Hittable>(items: &Vec<T>, time0: f32, time1: f32) -> Option<Aabb> {
    if items.is_empty() {
        return None;
    }

    let mut temp_box: Aabb;
    let mut output_box = Aabb::empty();
    let mut first_box = true;

    for object in items.iter() {
        temp_box = object.bounding_box(time0, time1)?;

        output_box = if first_box {
            temp_box
        } else {
            Aabb::surrounding_box(output_box, temp_box)
        };

        first_box = false;
    }

    Some(output_box)
}

impl Hittable for HittableList {
    fn hit(&self, ray: &Ray, t_min: f32, t_max: f32) -> Option<HitRecord> {
        zone!();
        let first = self.bvh_nodes.get(self.first_node_index);

        if first.is_none() {
            panic!("There should be nodes in the hittable list.");
        }

        self.hit_node(&first.unwrap(), ray, t_min, t_max)
    }

    fn bounding_box(&self, time0: f32, time1: f32) -> Option<Aabb> {
        if self.is_empty() {
            return None;
        }

        let spheres_box = get_objects_bounding_box(&self.spheres, time0, time1);
        let moving_spheres_box = get_objects_bounding_box(&self.moving_spheres, time0, time1);
        let xy_rectangles_box = get_objects_bounding_box(&self.xy_rectangles, time0, time1);
        let xz_rectangles_box = get_objects_bounding_box(&self.xz_rectangles, time0, time1);
        let yz_rectangles_box = get_objects_bounding_box(&self.yz_rectangles, time0, time1);
        let aabb_box_box = get_objects_bounding_box(&self.aabb_boxes, time0, time1);

        let a = Aabb::opt_surrounding_box(spheres_box, moving_spheres_box);
        let b = Aabb::opt_surrounding_box(a, xy_rectangles_box);
        let c = Aabb::opt_surrounding_box(b, xz_rectangles_box);
        let d = Aabb::opt_surrounding_box(c, yz_rectangles_box);

        Aabb::opt_surrounding_box(d, aabb_box_box)
    }
}

As you can see, there is a lot of boilerplate / repetition because I wanted to avoid having a Vec<Box<dyn Hittable>> which would make the program slower and less cache friendly.

Would there be a way to design my code in a way that makes it easier to deal with new objects ?

Answer 1

The HittableList identifer is good, but it seems to dwell on a datastructure aspect that's not essential, and it's maybe not your favorite marketing term for getting folks to use your library. Consider rethinking the name. IDK, HittableFinder, maybe?

I find this [benchmark] line wonderfully clear:

    let mut world = HittableList::new();

It would never have occurred to me to write that if all I saw was the definition. Likely I would have gone with let mut hl = ... or something equally horrible.

---

It's obvious to one skilled in the art.

But as I approached this, I was sad I didn't see a comment that mentioned "bounding volume hierarchy" or "axis aligned bounding box".

That said, don't necessarily change anything. The Gentle Reader will eventually be clued in, I suppose.

If you're relying on some text for the math, then it would be useful for the code to cite the reference.

---

In len(), apparently bvh nodes are special somehow, as they don't contribute to the length. Ok.

---

    pub fn hit_no_limit(&self, ray: &Ray) -> Option<HitRecord> {
        self.hit(ray, 0.001, f32::INFINITY)

Uggh! Magic number. Please define epsilon or something.

---

    pub fn hit_at(
        &self,
        hittable_object_index: &HittableObjectIndex,
        ray: &Ray,
        t_min: f32,
        t_max: f32,
    ) -> Option<HitRecord> {
        zone!();

The reader needs to understand that those last two args represent parameterized distance along a ray. Either a comment or a literature citation is warranted. It's not yet clear how caller would choose sensible t values. Based on distance from eye to viewport, maybe?

I think the tracy zone!() is similar to matlab's tic / toc. It seems out of place in a production release. Either that, or we'd want to see an explanation of how this package supports an app developer's desire to understand where the cycles went.

The match dispatches to one of seven methods, each conforming to the same interface. So this seems like a good place to assign a function pointer, and then there's just a single call site that specifies the three args.

Similar argument almost works in get_aabb, except for the bvh node case. And I confess I have no idea what time{0,1} are about.

---

The is_empty definition seems tedious, given that we could just compare len() == 0.

---

hit_node would benefit from a word picture, or URL of a diagram, explaining how left / right fit into the orientation of the ray.

---

    fn box_compare(
                ...
                panic!("no bounding box in bvh node")

This doesn't seem very diagnostic. If I were a maintainer chasing a bug report, I would probably want to know details like the two times passed in.

---

In init_bvh_nodes, similar DRY argument as above. It seems there's an opportunity to iterate over a bunch of object types, with just a single call site for push / new.

---

fn get_objects_bounding_box<T: Hittable>(items: &Vec<T>, time0: f32, time1: f32) -> Option<Aabb> {
    ...
    let mut temp_box: Aabb;

It seems like we should be able to defer the temp_box declaration until down where it's assigned. Sigh.

Maybe initialize output_box to be element zero, and lose the flag, relying on a self comparison to be identity?

---

Time to take a step back.

Having read the whole thing, I still don't see the motivation for the sphere / moving sphere distinction. Just fill in a velocity of zero and be done with it, no?

The three rectangle orientations seem to play into this design consideration:

I wanted to avoid having a Vec<Box> which would make the program slower and less cache friendly.

Ok, fair enough. But in the codebase I would really like to see benchmark figures justifying such decisions, as there truly is an impact on the code. For example, I imagine len() is correct, but that's not obvious to me from just reading the source. There's certainly a place for these optimizations, but from my reading I am not yet convinced they're needed here. Comparative timings would help.

---

LGTM, ship it!