pbrt/src/wavefront/aggregate.rs

use crate::globals::get_options;
use rayon::prelude::*;
use shared::core::geometry::{Bounds3f, Ray, VectorLike};
use shared::core::interaction::{InteractionTrait, SurfaceInteraction};
use shared::core::material::MaterialTrait;
use shared::core::primitive::{Primitive, PrimitiveTrait};
use shared::core::texture::BasicTextureEvaluator;
use shared::core::texture::TextureEvaluator;
use shared::wavefront::workitems::*;
use shared::wavefront::WavefrontAggregate;
use shared::{Float, Ptr};

pub struct CpuAggregate {
    pub aggregate: Primitive,
}

impl CpuAggregate {
    pub fn new(aggregate: Primitive) -> Self {
        Self { aggregate }
    }
}

// fn enqueue_after_intersection(
//     r: RayWorkItem, ray_medium: Medium, t_max: Float, intr: SurfaceInteraction,
//     mut medium_sample_queue: &MediumQueue, mut next_ray_queue: &RayQueue,
//     mut hit_area_light_queue: &HitAreaLightQueue, mut basic_eval_mtl_q: &MaterialEvalQueue,
//     mut universal_eval_mlt_q: MaterialEvalQueue) {
// }

impl WavefrontAggregate for CpuAggregate {
    fn bounds(&self) -> Bounds3f {
        self.aggregate.bounds()
    }

    fn intersect_closest(
        &self,
        max_rays: usize,
        ray_q: &RayQueue,
        escaped_ray_q: &EscapedRayQueue,
        hit_area_light_q: &HitAreaLightQueue,
        basic_eval_mtl_q: &MaterialEvalQueue,
        universal_eval_mtl_q: &MaterialEvalQueue,
        next_ray_q: &RayQueue,
        _pixel_sample_state: &PixelSampleState,
    ) {
        let n_rays = ray_q.size().min(max_rays as u32);

        // Intersect ray with the scene and enqueue resulting work
        (0..n_rays as usize).into_par_iter().for_each(|i| {
            let r = unsafe { ray_q.get(i) };

            let Some(si) = self.aggregate.intersect(&r.ray, None) else {
                // EnqueueMiss
                escaped_ray_q.push(EscapedRayWorkItem {
                    ray_o: r.ray.o,
                    ray_d: r.ray.d,
                    lambda: r.lambda,
                    pixel_index: r.pixel_index,
                    beta: r.beta,
                    r_u: r.r_u,
                    r_l: r.r_l,
                    depth: r.depth,
                    specular_bounce: r.specular_bounce != 0,
                    prev_intr_ctx: r.prev_intr_ctx,
                });
                return;
            };

            let intr = &si.intr;

            // Medium transition
            if intr.material.is_null() {
                let mut next = r;
                next.ray = intr.spawn_ray(r.ray.d);
                next_ray_q.push(next);
                return;
            }

            // Area light hit
            if !intr.area_light.is_null() {
                hit_area_light_q.push(HitAreaLightWorkItem {
                    area_light: intr.area_light,
                    p: intr.p(),
                    n: intr.n(),
                    uv: intr.common.uv,
                    wo: -r.ray.d,
                    lambda: r.lambda,
                    pixel_index: r.pixel_index,
                    beta: r.beta,
                    r_u: r.r_u,
                    r_l: r.r_l,
                    depth: r.depth,
                    specular_bounce: r.specular_bounce != 0,
                    prev_intr_ctx: r.prev_intr_ctx,
                });
            }

            // Material eval queue dispatch
            let material = *intr.material.get().unwrap();
            let eval_q = if material.can_evaluate_textures(&BasicTextureEvaluator) {
                basic_eval_mtl_q
            } else {
                universal_eval_mtl_q
            };

            eval_q.push(MaterialEvalWorkItem {
                p: intr.pi(),
                n: intr.n(),
                ns: intr.shading.n,
                dpdu: intr.shading.dpdu,
                dpdv: intr.shading.dpdv,
                uv: intr.common.uv,
                wo: intr.wo(),
                time: r.ray.time,
                face_index: intr.face_index,
                material: intr.material,
                area_light: intr.area_light,
                medium_interface: intr.common.medium_interface,
                pixel_index: r.pixel_index,
                lambda: r.lambda,
                beta: r.beta,
                r_u: r.r_u,
                any_non_specular_bounces: r.any_non_specular_bounces != 0,
                depth: r.depth,
                eta_scale: r.eta_scale,
                dpdus: intr.shading.dpdu,
                dpdvs: intr.shading.dpdv,
                dndus: intr.shading.dndu,
                dndvs: intr.shading.dndv,
            });
        });
    }

    fn intersect_shadow(
        &self,
        max_rays: usize,
        shadow_ray_q: &ShadowRayQueue,
        pixel_sample_state: &PixelSampleState,
    ) {
        let n_rays = shadow_ray_q.size().min(max_rays as u32);

        (0..n_rays as usize).into_par_iter().for_each(|i| {
            let work = unsafe { shadow_ray_q.get(i) };

            let ray = Ray::new(work.ray.o, work.ray.d, Some(work.ray.time), Ptr::null());
            if !self.aggregate.intersect_p(&ray, Some(work.t_max)) {
                let pi = work.pixel_index as usize;
                let ld = work.l_d / (work.r_u + work.r_l).average();
                let mut l = pixel_sample_state.l.get(pi);
                l += ld;
                pixel_sample_state.l.set(pi, l);
            }
        });
    }

    fn intersect_shadow_tr(
        &self,
        max_rays: usize,
        shadow_ray_q: &ShadowRayQueue,
        pixel_sample_state: &PixelSampleState,
    ) {
        self.intersect_shadow(max_rays, shadow_ray_q, pixel_sample_state);
    }
}