pbrt/shared/src/core/bxdf.rs

use crate::bxdfs::*;
use crate::core::bsdf::BSDFSample;
use crate::core::geometry::{Point2f, Vector3f, abs_cos_theta};
use crate::spectra::SampledSpectrum;
use crate::utils::sampling::{sample_uniform_hemisphere, uniform_hemisphere_pdf};
use crate::{Float, PI};
use bitflags::bitflags;
use core::any::Any;
use core::ops::Not;
use enum_dispatch::enum_dispatch;

bitflags! {
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub struct BxDFReflTransFlags: u8 {
        const REFLECTION   = 1 << 0;
        const TRANSMISSION = 1 << 1;
        const ALL = Self::REFLECTION.bits() | Self::TRANSMISSION.bits();
    }
}

bitflags! {
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub struct BxDFFlags: u8 {
        const UNSET = 0;
        const REFLECTION   = 1 << 0;
        const TRANSMISSION = 1 << 1;
        const DIFFUSE      = 1 << 2;
        const GLOSSY       = 1 << 3;
        const SPECULAR     = 1 << 4;

        const DIFFUSE_REFLECTION     = Self::DIFFUSE.bits() | Self::REFLECTION.bits();
        const DIFFUSE_TRANSMISSION   = Self::DIFFUSE.bits() | Self::TRANSMISSION.bits();
        const GLOSSY_REFLECTION      = Self::GLOSSY.bits() | Self::REFLECTION.bits();
        const GLOSSY_TRANSMISSION    = Self::GLOSSY.bits() | Self::TRANSMISSION.bits();
        const SPECULAR_REFLECTION    = Self::SPECULAR.bits() | Self::REFLECTION.bits();
        const SPECULAR_TRANSMISSION  = Self::SPECULAR.bits() | Self::TRANSMISSION.bits();

        const SCATTERING = Self::DIFFUSE.bits() | Self::GLOSSY.bits() | Self::SPECULAR.bits();
        const ALL = Self::REFLECTION.bits() | Self::TRANSMISSION.bits() | Self::SCATTERING.bits();
    }
}

impl BxDFFlags {
    #[inline]
    pub fn is_reflective(&self) -> bool {
        self.contains(Self::REFLECTION)
    }
    #[inline]
    pub fn is_transmissive(&self) -> bool {
        self.contains(Self::TRANSMISSION)
    }
    #[inline]
    pub fn is_diffuse(&self) -> bool {
        self.contains(Self::DIFFUSE)
    }
    #[inline]
    pub fn is_glossy(&self) -> bool {
        self.contains(Self::GLOSSY)
    }
    #[inline]
    pub fn is_specular(&self) -> bool {
        self.contains(Self::SPECULAR)
    }
    #[inline]
    pub fn is_non_specular(&self) -> bool {
        self.intersects(Self::DIFFUSE | Self::GLOSSY)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransportMode {
    Radiance,
    Importance,
}

impl Not for TransportMode {
    type Output = Self;
    fn not(self) -> Self::Output {
        match self {
            TransportMode::Radiance => TransportMode::Importance,
            TransportMode::Importance => TransportMode::Radiance,
        }
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct FArgs {
    pub mode: TransportMode,
    pub sample_flags: BxDFReflTransFlags,
}

impl Default for FArgs {
    fn default() -> Self {
        Self {
            mode: TransportMode::Radiance,
            sample_flags: BxDFReflTransFlags::ALL,
        }
    }
}

#[enum_dispatch]
pub trait BxDFTrait: Any {
    fn flags(&self) -> BxDFFlags;

    fn f(&self, wo: Vector3f, wi: Vector3f, mode: TransportMode) -> SampledSpectrum;

    fn sample_f(&self, wo: Vector3f, uc: Float, u: Point2f, f_args: FArgs) -> Option<BSDFSample>;

    fn pdf(&self, wo: Vector3f, wi: Vector3f, f_args: FArgs) -> Float;

    fn rho_wo(&self, wo: Vector3f, uc: &[Float], u2: &[Point2f]) -> SampledSpectrum {
        let mut r = SampledSpectrum::new(0.);
        for i in 0..uc.len() {
            if let Some(bs) = self.sample_f(wo, uc[i], u2[i], FArgs::default()) {
                r += bs.f * abs_cos_theta(bs.wi) / bs.pdf;
            }
        }
        r / uc.len() as Float
    }

    fn rho_u(&self, u1: &[Point2f], uc: &[Float], u2: &[Point2f]) -> SampledSpectrum {
        let mut r = SampledSpectrum::new(0.);
        for i in 0..uc.len() {
            let wo = sample_uniform_hemisphere(u1[i]);
            if wo.z() == 0. {
                continue;
            }
            let pdfo = uniform_hemisphere_pdf();
            if let Some(bs) = self.sample_f(wo, uc[i], u2[i], FArgs::default()) {
                r += bs.f * abs_cos_theta(bs.wi) * abs_cos_theta(wo) / (pdfo * bs.pdf);
            }
        }
        r / (PI * uc.len() as Float)
    }

    fn regularize(&mut self) {
        todo!();
    }

    fn as_any(&self) -> &dyn Any;
}

#[repr(C)]
#[enum_dispatch(BxDFTrait)]
#[derive(Debug, Clone, Copy)]
pub enum BxDF {
    Diffuse(DiffuseBxDF),
    Dielectric(DielectricBxDF),
    ThinDielectric(ThinDielectricBxDF),
    Conductor(ConductorBxDF),
    Measured(MeasuredBxDF),
    Hair(HairBxDF),
    CoatedDiffuse(CoatedDiffuseBxDF),
    CoatedConductor(CoatedConductorBxDF),
    NormalizedFresnel(NormalizedFresnelBxDF),
}