pbrt/shared/src/core/primitive.rs

use crate::core::geometry::{Bounds3f, Ray};
use crate::core::aggregates::DeviceBVHAggregate;
use crate::core::interaction::{Interaction, InteractionTrait, SurfaceInteraction};
use crate::core::light::Light;
use crate::core::material::Material;
use crate::core::medium::{Medium, MediumInterface};
use crate::core::pbrt::Float;
use crate::core::shape::{Shape, ShapeIntersection, ShapeTrait};
use crate::core::texture::{GPUFloatTexture, TextureEvalContext};
use crate::utils::hash::hash_float;
use crate::utils::transform::{AnimatedTransform, Transform};
use crate::utils::Ptr;
use alloc::boxed::Box;
use alloc::sync::Arc;

use enum_dispatch::enum_dispatch;

#[enum_dispatch]
pub trait PrimitiveTrait: Send + Sync {
    fn bounds(&self) -> Bounds3f;
    fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection>;
    fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool;
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct GeometricPrimitive {
    pub shape: Ptr<Shape>,
    pub material: Ptr<Material>,
    pub area_light: Ptr<Light>,
    pub medium_interface: MediumInterface,
    pub alpha: Ptr<GPUFloatTexture>,
}

unsafe impl Send for GeometricPrimitive {}
unsafe impl Sync for GeometricPrimitive {}

impl PrimitiveTrait for GeometricPrimitive {
    fn bounds(&self) -> Bounds3f {
        self.shape.bounds()
    }

    fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        let mut si = self.shape.intersect(r, t_max)?;
        if !self.alpha.is_null() {
            let alpha = unsafe { &self.alpha.as_ref() };
            let ctx = TextureEvalContext::from(&si.intr);
            let a = alpha.evaluate(&ctx);
            if a < 1.0 {
                let u = if a <= 0.0 {
                    1.0
                } else {
                    hash_float(&(r.o, r.d))
                };

                if u > a {
                    let r_next = si.intr.spawn_ray(r.d);

                    let new_t_max = t_max.map(|t| t - si.t_hit());

                    if let Some(mut si_next) = self.intersect(&r_next, new_t_max) {
                        si_next.set_t_hit(si_next.t_hit() + si.t_hit());
                        return Some(si_next);
                    } else {
                        return None;
                    }
                }
            }
        }

        if r.medium.is_null() {
            return None;
        }
        si.set_intersection_properties(
            self.material,
            self.area_light,
            self.medium_interface,
            r.medium,
        );

        Some(si)
    }

    fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
        if !self.alpha.is_null() {
            self.intersect(r, t_max).is_some()
        } else {
            self.shape.intersect_p(r, t_max)
        }
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct SimplePrimitive {
    pub shape: Ptr<Shape>,
    pub material: Ptr<Material>,
}

impl PrimitiveTrait for SimplePrimitive {
    fn bounds(&self) -> Bounds3f {
        self.shape.bounds()
    }

    fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        let mut si = self.shape.intersect(r, t_max)?;
        si.set_intersection_properties(self.material, Ptr::null(), MediumInterface::default(), r.medium);
        Some(si)
    }

    fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
        self.shape.intersect_p(r, t_max)
    }
}

#[derive(Debug, Clone, Copy)]
pub struct TransformedPrimitive {
    pub primitive: Ptr<Primitive>,
    pub render_from_primitive: Ptr<Transform>,
}

impl PrimitiveTrait for TransformedPrimitive {
    fn bounds(&self) -> Bounds3f {
        self.render_from_primitive
            .apply_to_bounds(self.primitive.bounds())
    }

    fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        let (ray, t_max) = self.render_from_primitive.apply_inverse_ray(r, t_max);
        let mut si = self.primitive.intersect(&ray, Some(t_max))?;
        let intr_wrapper = Interaction::Surface(si.intr);
        let new_render_enum = self
            .render_from_primitive
            .apply_to_interaction(&intr_wrapper);

        if let Interaction::Surface(new_surf) = new_render_enum {
            si.intr = new_surf;
        } else {
            panic!("TransformedPrimitive: Transform changed interaction type (impossible)");
        }

        Some(si)
    }

    fn intersect_p(&self, r: &Ray, t_max: Option<Float>) -> bool {
        let (ray, t_max) = self.render_from_primitive.apply_inverse_ray(r, t_max);
        self.primitive.intersect_p(&ray, Some(t_max))
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct AnimatedPrimitive {
    pub primitive: Ptr<Primitive>,
    pub render_from_primitive: Ptr<AnimatedTransform>,
}

impl PrimitiveTrait for AnimatedPrimitive {
    fn bounds(&self) -> Bounds3f {
        self.render_from_primitive
            .motion_bounds(&self.primitive.bounds())
    }

    fn intersect(&self, r: &Ray, t_max: Option<Float>) -> Option<ShapeIntersection> {
        let interp_render_from_primitive = self.render_from_primitive.interpolate(r.time);
        let (ray, t_max) = interp_render_from_primitive.apply_inverse_ray(r, t_max);
        let mut si = self.primitive.intersect(&ray, Some(t_max))?;
        let wrapper = si.intr.into();
        let new_wrapper = interp_render_from_primitive.apply_to_interaction(&wrapper);

        if let Interaction::Surface(new_surf) = new_wrapper {
            si.intr = new_surf;
        } else {
            panic!("TransformedPrimitive: Interaction type changed unexpectedly!");
        }

        Some(si)
    }

    fn intersect_p(&self, _r: &Ray, _t_max: Option<Float>) -> bool {
        todo!()
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct LinearBVHNode {
    bounds: Bounds3f,
}

#[derive(Debug, Clone, Copy)]
pub struct KdTreeAggregate;

impl PrimitiveTrait for KdTreeAggregate {
    fn bounds(&self) -> Bounds3f {
        todo!()
    }

    fn intersect(&self, _r: &Ray, _t_max: Option<Float>) -> Option<ShapeIntersection> {
        todo!()
    }

    fn intersect_p(&self, _r: &Ray, _t_max: Option<Float>) -> bool {
        todo!()
    }
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
#[enum_dispatch(PrimitiveTrait)]
pub enum Primitive {
    Simple(SimplePrimitive),
    Geometric(GeometricPrimitive),
    Transformed(TransformedPrimitive),
    Animated(AnimatedPrimitive),
    BVH(DeviceBVHAggregate),
    KdTree(KdTreeAggregate),
}