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

use rand::seq::SliceRandom;
use ray_tracing_core::{
    prelude::*,
    scene::{Intersection, LightSample, Scene},
};
use sampling::sample_triangle;
use std::sync::Arc;

type Index = u32;

#[derive(Debug)]
pub struct TriangleBVH<R: Rng> {
    vertices: Vec<Pos3>,
    triangles: Vec<Triangle>,
    materials: Arc<[BVHMaterial<R>]>,
    bvh: Vec<Node>,
    lights: Vec<Triangle>,
}

impl<R: Rng> Clone for TriangleBVH<R> {
    fn clone(&self) -> Self {
        Self {
            vertices: self.vertices.clone(),
            triangles: self.triangles.clone(),
            materials: Arc::clone(&self.materials),
            bvh: self.bvh.clone(),
            lights: self.lights.clone(),
        }
    }
}

#[derive(Debug)]
pub struct BVHMaterial<R: Rng> {
    pub material: Option<Box<dyn Material<R>>>,
    pub light: Option<Box<dyn Light<R>>>,
}

impl<R: Rng> BVHMaterial<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)),
        }
    }

    pub fn new_empty() -> Self {
        Self {
            material: None,
            light: None,
        }
    }
}

#[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 {
    vertices: [Index; 3],
    material: Index,
}

impl Triangle {
    pub fn new(vertices: [Index; 3], material: Index) -> Self {
        Triangle { vertices, material }
    }
}

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
    }
}

fn triangle_normal(v: [Pos3; 3]) -> Dir3 {
    let e1 = v[1] - v[0];
    let e2 = v[2] - v[0];

    Dir3::cross(e1, e2)
}

fn calculate_aabb(vertices: &[Pos3], triangles: &[Triangle]) -> 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],
    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.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 = calculate_aabb(vertices, left_range);
    let right_aabb = 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,
    });

    build_bvh(vertices, left_range, bvh, i as usize, left_aabb);
    build_bvh(vertices, right_range, bvh, i as usize + 1, right_aabb);
}

impl<R: Rng> TriangleBVH<R> {
    pub fn new(
        vertices: Vec<Pos3>,
        mut triangles: Vec<Triangle>,
        materials: Vec<BVHMaterial<R>>,
    ) -> Self {
        let mut bvh = vec![Node::Leaf {
            start: 0,
            count: triangles.len() as Index,
        }];
        let aabb = calculate_aabb(&vertices, &triangles);
        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,
            triangles,
            materials: materials.into(),
            lights,
        }
    }

    fn get_vertices_from_index(&self, triangle: Index) -> [Pos3; 3] {
        self.get_vertices(&self.triangles[triangle as usize])
    }

    fn get_vertices(&self, triangle: &Triangle) -> [Pos3; 3] {
        [
            self.vertices[triangle.vertices[0] as usize],
            self.vertices[triangle.vertices[1] as usize],
            self.vertices[triangle.vertices[2] as usize],
        ]
    }

    fn intersect_bvh(
        &self,
        node: Index,
        ray: Ray,
        min: Float,
        max: Float,
    ) -> Option<(Index, Float)> {
        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.1))
                                .filter(|far_intersect| far_intersect.1 < close_intersect.1)
                            {
                                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) = triangle_intersection(ray, self.get_vertices_from_index(i)) {
                        if min <= t && t <= max && !intersection.is_some_and(|(_, old_t)| t > old_t)
                        {
                            intersection.replace((i, t));
                        }
                    }
                }

                intersection
            }
        }
    }
}

impl<R: Rng> Scene<R> for TriangleBVH<R> {
    fn intersect(
        &self,
        ray: Ray,
        min: Float,
        max: Float,
    ) -> Option<ray_tracing_core::scene::Intersection<'_, R>> {
        let (i, t) = self.intersect_bvh(0, ray, min, max)?;

        let triangle = self.triangles[i as usize];
        let material = &self.materials[triangle.material as usize];

        let n = triangle_normal(self.get_vertices_from_index(i));

        let area = n.length() * 0.5;

        Some(Intersection::new(
            t,
            n.normalize(),
            material.material.as_deref(),
            material.light.as_deref(),
            1.0 / ((self.lights.len() as Float) * area),
        ))
    }

    fn sample_light(
        &self,
        _w_in: Dir3,
        _intersection: &Intersection<'_, R>,
        rng: &mut R,
    ) -> Option<ray_tracing_core::scene::LightSample<'_, R>> {
        let t = self.lights.choose(rng);

        if let Some(t) = t {
            let light = self.materials[t.material as usize]
                .light
                .as_ref()
                .unwrap()
                .as_ref();
            let b = sample_triangle(rng.gen());

            let n = triangle_normal(self.get_vertices(t));
            let area = n.length() * 0.5;

            Some(LightSample::new(
                Pos3::from_barycentric(self.get_vertices(t), b),
                1.0 / ((self.lights.len() as Float) * area),
                n.normalize(),
                light,
            ))
        } else {
            None
        }
    }
}