Add Acceleration Structure Interface

ray-tracing-scene/src/acceleration_structure/mod.rs

@@ -0,0 +1,119 @@
+use std::sync::Arc;
+
+use rand::seq::SliceRandom;
+use ray_tracing_core::{
+    prelude::*,
+    scene::{Intersection, LightSample, Scene},
+};
+use sampling::sample_triangle;
+
+use crate::util::triangle_normal;
+
+pub mod triangle_bvh;
+
+pub trait AccelerationStructure<T> {
+    fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<(Float, Dir3, T)>;
+}
+
+#[derive(Debug)]
+pub struct ASMaterial<R: Rng> {
+    pub material: Option<Box<dyn Material<R>>>,
+    pub light: Option<Box<dyn Light<R>>>,
+}
+
+impl<R: Rng> ASMaterial<R> {
+    pub fn new_material<M: Material<R> + 'static>(material: M) -> Self {
+        Self {
+            material: Some(Box::new(material)),
+            light: None,
+        }
+    }
+    pub fn new_light<L: Light<R> + 'static>(light: L) -> Self {
+        Self {
+            material: None,
+            light: Some(Box::new(light)),
+        }
+    }
+    pub fn new_material_light<M: Material<R> + 'static, L: Light<R> + 'static>(
+        material: M,
+        light: L,
+    ) -> Self {
+        Self {
+            material: Some(Box::new(material)),
+            light: Some(Box::new(light)),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct AccelerationStructureScene<A, R: Rng> {
+    acceleration_structure: A,
+    materials: Arc<[ASMaterial<R>]>,
+    lights: Vec<(usize, [Pos3; 3])>,
+}
+
+impl<A: Clone, R: Rng> Clone for AccelerationStructureScene<A, R> {
+    fn clone(&self) -> Self {
+        Self {
+            acceleration_structure: self.acceleration_structure.clone(),
+            materials: Arc::clone(&self.materials),
+            lights: self.lights.clone(),
+        }
+    }
+}
+
+impl<A, R: Rng> AccelerationStructureScene<A, R> {
+    pub fn new(a: A, materials: Arc<[ASMaterial<R>]>, lights: Vec<(usize, [Pos3; 3])>) -> Self {
+        Self {
+            acceleration_structure: a,
+            materials,
+            lights,
+        }
+    }
+}
+
+impl<A: AccelerationStructure<u32>, R: Rng> Scene<R> for AccelerationStructureScene<A, R> {
+    fn intersect(
+        &self,
+        ray: Ray,
+        min: Float,
+        max: Float,
+    ) -> Option<ray_tracing_core::scene::Intersection<'_, R>> {
+        let (t, n, i) = self.acceleration_structure.intersect(ray, min, max)?;
+
+        let material = &self.materials[i as usize];
+
+        Some(Intersection::new(
+            t,
+            n,
+            material.material.as_deref(),
+            material.light.as_deref(),
+        ))
+    }
+
+    fn sample_light(
+        &self,
+        _w_in: Dir3,
+        _intersection: &ray_tracing_core::scene::Intersection<'_, R>,
+        rng: &mut R,
+    ) -> Option<ray_tracing_core::scene::LightSample<'_, R>> {
+        let t = self.lights.choose(rng);
+
+        if let Some(&(t, v)) = t {
+            let light = self.materials[t].light.as_ref().unwrap().as_ref();
+            let b = sample_triangle(rng.gen());
+
+            let n = triangle_normal(v);
+            let area = n.length() * 0.5;
+
+            Some(LightSample::new(
+                Pos3::from_barycentric(v, b),
+                1.0 / ((self.lights.len() as Float) * area),
+                n.normalize(),
+                light,
+            ))
+        } else {
+            None
+        }
+    }
+}

ray-tracing-scene/src/acceleration_structure/triangle_bvh.rs

@@ -0,0 +1,222 @@
+use ray_tracing_core::prelude::*;
+
+use crate::util::triangle_normal;
+
+use super::AccelerationStructure;
+
+type Index = u32;
+
+#[derive(Debug, Clone)]
+pub struct TriangleBVH<T> {
+    vertices: Vec<Pos3>,
+    triangles: Vec<Triangle<T>>,
+    bvh: Vec<Node>,
+}
+
+#[derive(Debug, Clone, Copy)]
+enum Node {
+    Inner {
+        left: Index,
+        left_aabb: AABB,
+        right: Index,
+        right_aabb: AABB,
+    },
+    Leaf {
+        start: Index,
+        count: Index,
+    },
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct Triangle<T> {
+    pub vertices: [Index; 3],
+    pub data: T,
+}
+
+impl<T> Triangle<T> {
+    pub fn new(vertices: [Index; 3], data: T) -> Self {
+        Self { vertices, data }
+    }
+}
+
+impl<T: Copy> TriangleBVH<T> {
+    pub fn new(vertices: Vec<Pos3>, mut triangles: Vec<Triangle<T>>) -> Self {
+        let mut bvh = vec![Node::Leaf {
+            start: 0,
+            count: triangles.len() as Index,
+        }];
+
+        let aabb = Self::calculate_aabb(&vertices, &triangles);
+
+        Self::build_bvh(&vertices, &mut triangles, &mut bvh, 0, aabb);
+
+        Self {
+            vertices,
+            triangles,
+            bvh,
+        }
+    }
+
+    fn calculate_aabb(vertices: &[Pos3], triangles: &[Triangle<T>]) -> AABB {
+        let mut aabb = AABB::new(
+            vertices[triangles[0].vertices[0] as usize],
+            vertices[triangles[0].vertices[1] as usize],
+        );
+        aabb = aabb.extend(vertices[triangles[0].vertices[2] as usize]);
+
+        for t in triangles.iter().skip(1) {
+            for v in t.vertices {
+                aabb = aabb.extend(vertices[v as usize]);
+            }
+        }
+
+        aabb
+    }
+
+    fn build_bvh(
+        vertices: &[Pos3],
+        triangles: &mut [Triangle<T>],
+        bvh: &mut Vec<Node>,
+        node: usize,
+        aabb: AABB,
+    ) {
+        let (start, count) = if let Node::Leaf { start, count } = bvh[node] {
+            (start, count)
+        } else {
+            unreachable!()
+        };
+
+        if count < 8 {
+            return;
+        }
+
+        let size = aabb.size();
+
+        let dim = if size.x() > Float::max(size.y(), size.z()) {
+            0
+        } else if size.y() > Float::max(size.x(), size.z()) {
+            1
+        } else {
+            2
+        };
+
+        let get_key = |t: &Triangle<T>| {
+            t.vertices
+                .iter()
+                .map(|&v| vertices[v as usize][dim])
+                .sum::<Float>()
+                / 3.0
+        };
+
+        triangles.sort_by(|a, b| get_key(a).partial_cmp(&get_key(b)).unwrap());
+
+        let (left_range, right_range) = triangles.split_at_mut((count / 2) as usize);
+
+        let left_aabb = Self::calculate_aabb(vertices, left_range);
+        let right_aabb = Self::calculate_aabb(vertices, right_range);
+
+        let i = bvh.len() as Index;
+        bvh[node] = Node::Inner {
+            left: i,
+            left_aabb,
+            right: i + 1,
+            right_aabb,
+        };
+
+        bvh.push(Node::Leaf {
+            start,
+            count: left_range.len() as Index,
+        });
+        bvh.push(Node::Leaf {
+            start: start + left_range.len() as Index,
+            count: right_range.len() as Index,
+        });
+
+        Self::build_bvh(vertices, left_range, bvh, i as usize, left_aabb);
+        Self::build_bvh(vertices, right_range, bvh, i as usize + 1, right_aabb);
+    }
+
+    fn intersect_bvh(
+        &self,
+        node: Index,
+        ray: Ray,
+        min: Float,
+        max: Float,
+    ) -> Option<(Float, Dir3, T)> {
+        match self.bvh[node as usize] {
+            Node::Inner {
+                left,
+                left_aabb,
+                right,
+                right_aabb,
+            } => {
+                let left_intersect = left_aabb.intersect_ray(ray, min, max);
+                let right_intersect = right_aabb.intersect_ray(ray, min, max);
+
+                match (left_intersect, right_intersect) {
+                    (None, None) => None,
+                    (None, Some(_)) => self.intersect_bvh(right, ray, min, max),
+                    (Some(_), None) => self.intersect_bvh(left, ray, min, max),
+                    (Some(l), Some(r)) => {
+                        let close;
+                        let far;
+                        if l < r {
+                            close = left;
+                            far = right;
+                        } else {
+                            close = right;
+                            far = left;
+                        }
+
+                        if let Some(close_intersect) = self.intersect_bvh(close, ray, min, max) {
+                            if let Some(far_intersect) = self
+                                .intersect_bvh(far, ray, min, Float::min(max, close_intersect.0))
+                                .filter(|far_intersect| far_intersect.0 < close_intersect.0)
+                            {
+                                Some(far_intersect)
+                            } else {
+                                Some(close_intersect)
+                            }
+                        } else {
+                            self.intersect_bvh(far, ray, min, max)
+                        }
+                    }
+                }
+            }
+            Node::Leaf { start, count } => {
+                let mut intersection = None;
+                for i in start..(start + count) {
+                    if let Some(t) = crate::util::triangle_intersection(
+                        ray,
+                        self.get_vertices(&self.triangles[i as usize]),
+                    ) {
+                        if min <= t
+                            && t <= max
+                            && !intersection.is_some_and(|(old_t, _, _)| t > old_t)
+                        {
+                            let n = triangle_normal(self.get_vertices(&self.triangles[i as usize]))
+                                .normalize();
+                            intersection.replace((t, n, self.triangles[i as usize].data));
+                        }
+                    }
+                }
+
+                intersection
+            }
+        }
+    }
+
+    fn get_vertices(&self, triangle: &Triangle<T>) -> [Pos3; 3] {
+        [
+            self.vertices[triangle.vertices[0] as usize],
+            self.vertices[triangle.vertices[1] as usize],
+            self.vertices[triangle.vertices[2] as usize],
+        ]
+    }
+}
+
+impl<T: Copy> AccelerationStructure<T> for TriangleBVH<T> {
+    fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<(Float, Dir3, T)> {
+        self.intersect_bvh(0, ray, min, max)
+    }
+}

ray-tracing-scene/src/examples/cornell2.rs

@@ -1,7 +1,7 @@
 use rand::Rng;
 use ray_tracing_core::{light::AreaLight, prelude::*, scene::Scene};
 use ray_tracing_material::{diffuse::DiffuseMaterial, mirror::Mirror};
-use std::{cell::OnceCell, fmt::Debug};
+use std::cell::OnceCell;
 
 use crate::{
     parse_obj::ObjData,

ray-tracing-scene/src/examples/mod.rs

@@ -23,9 +23,14 @@ pub fn example_scenes<R: Rng + Debug + 'static>() -> HashMap<&'static str, Box<d
         "stanford_dragon",
         Box::new(stanford_dragon::StanfordDragon::new()),
     );
+    map.insert(
+        "stanford_dragon_as",
+        Box::new(stanford_dragon_as::StanfordDragon::new()),
+    );
     map
 }
 
 pub mod basic_cornell;
 pub mod cornell2;
 pub mod stanford_dragon;
+pub mod stanford_dragon_as;

ray-tracing-scene/src/examples/stanford_dragon_as.rs

@@ -0,0 +1,126 @@
+use rand::Rng;
+use ray_tracing_core::{light::AreaLight, prelude::*, scene::Scene};
+use ray_tracing_material::{
+    microfacet::{BeckmannDistribution, Microfacet},
+    oren_nayar::OrenNayar,
+};
+use std::cell::OnceCell;
+
+use crate::{
+    acceleration_structure::{
+        triangle_bvh::{Triangle, TriangleBVH},
+        ASMaterial, AccelerationStructureScene,
+    },
+    parse_obj::ObjData,
+};
+
+use super::ExampleScene;
+
+pub struct StanfordDragon<R: Rng> {
+    scene: OnceCell<AccelerationStructureScene<TriangleBVH<u32>, R>>,
+}
+
+impl<R: Rng> StanfordDragon<R> {
+    pub fn new() -> Self {
+        Self {
+            scene: OnceCell::new(),
+        }
+    }
+}
+
+impl<R: Rng + 'static> ExampleScene<R> for StanfordDragon<R> {
+    fn get_scene(&self) -> Box<dyn Scene<R> + Sync> {
+        let s = self.scene.get_or_init(|| {
+            let obj = ObjData::new("ray-tracing-scene/obj/stanford_dragon.obj").unwrap();
+
+            let mut triangles = obj
+                .triangle_faces()
+                .map(|t| Triangle::new(t.v, 0))
+                .collect::<Vec<_>>();
+
+            let color = Color::new(0.2, 0.2, 0.9);
+            let materials = vec![
+                ASMaterial::new_material(Microfacet::new(BeckmannDistribution::new(0.01), color)),
+                ASMaterial::new_material(OrenNayar::new(0.5, Color::new(0.8, 0.8, 0.8))),
+                ASMaterial::new_material(OrenNayar::new(0.5, Color::new(0.9, 0.0, 0.0))),
+                ASMaterial::new_material(OrenNayar::new(0.5, Color::new(0.0, 0.9, 0.0))),
+                ASMaterial::new_light(AreaLight::new(Color::white() * 30.0)),
+            ];
+
+            let mut vertices = obj.vertices;
+            for v in &mut vertices {
+                *v = Pos3::new(v.x(), v.z() + 60.0, -v.y());
+            }
+
+            let side_length = 400.0;
+            let light_size = 150.0;
+            let light_offset = 0.01;
+            let offset = vertices.len() as u32;
+            vertices.extend_from_slice(&[
+                Pos3::new(side_length, 2.0 * side_length, side_length),
+                Pos3::new(side_length, 2.0 * side_length, -side_length),
+                Pos3::new(side_length, 0.0, side_length),
+                Pos3::new(side_length, 0.0, -side_length),
+                Pos3::new(-side_length, 2.0 * side_length, side_length),
+                Pos3::new(-side_length, 2.0 * side_length, -side_length),
+                Pos3::new(-side_length, 0.0, side_length),
+                Pos3::new(-side_length, 0.0, -side_length),
+                Pos3::new(light_size, 2.0 * side_length - light_offset, light_size),
+                Pos3::new(light_size, 2.0 * side_length - light_offset, -light_size),
+                Pos3::new(-light_size, 2.0 * side_length - light_offset, light_size),
+                Pos3::new(-light_size, 2.0 * side_length - light_offset, -light_size),
+            ]);
+
+            triangles.extend_from_slice(&[
+                Triangle::new([offset, 1 + offset, 2 + offset], 1),
+                Triangle::new([1 + offset, 3 + offset, 2 + offset], 1),
+                Triangle::new([offset, 4 + offset, 1 + offset], 1),
+                Triangle::new([1 + offset, 4 + offset, 5 + offset], 1),
+                Triangle::new([2 + offset, 3 + offset, 6 + offset], 1),
+                Triangle::new([3 + offset, 7 + offset, 6 + offset], 1),
+                Triangle::new([offset, 2 + offset, 4 + offset], 2),
+                Triangle::new([6 + offset, 4 + offset, 2 + offset], 2),
+                Triangle::new([1 + offset, 5 + offset, 3 + offset], 3),
+                Triangle::new([7 + offset, 3 + offset, 5 + offset], 3),
+                Triangle::new([8 + offset, 10 + offset, 9 + offset], 4),
+                Triangle::new([11 + offset, 9 + offset, 10 + offset], 4),
+            ]);
+
+            let lights: Vec<_> = triangles
+                .iter()
+                .filter(|f| materials[f.data as usize].light.is_some())
+                .map(|t| {
+                    (
+                        t.data as usize,
+                        [
+                            vertices[t.vertices[0] as usize],
+                            vertices[t.vertices[1] as usize],
+                            vertices[t.vertices[2] as usize],
+                        ],
+                    )
+                })
+                .collect();
+            let a = TriangleBVH::new(vertices, triangles);
+            AccelerationStructureScene::new(a, materials.into(), lights)
+        });
+
+        let n: AccelerationStructureScene<_, _> = s.clone();
+        Box::new(n)
+    }
+
+    fn get_camera_pos(&self) -> Pos3 {
+        Pos3::new(-150.0, 160.0, 250.0)
+    }
+
+    fn get_camera_look_at(&self) -> Pos3 {
+        Pos3::new(0.0, 60.0, 0.0)
+    }
+
+    fn get_camera_up(&self) -> Dir3 {
+        Dir3::up()
+    }
+
+    fn get_horizontal_fov(&self) -> Float {
+        Float::to_radians(90.0)
+    }
+}

ray-tracing-scene/src/lib.rs

@@ -1,4 +1,6 @@
+pub mod acceleration_structure;
 pub mod basic_scene;
 pub mod examples;
 pub mod parse_obj;
 pub mod triangle_bvh;
+pub mod util;

ray-tracing-scene/src/triangle_bvh.rs

@@ -226,12 +226,12 @@ impl<R: Rng> TriangleBVH<R> {
             count: triangles.len() as Index,
         }];
         let aabb = calculate_aabb(&vertices, &triangles);
-        build_bvh(&vertices, &mut triangles, &mut bvh, 0, aabb);
         let lights: Vec<_> = triangles
             .iter()
             .filter(|f| materials[f.material as usize].light.is_some())
             .map(Triangle::clone)
             .collect();
+        build_bvh(&vertices, &mut triangles, &mut bvh, 0, aabb);
         Self {
             vertices,
             bvh,

ray-tracing-scene/src/util.rs

@@ -0,0 +1,43 @@
+use ray_tracing_core::prelude::*;
+
+pub fn triangle_intersection(ray: Ray, v: [Pos3; 3]) -> Option<Float> {
+    let e1 = v[1] - v[0];
+    let e2 = v[2] - v[0];
+
+    let ray_cross_e2 = Dir3::cross(ray.dir(), e2);
+    let det = e1.dot(ray_cross_e2);
+
+    if det > -f32::EPSILON && det < f32::EPSILON {
+        return None; // This ray is parallel to this triangle.
+    }
+
+    let inv_det = 1.0 / det;
+    let s = ray.start() - v[0];
+    let u = inv_det * s.dot(ray_cross_e2);
+    if !(0.0..=1.0).contains(&u) {
+        return None;
+    }
+
+    let s_cross_e1 = s.cross(e1);
+    let v = inv_det * Dir3::dot(ray.dir(), s_cross_e1);
+    if v < 0.0 || u + v > 1.0 {
+        return None;
+    }
+    // At this stage we can compute t to find out where the intersection point is on the line.
+    let t = inv_det * e2.dot(s_cross_e1);
+
+    if t > Float::EPSILON {
+        // ray intersection
+        Some(t)
+    } else {
+        // This means that there is a line intersection but not a ray intersection.
+        None
+    }
+}
+
+pub fn triangle_normal(v: [Pos3; 3]) -> Dir3 {
+    let e1 = v[1] - v[0];
+    let e2 = v[2] - v[0];
+
+    Dir3::cross(e1, e2)
+}