Use bvh for scene shapes

ray-tracing-pbrt-scene/src/bvh.rs

@@ -2,8 +2,14 @@ use ray_tracing_core::{aabb::AABB, prelude::*};
 
 pub type Index = u32;
 
+pub trait BvhIntersection {
+    fn get_t(&self) -> Float;
+}
+
 pub trait BvhData {
-    type Intersect: AsRef<Float>;
+    type Intersect<'a>: BvhIntersection
+    where
+        Self: 'a;
 
     // fn get_slice(&mut self) -> &mut [Self::Sortable];
 
@@ -19,7 +25,8 @@ pub trait BvhData {
             .unwrap()
     }
 
-    fn intersect(&self, id: Index, ray: Ray, min: Float, max: Float) -> Option<Self::Intersect>;
+    fn intersect(&self, id: Index, ray: Ray, min: Float, max: Float)
+    -> Option<Self::Intersect<'_>>;
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -57,6 +64,10 @@ impl<T: BvhData> Bvh<T> {
         Self { data, bvh }
     }
 
+    pub fn inner(&self) -> &T {
+        &self.data
+    }
+
     fn build_bvh(data: &mut T, bvh: &mut Vec<Node>, node: usize, aabb: AABB, min: Index) {
         let (start, count) = if let Node::Leaf { start, count } = bvh[node] {
             (start, count)
@@ -112,7 +123,7 @@ impl<T: BvhData> Bvh<T> {
         ray: Ray,
         min: Float,
         max: Float,
-    ) -> Option<T::Intersect> {
+    ) -> Option<T::Intersect<'_>> {
         match self.bvh[node as usize] {
             Node::Inner {
                 left,
@@ -144,10 +155,10 @@ impl<T: BvhData> Bvh<T> {
                                     far,
                                     ray,
                                     min,
-                                    Float::min(max, *close_intersect.as_ref()),
+                                    Float::min(max, close_intersect.get_t()),
                                 )
                                 .filter(|far_intersect| {
-                                    *far_intersect.as_ref() < *close_intersect.as_ref()
+                                    far_intersect.get_t() < close_intersect.get_t()
                                 })
                             {
                                 Some(far_intersect)
@@ -164,10 +175,10 @@ impl<T: BvhData> Bvh<T> {
                 let mut intersection = None;
                 for i in start..(start + count) {
                     if let Some(t) = self.data.intersect(i, ray, min, max)
-                        && min <= *t.as_ref()
+                        && min <= t.get_t()
-                        && *t.as_ref() <= max
+                        && t.get_t() <= max
                         && !intersection.as_ref().is_some_and(
-                            |old_t: &<T as BvhData>::Intersect| *t.as_ref() > *old_t.as_ref(),
+                            |old_t: &<T as BvhData>::Intersect<'_>| t.get_t() > old_t.get_t(),
                         )
                     {
                         intersection.replace(t);
@@ -179,7 +190,7 @@ impl<T: BvhData> Bvh<T> {
         }
     }
 
-    pub fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<T::Intersect> {
+    pub fn intersect(&self, ray: Ray, min: Float, max: Float) -> Option<T::Intersect<'_>> {
         self.intersect_inner(0, ray, min, max)
     }
 }

ray-tracing-pbrt-scene/src/lib.rs

@@ -636,16 +636,8 @@ pub struct Pbrt<R: Rng> {
 }
 
 impl<R: Rng> Pbrt<R> {
-    fn new(settings: PbrtWorldSettings) -> Self {
+    fn new(settings: PbrtWorldSettings, scene: PbrtScene<R>) -> Self {
-        Self {
+        Self { settings, scene }
-            settings,
-            scene: PbrtScene {
-                shapes: Vec::new(),
-                infinite_light: Some(scene::PbrtInfiniteLight {
-                    color: Color::new(0.4, 0.45, 0.5),
-                }),
-            },
-        }
     }
 }
 
@@ -779,12 +771,12 @@ fn inner_parse_pbrt<R: Rng + std::fmt::Debug>(
 
     let (camera, camera_ctm) = camera.ok_or(miette!("A camera has to be specified"))?;
 
-    let mut pbrt = Pbrt::new(PbrtWorldSettings { camera, camera_ctm });
-
     context.ctm = vec![AffineTransform::identity()];
 
     // let mut context_material = vec![];
 
+    let mut shapes = Vec::new();
+
     while let Some(p) = parser.next(&context).transpose()? {
         match p {
             Statement::AttributeBegin => context.push(),
@@ -801,24 +793,24 @@ fn inner_parse_pbrt<R: Rng + std::fmt::Debug>(
             Statement::Shape(shape_type, shape_alpha) => {
                 // dbg!(&context);
                 if context.area_light.is_empty() {
-                    pbrt.scene.shapes.push(Shape {
+                    shapes.push(Shape::new(
-                        ctm: context.get_ctm(),
+                        context.get_ctm(),
-                        material: Either::A(Arc::clone(
+                        Either::A(Arc::clone(
                             context
                                 .material
                                 .last()
                                 .ok_or_else(|| miette!("No material specified"))?,
                         )),
-                        obj: shape_type,
+                        shape_type,
-                        alpha: shape_alpha,
+                        shape_alpha,
-                    });
+                    ));
                 } else {
-                    pbrt.scene.shapes.push(Shape {
+                    shapes.push(Shape::new(
-                        ctm: context.get_ctm(),
+                        context.get_ctm(),
-                        material: Either::B(context.area_light.last().unwrap().clone()),
+                        Either::B(context.area_light.last().unwrap().clone()),
-                        obj: shape_type,
+                        shape_type,
-                        alpha: shape_alpha,
+                        shape_alpha,
-                    });
+                    ));
                 }
             }
             Statement::CoordinateSystem(s) => {
@@ -868,6 +860,15 @@ fn inner_parse_pbrt<R: Rng + std::fmt::Debug>(
         }
     }
 
+    let pbrt = Pbrt::new(
+        PbrtWorldSettings { camera, camera_ctm },
+        PbrtScene {
+            shapes: Bvh::new(shapes, 1),
+            infinite_light: Some(scene::PbrtInfiniteLight {
+                color: Color::white(),
+            }),
+        },
+    );
     // dbg!(context);
 
     Ok(pbrt)

ray-tracing-pbrt-scene/src/scene.rs

@@ -5,14 +5,61 @@ use ray_tracing_core::{
     scene::{Intersection, Scene},
 };
 
-use crate::{material::PbrtMaterial, shape::Shape};
+use crate::{
+    bvh::{Bvh, BvhData, BvhIntersection},
+    material::PbrtMaterial,
+    shape::Shape,
+};
 
 #[derive(Debug)]
 pub struct PbrtScene<R: Rng> {
-    pub(crate) shapes: Vec<Shape<R>>,
+    pub(crate) shapes: Bvh<Vec<Shape<R>>>,
     pub(crate) infinite_light: Option<PbrtInfiniteLight>,
 }
 
+pub struct SceneIntersectionWrapper<'a, R: Rng>(
+    ray_tracing_core::scene::Intersection<R, UVMaterial<'a, R>, &'a dyn Light<R>>,
+);
+
+impl<R: Rng> BvhIntersection for SceneIntersectionWrapper<'_, R> {
+    fn get_t(&self) -> Float {
+        self.0.t()
+    }
+}
+
+impl<R: Rng + std::fmt::Debug> BvhData for Vec<Shape<R>> {
+    type Intersect<'a>
+        = SceneIntersectionWrapper<'a, R>
+    where
+        R: 'a;
+
+    fn len(&self) -> crate::bvh::Index {
+        self.len() as crate::bvh::Index
+    }
+
+    fn sort_range_dim(&mut self, id: std::ops::Range<crate::bvh::Index>, dim: u8) {
+        let get_key = |t: &Shape<R>| t.aabb.max()[dim as usize] - t.aabb.min()[dim as usize];
+        self[id.start as usize..id.end as usize]
+            .sort_by(|a, b| get_key(a).partial_cmp(&get_key(b)).unwrap());
+    }
+
+    fn get_aabb(&self, id: crate::bvh::Index) -> ray_tracing_core::prelude::AABB {
+        self[id as usize].aabb
+    }
+
+    fn intersect(
+        &self,
+        id: crate::bvh::Index,
+        ray: ray_tracing_core::prelude::Ray,
+        min: Float,
+        max: Float,
+    ) -> Option<Self::Intersect<'_>> {
+        self[id as usize]
+            .intersect(ray, min, max)
+            .map(|i| SceneIntersectionWrapper(i))
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct PbrtInfiniteLight {
     pub(crate) color: Color,
@@ -26,14 +73,23 @@ impl<R: Rng> Light<R> for PbrtInfiniteLight {
     }
 }
 
-#[derive(Debug)]
+pub struct UVMaterial<'a, R: Rng> {
-pub struct UVMaterial<'a, R: Rng + std::fmt::Debug> {
     pub(crate) u: Float,
     pub(crate) v: Float,
     pub(crate) material: &'a dyn PbrtMaterial<R>,
 }
 
-impl<R: Rng + std::fmt::Debug> Material<R> for UVMaterial<'_, R> {
+impl<'a, R: Rng> std::fmt::Debug for UVMaterial<'a, R> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("UVMaterial")
+            .field("u", &self.u)
+            .field("v", &self.v)
+            .field("material", &self.material)
+            .finish()
+    }
+}
+
+impl<R: Rng> Material<R> for UVMaterial<'_, R> {
     fn eval(
         &self,
         w_in: ray_tracing_core::prelude::Dir3,
@@ -77,20 +133,7 @@ impl<R: Rng + std::fmt::Debug> Scene<R> for PbrtScene<R> {
         min: ray_tracing_core::prelude::Float,
         max: ray_tracing_core::prelude::Float,
     ) -> Option<ray_tracing_core::scene::Intersection<R, Self::Mat<'_>, Self::Light<'_>>> {
-        let mut i = None;
+        let i = self.shapes.intersect(ray, min, max).map(|i| i.0);
-        for s in &self.shapes {
-            if let Some(new_i) = s.intersect(ray, min, max)
-                && i.as_ref().is_none_or(
-                    |i: &ray_tracing_core::scene::Intersection<
-                        R,
-                        Self::Mat<'_>,
-                        Self::Light<'_>,
-                    >| { i.t() > new_i.t() },
-                )
-            {
-                i = Some(new_i);
-            }
-        }
 
         i.or_else(|| {
             self.infinite_light.as_ref().map(|l| {

ray-tracing-pbrt-scene/src/shape.rs

@@ -4,7 +4,7 @@ use ray_tracing_core::{affine_transform::AffineTransform, prelude::*, scene::Int
 
 use crate::{
     AreaLight,
-    bvh::{Bvh, BvhData},
+    bvh::{Bvh, BvhData, BvhIntersection},
     either::Either,
     material::PbrtMaterial,
     scene::UVMaterial,
@@ -14,11 +14,40 @@ use crate::{
 #[allow(dead_code)]
 pub(crate) struct Shape<R: Rng> {
     pub(crate) ctm: AffineTransform,
+    pub(crate) aabb: AABB,
     pub(crate) material: Either<Arc<dyn PbrtMaterial<R>>, AreaLight>,
     pub(crate) obj: ShapeType,
     pub(crate) alpha: ShapeAlpha,
 }
 
+impl<R: Rng> Shape<R> {
+    pub(crate) fn new(
+        ctm: AffineTransform,
+        material: Either<Arc<dyn PbrtMaterial<R>>, AreaLight>,
+        obj: ShapeType,
+        alpha: ShapeAlpha,
+    ) -> Self {
+        let aabb = match &obj {
+            ShapeType::TriangleMesh(bvh) => bvh.inner().p[2..]
+                .iter()
+                .fold(AABB::new(bvh.inner().p[0], bvh.inner().p[1]), |a, &p| {
+                    a.extend(p)
+                }),
+            ShapeType::BilinearMesh { p, .. } => p[2..]
+                .iter()
+                .fold(AABB::new(p[0], p[1]), |a, &p| a.extend(p)),
+            _ => todo!(),
+        };
+        Shape {
+            ctm,
+            aabb,
+            material,
+            obj,
+            alpha,
+        }
+    }
+}
+
 #[derive(Debug)]
 #[allow(dead_code)]
 pub(crate) enum ShapeAlpha {
@@ -34,9 +63,9 @@ pub(crate) struct InnerIntersect {
     v: Float,
 }
 
-impl AsRef<Float> for InnerIntersect {
+impl BvhIntersection for InnerIntersect {
-    fn as_ref(&self) -> &Float {
+    fn get_t(&self) -> Float {
-        &self.t
+        self.t
     }
 }
 #[derive(Debug)]
@@ -50,7 +79,7 @@ pub(crate) struct TriangleMesh {
 }
 
 impl BvhData for TriangleMesh {
-    type Intersect = InnerIntersect;
+    type Intersect<'a> = InnerIntersect;
 
     fn len(&self) -> crate::bvh::Index {
         self.indices.len() as crate::bvh::Index
@@ -75,7 +104,7 @@ impl BvhData for TriangleMesh {
         ray: Ray,
         min: Float,
         max: Float,
-    ) -> Option<Self::Intersect> {
+    ) -> Option<Self::Intersect<'_>> {
         let _ = max;
         let _ = min;
         let v = [