pbrt/shared/src/core/aggregates.rs

use crate::core::geometry::{Bounds3f, Point3f, Ray, Vector3f};
use crate::core::primitive::{Primitive, PrimitiveTrait};
use crate::core::shape::ShapeIntersection;
use crate::{gvec, Float, GVec, Ptr};

#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SplitMethod {
    SAH,
    Hlbvh,
    Middle,
    EqualCounts,
}

#[repr(C)]
#[derive(Default, Debug, Clone, Copy)]
pub struct LinearBVHNode {
    pub bounds: Bounds3f,
    pub primitives_offset: usize,
    pub n_primitives: u16,
    pub axis: u8,
    pub pad: u8,
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct BVHAggregate {
    pub node_count: u32,
    pub max_prims_in_node: u32,
    pub split_method: SplitMethod,
    pub primitives: GVec<Primitive>,
    pub nodes: GVec<LinearBVHNode>,
}

impl BVHAggregate {
    pub fn empty() -> Self {
        Self {
            node_count: 0,
            max_prims_in_node: 0,
            split_method: SplitMethod::SAH,
            primitives: gvec(),
            nodes: gvec(),
        }
    }

    #[inline(always)]
    fn node(&self, i: usize) -> &LinearBVHNode {
        unsafe { self.nodes.get_unchecked(i) }
    }

    #[inline(always)]
    fn primitive(&self, i: usize) -> &Primitive {
        unsafe { self.primitives.get_unchecked(i) }
    }
}

impl PrimitiveTrait for BVHAggregate {
    fn bounds(&self) -> Bounds3f {
        if self.nodes.is_empty() || self.node_count == 0 {
            Bounds3f::default()
        } else {
            self.node(0).bounds
        }
    }

    fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        if self.nodes.is_empty() {
            return None;
        }

        let mut best_si: Option<ShapeIntersection> = None;

        let mut hit_t = t_max.unwrap_or(Float::INFINITY);

        let inv_dir = Vector3f::new(1.0 / r.d.x(), 1.0 / r.d.y(), 1.0 / r.d.z());
        let dir_is_neg = [
            if inv_dir.x() < 0.0 { 1 } else { 0 },
            if inv_dir.y() < 0.0 { 1 } else { 0 },
            if inv_dir.z() < 0.0 { 1 } else { 0 },
        ];

        let mut to_visit_offset = 0;
        let mut current_node_index = 0;
        let mut nodes_to_visit = [0usize; 64];

        loop {
            let node = &self.nodes[current_node_index];

            // Check ray against BVH node bounds using the current closest hit_t
            if node
                .bounds
                .intersect_p(r.o, hit_t, inv_dir, &dir_is_neg)
                .is_some()
            {
                if node.n_primitives > 0 {
                    // Intersect ray with all primitives in this leaf
                    for i in 0..node.n_primitives {
                        let prim_idx = node.primitives_offset + i as usize;
                        let prim = &self.primitives[prim_idx];

                        if let Some(si) = prim.intersect(r, Some(hit_t)) {
                            hit_t = si.t_hit();
                            best_si = Some(si);
                        }
                    }

                    if to_visit_offset == 0 {
                        break;
                    }
                    to_visit_offset -= 1;
                    current_node_index = nodes_to_visit[to_visit_offset];
                } else {
                    // Check the sign of the ray direction against the split axis
                    if dir_is_neg[node.axis as usize] == 1 {
                        // Ray is negative (Right -> Left).
                        // Near child is Second Child (stored in primitives_offset).
                        // Far child is First Child (current + 1).

                        // Push Far
                        nodes_to_visit[to_visit_offset] = current_node_index + 1;
                        to_visit_offset += 1;

                        // Visit Near immediately
                        current_node_index = node.primitives_offset;
                    } else {
                        // Ray is positive (Left -> Right).
                        // Push Far
                        nodes_to_visit[to_visit_offset] = node.primitives_offset;
                        to_visit_offset += 1;

                        current_node_index += 1;
                    }
                }
            } else {
                // The ray missed the AABB of this node. Pop stack to try the next node.
                if to_visit_offset == 0 {
                    break;
                }
                to_visit_offset -= 1;
                current_node_index = nodes_to_visit[to_visit_offset];
            }
        }

        best_si
    }

    fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
        if self.nodes.is_empty() {
            return false;
        }

        let t_max = t_max.unwrap_or(Float::INFINITY);

        let inv_dir = Vector3f::new(1.0 / r.d.x(), 1.0 / r.d.y(), 1.0 / r.d.z());
        let dir_is_neg = [
            if inv_dir.x() < 0.0 { 1 } else { 0 },
            if inv_dir.y() < 0.0 { 1 } else { 0 },
            if inv_dir.z() < 0.0 { 1 } else { 0 },
        ];

        let mut to_visit_offset = 0;
        let mut current_node_index = 0;
        let mut nodes_to_visit = [0usize; 64];

        loop {
            let node = &self.nodes[current_node_index];

            // Check AABB
            if node
                .bounds
                .intersect_p(r.o, t_max, inv_dir, &dir_is_neg)
                .is_some()
            {
                if node.n_primitives > 0 {
                    for i in 0..node.n_primitives {
                        let prim_idx = node.primitives_offset + i as usize;
                        let prim = &self.primitives[prim_idx];

                        if prim.intersect_p(r, Some(t_max)) {
                            return true;
                        }
                    }

                    // No intersection in this leaf, try next node in stack
                    if to_visit_offset == 0 {
                        break;
                    }
                    to_visit_offset -= 1;
                    current_node_index = nodes_to_visit[to_visit_offset];
                } else {
                    // Standard front-to-back traversal order helps find an occlusion
                    // closer to the origin faster, potentially saving work.

                    if dir_is_neg[node.axis as usize] == 1 {
                        nodes_to_visit[to_visit_offset] = current_node_index + 1;
                        to_visit_offset += 1;
                        current_node_index = node.primitives_offset;
                    } else {
                        nodes_to_visit[to_visit_offset] = node.primitives_offset;
                        to_visit_offset += 1;
                        current_node_index += 1;
                    }
                }
            } else {
                if to_visit_offset == 0 {
                    break;
                }
                to_visit_offset -= 1;
                current_node_index = nodes_to_visit[to_visit_offset];
            }
        }
        false
    }
}