monte-carlo/src/bvh.rs

// ┏┓ ╻ ╻╻ ╻
// ┣┻┓┃┏┛┣━┫
// ┗━┛┗┛ ╹ ╹


use crate::geometry::Vec3;
use crate::geometry::Axis;
use crate::geometry::Ray;
use crate::mesh::Face;
use crate::mesh::Mesh;

pub struct Intersection<'a> {
    pub t: f32,
    pub fr: FaceReference<'a>
}

impl PartialOrd for Intersection<'_> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        return self.t.partial_cmp(&other.t);
    }
}

impl PartialEq for Intersection<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.t == other.t
    }
}

/**
 * Axis-Aligned Bounding Box.
 */
pub struct AABB {
    pub min: Vec3,
    pub max: Vec3
}

fn f32_min(a: f32, b: f32) -> f32 {
    return if a < b { a } else { b };
}

fn f32_max(a: f32, b: f32) -> f32 {
    return if a >= b { a } else { b };
}

impl AABB {
    pub fn extend(&mut self, m: &Mesh, f: &Face) {
        let coords = m.face_coords(f);

        self.min.x = f32::min(self.min.x, coords.iter().map(|c| c.x).reduce(f32_min).unwrap());
        self.min.y = f32::min(self.min.y, coords.iter().map(|c| c.y).reduce(f32_min).unwrap());
        self.min.z = f32::min(self.min.z, coords.iter().map(|c| c.z).reduce(f32_min).unwrap());

        self.max.x = f32::max(self.max.x, coords.iter().map(|c| c.x).reduce(f32_max).unwrap());
        self.max.y = f32::max(self.max.y, coords.iter().map(|c| c.y).reduce(f32_max).unwrap());
        self.max.z = f32::max(self.max.z, coords.iter().map(|c| c.z).reduce(f32_max).unwrap());
    }

    pub fn from(fr: &[FaceReference]) -> AABB {
        let mut aabb = AABB {
            min: Vec3::new(f32::INFINITY, f32::INFINITY, f32::INFINITY),
            max: Vec3::new(-f32::INFINITY, -f32::INFINITY, -f32::INFINITY),
        };

        for face_reference in fr {
            aabb.extend(face_reference.m, face_reference.f);
        }

        return aabb;
    }
    
    /**
     * Test for intersection of ABB with Ray r.
     * t is the maximum distance of the ray from its origin
     */
    pub fn intersect(&self, r: &Ray, t: f32) -> bool {
        // See [Tavian Barnes' site](https://tavianator.com/2011/ray_box.html) for details.
        let tx1 = (self.min.x - r.origin.x) * r.n_inv.x;
        let tx2 = (self.max.x - r.origin.x) * r.n_inv.x;

        let mut tmin = f32::min(tx1, tx2);
        let mut tmax = f32::max(tx1, tx2);

        let ty1 = (self.min.y - r.origin.y) * r.n_inv.y;
        let ty2 = (self.max.y - r.origin.y) * r.n_inv.y;

        tmin = f32::max(tmin, f32::min(ty1, ty2));
        tmax = f32::min(tmax, f32::max(ty1, ty2));

        let tz1 = (self.min.z - r.origin.z) * r.n_inv.z;
        let tz2 = (self.max.z - r.origin.z) * r.n_inv.z;

        tmin = f32::max(tmin, f32::min(tz1, tz2));
        tmax = f32::min(tmax, f32::max(tz1, tz2));

        return tmax >= f32::max(0.0, tmin) && tmin < t;
    }
}



/**
 * Bounding Volume Hierarchy
 */
pub enum BVH<'a> {
    Node(AABB, Box<BVH<'a>>, Box<BVH<'a>>),
    Leaf(AABB, Vec<FaceReference<'a>>)
}


#[derive(Copy, Clone)]
pub struct FaceReference<'a> {
    m: &'a Mesh,
    f: &'a Face
}

impl BVH<'_> {
    fn build_recursive<'a, 'b>(frs: &'b mut [FaceReference<'a>], recursion_limit: u16, split_axis: Axis) -> BVH<'a> {
        let aabb = AABB::from(frs);
        let frs_count = frs.len();

        if recursion_limit == 0 || frs_count <= 3 {
            // We have reached the maximum of recursions necessary.
            return BVH::Leaf(aabb, Vec::from(frs));
        }

        // Sort along the current axis
        frs.sort_by(|a,b| a.f.center.select(split_axis).partial_cmp(
                              &b.f.center.select(split_axis)).unwrap());

        // Select using surface area heuristic
        let total_surface: f32= frs.iter().map(|fr| fr.f.area).sum();

        let mut split_index : usize = 0;
        let mut surface_area = 0.0;
        while surface_area < total_surface / 2.0 {
            surface_area += frs[split_index].f.area;
            split_index += 1;
        }

        // Recurse
        let (frs1, frs2) = frs.split_at_mut(split_index);
        let left = Self::build_recursive(frs1, recursion_limit - 1, split_axis.next());
        let right = Self::build_recursive(frs2, recursion_limit - 1, split_axis.next());

        return BVH::Node(aabb, Box::new(left), Box::new(right));
    }

    pub fn from<'a>(meshes: &'a Vec<Mesh>, recursion_depth: u16) -> BVH {
        // Create references for each triangle
        let mut references = vec!();

        for m in meshes {
            for f in &m.faces {
                references.push(FaceReference {
                    m: &m,
                    f: &f
                });
            }
        }


        let bvh = Self::build_recursive(&mut references, recursion_depth, Axis::x);

        return bvh;
    }

    fn intersect_ray_triangle<'a>(ray: &Ray, fr: &'a FaceReference) -> Option<Intersection<'a>> {
        // See Akenine-Möller et al. Real-Time Rendering (2018), 22.8 Ray/Triangle Intersection
        let p = fr.m.face_coords(fr.f);
        let e1 = p[1] - p[0];
        let e2 = p[2] - p[0];
        let q = ray.direction.cross(e2);

        let a = e1 * q;
        if f32::abs(a) < 1e-9 {
            return None;
        }

        let f = 1.0 / a;
        let s = &ray.origin - p[0];

        let u = f * (s * q);
        if u < 0.0 {
            return None
        };

        let r = s.cross(e1);
        let v = f * (ray.direction * r);
        if v < 0.0 || u + v > 1.0 {
            return None;
        }

        let t = f * (e2 * r);

        return Some(Intersection {
            t: t,
            fr: *fr
        });

    }

    pub fn intersect<'a>(&'a self, r : &Ray, counter: usize) -> (usize, Option<Intersection<'a>>) {
        match self {
            BVH::Leaf(aabb, frs) => {
                // Early-out if the ray does not intersect the AABB
                if  !aabb.intersect(r, f32::INFINITY) {
                    return (counter+1, None);
                }

                // Make intersection tests with individual triangles
                let closest_intersection : Option<Intersection> = frs.iter()
                    .map(|fr| Self::intersect_ray_triangle(r, fr))
                    .filter(|i| i.is_some())
                    .map(|i| i.unwrap())
                    .filter(|i| i.t >= 0.0)
                    .reduce(|a, b| if a < b { a } else { b});

                return (counter+1, closest_intersection);
            },
            BVH::Node(aabb, left, right) => {
                // Early-out if the ray does not intersect the AABB
                if !aabb.intersect(r, f32::INFINITY) {
                    return (counter + 1, None);
                }

                // Recurse down the tree
                let (c1, left_intersection) = left.intersect(r, counter + 1);
                let (c2, right_intersection) = right.intersect(r, counter + 1);

                if left_intersection.is_some() && right_intersection.is_some() {
                    let (li, ri) = (left_intersection.unwrap(), right_intersection.unwrap());
                    return if li < ri { (c1 + c2, Some(li)) } else {(c1 + c2, Some(ri))};
                } else {
                    // At least one intersection is None.
                    return (c1+c2, left_intersection.or(right_intersection));
                }
            }
        }
    }

}

impl std::fmt::Display for AABB {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        return write!(f, "{} -> {}", self.min, self.max);
    }
}
impl std::fmt::Display for BVH<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            BVH::Leaf(aabb,_) => {
                return write!(f, "(Leaf {})", aabb);
            }
            BVH::Node(aabb, left, right) => {
                return write!(f, "(Node {}; {} {})", aabb, left, right);
            }
        }
    }
}