pbrt/shared/src/core/texture.rs

use crate::core::color::ColorEncoding;
use crate::core::geometry::{
    Normal3f, Point2f, Point3f, Vector2f, Vector3f, VectorLike, spherical_phi, spherical_theta,
};
use crate::core::image::WrapMode;
use crate::core::interaction::{Interaction, InteractionTrait, SurfaceInteraction};
use crate::spectra::{
    RGBAlbedoSpectrum, RGBIlluminantSpectrum, RGBUnboundedSpectrum, SampledSpectrum,
    SampledWavelengths,
};
use crate::textures::*;
use crate::utils::Ptr;
use crate::utils::Transform;
use crate::utils::math::square;
use crate::{Float, INV_2_PI, INV_PI, PI};
use enum_dispatch::enum_dispatch;

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct TexCoord2D {
    pub st: Point2f,
    pub dsdx: Float,
    pub dsdy: Float,
    pub dtdx: Float,
    pub dtdy: Float,
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub enum TextureMapping2D {
    UV(UVMapping),
    Spherical(SphericalMapping),
    Cylindrical(CylindricalMapping),
    Planar(PlanarMapping),
}

impl TextureMapping2D {
    pub fn map(&self, ctx: &TextureEvalContext) -> TexCoord2D {
        match self {
            TextureMapping2D::UV(t) => t.map(ctx),
            TextureMapping2D::Spherical(t) => t.map(ctx),
            TextureMapping2D::Cylindrical(t) => t.map(ctx),
            TextureMapping2D::Planar(t) => t.map(ctx),
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct UVMapping {
    pub su: Float,
    pub sv: Float,
    pub du: Float,
    pub dv: Float,
}

impl Default for UVMapping {
    fn default() -> Self {
        Self {
            su: 1.0,
            sv: 1.0,
            du: 0.0,
            dv: 0.0,
        }
    }
}

impl UVMapping {
    pub fn new(su: Float, sv: Float, du: Float, dv: Float) -> Self {
        Self { su, sv, du, dv }
    }

    pub fn map(&self, ctx: &TextureEvalContext) -> TexCoord2D {
        let dsdx = self.su * ctx.dudx;
        let dsdy = self.su * ctx.dudy;
        let dtdx = self.sv * ctx.dvdx;
        let dtdy = self.sv * ctx.dvdy;
        let st = Point2f::new(self.su * ctx.uv[0] + self.du, self.sv * ctx.uv[1] * self.dv);
        TexCoord2D {
            st,
            dsdx,
            dsdy,
            dtdx,
            dtdy,
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct SphericalMapping {
    texture_from_render: Transform,
}

impl SphericalMapping {
    pub fn new(texture_from_render: &Transform) -> Self {
        Self {
            texture_from_render: *texture_from_render,
        }
    }

    pub fn map(&self, ctx: &TextureEvalContext) -> TexCoord2D {
        let pt = self.texture_from_render.apply_to_point(ctx.p);
        let x2y2 = square(pt.x()) + square(pt.y());
        let sqrtx2y2 = x2y2.sqrt();
        let dsdp = Vector3f::new(-pt.y(), pt.x(), 0.) / (2. * PI * x2y2);
        let dtdp = 1. / (PI * (x2y2 * square(pt.z())))
            * Vector3f::new(
                pt.x() * pt.z() / sqrtx2y2,
                pt.y() * pt.z() / sqrtx2y2,
                -sqrtx2y2,
            );
        let dpdx = self.texture_from_render.apply_to_vector(ctx.dpdx);
        let dpdy = self.texture_from_render.apply_to_vector(ctx.dpdy);
        let dsdx = dsdp.dot(dpdx);
        let dsdy = dsdp.dot(dpdy);
        let dtdx = dtdp.dot(dpdx);
        let dtdy = dtdp.dot(dpdy);
        let vec = (pt - Point3f::default()).normalize();
        let st = Point2f::new(spherical_theta(vec) * INV_PI, spherical_phi(vec) * INV_2_PI);
        TexCoord2D {
            st,
            dsdx,
            dsdy,
            dtdx,
            dtdy,
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct CylindricalMapping {
    texture_from_render: Transform,
}

impl CylindricalMapping {
    pub fn new(texture_from_render: &Transform) -> Self {
        Self {
            texture_from_render: *texture_from_render,
        }
    }

    pub fn map(&self, ctx: &TextureEvalContext) -> TexCoord2D {
        let pt = self.texture_from_render.apply_to_point(ctx.p);
        let x2y2 = square(pt.x()) + square(pt.y());
        let dsdp = Vector3f::new(-pt.y(), pt.x(), 0.) / (2. * PI * x2y2);
        let dtdp = Vector3f::new(1., 0., 0.);
        let dpdx = self.texture_from_render.apply_to_vector(ctx.dpdx);
        let dpdy = self.texture_from_render.apply_to_vector(ctx.dpdy);
        let dsdx = dsdp.dot(dpdx);
        let dsdy = dsdp.dot(dpdy);
        let dtdx = dtdp.dot(dpdx);
        let dtdy = dtdp.dot(dpdy);
        let st = Point2f::new((PI * pt.y().atan2(pt.x())) * INV_2_PI, pt.z());
        TexCoord2D {
            st,
            dsdx,
            dsdy,
            dtdx,
            dtdy,
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct PlanarMapping {
    texture_from_render: Transform,
    vs: Vector3f,
    vt: Vector3f,
    ds: Float,
    dt: Float,
}

impl PlanarMapping {
    pub fn new(
        texture_from_render: &Transform,
        vs: Vector3f,
        vt: Vector3f,
        ds: Float,
        dt: Float,
    ) -> Self {
        Self {
            texture_from_render: *texture_from_render,
            vs,
            vt,
            ds,
            dt,
        }
    }

    pub fn map(&self, ctx: &TextureEvalContext) -> TexCoord2D {
        let vec: Vector3f = self.texture_from_render.apply_to_point(ctx.p).into();
        let dpdx = self.texture_from_render.apply_to_vector(ctx.dpdx);
        let dpdy = self.texture_from_render.apply_to_vector(ctx.dpdy);
        let dsdx = self.vs.dot(dpdx);
        let dsdy = self.vs.dot(dpdy);
        let dtdx = self.vt.dot(dpdx);
        let dtdy = self.vt.dot(dpdy);
        let st = Point2f::new(self.ds + vec.dot(self.vs), self.dt + vec.dot(self.vt));
        TexCoord2D {
            st,
            dsdx,
            dsdy,
            dtdx,
            dtdy,
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct TexCoord3D {
    pub p: Point3f,
    pub dpdx: Vector3f,
    pub dpdy: Vector3f,
}

pub trait TextureMapping3DTrait {
    fn map(&self, ctx: &TextureEvalContext) -> TexCoord3D;
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub enum TextureMapping3D {
    PointTransform(PointTransformMapping),
}

impl TextureMapping3D {
    pub fn map(&self, ctx: &TextureEvalContext) -> TexCoord3D {
        match self {
            TextureMapping3D::PointTransform(t) => t.map(ctx),
        }
    }
}

#[repr(C)]
#[derive(Clone, Debug, Copy)]
pub struct PointTransformMapping {
    pub texture_from_render: Transform,
}

impl PointTransformMapping {
    #[cfg(not(target_os = "cuda"))]
    pub fn new(texture_from_render: Transform) -> Self {
        Self {
            texture_from_render,
        }
    }

    pub fn map(&self, ctx: &TextureEvalContext) -> TexCoord3D {
        TexCoord3D {
            p: self.texture_from_render.apply_to_point(ctx.p),
            dpdx: self.texture_from_render.apply_to_vector(ctx.dpdx),
            dpdy: self.texture_from_render.apply_to_vector(ctx.dpdy),
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Default, Debug)]
pub struct TextureEvalContext {
    pub p: Point3f,
    pub dpdx: Vector3f,
    pub dpdy: Vector3f,
    pub n: Normal3f,
    pub uv: Point2f,
    pub dudx: Float,
    pub dudy: Float,
    pub dvdx: Float,
    pub dvdy: Float,
    pub face_index: u32,
}

impl TextureEvalContext {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        p: Point3f,
        dpdx: Vector3f,
        dpdy: Vector3f,
        n: Normal3f,
        uv: Point2f,
        dudx: Float,
        dudy: Float,
        dvdx: Float,
        dvdy: Float,
        face_index: u32,
    ) -> Self {
        Self {
            p,
            dpdx,
            dpdy,
            n,
            uv,
            dudx,
            dudy,
            dvdx,
            dvdy,
            face_index,
        }
    }
}

impl From<&SurfaceInteraction> for TextureEvalContext {
    fn from(si: &SurfaceInteraction) -> Self {
        Self {
            p: si.p(),
            dpdx: si.dpdx,
            dpdy: si.dpdy,
            n: si.common.n,
            uv: si.common.uv,
            dudx: si.dudx,
            dudy: si.dudy,
            dvdx: si.dvdx,
            dvdy: si.dvdy,
            face_index: si.face_index,
        }
    }
}

impl From<&Interaction> for TextureEvalContext {
    fn from(intr: &Interaction) -> Self {
        match intr {
            Interaction::Surface(si) => TextureEvalContext::from(si),
            Interaction::Medium(mi) => TextureEvalContext {
                p: mi.p(),
                ..Default::default()
            },

            Interaction::Simple(si) => TextureEvalContext {
                p: si.p(),
                ..Default::default()
            },
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub enum GPUFloatTexture {
    Constant(FloatConstantTexture),
    DirectionMix(GPUFloatDirectionMixTexture),
    Scaled(GPUFloatScaledTexture),
    Bilerp(FloatBilerpTexture),
    Checkerboard(FloatCheckerboardTexture),
    Dots(FloatDotsTexture),
    FBm(FBmTexture),
    Windy(WindyTexture),
    Wrinkled(WrinkledTexture),
    Ptex(GPUFloatPtexTexture),
    Image(GPUFloatImageTexture),
    Mix(GPUFloatMixTexture),
}

impl GPUFloatTexture {
    pub fn evaluate(&self, ctx: &TextureEvalContext) -> Float {
        match self {
            GPUFloatTexture::Constant(t) => t.evaluate(ctx),
            GPUFloatTexture::DirectionMix(t) => t.evaluate(ctx),
            GPUFloatTexture::Scaled(t) => t.evaluate(ctx),
            GPUFloatTexture::Bilerp(t) => t.evaluate(ctx),
            GPUFloatTexture::Checkerboard(t) => t.evaluate(ctx),
            GPUFloatTexture::Dots(t) => t.evaluate(ctx),
            GPUFloatTexture::FBm(t) => t.evaluate(ctx),
            GPUFloatTexture::Windy(t) => t.evaluate(ctx),
            GPUFloatTexture::Wrinkled(t) => t.evaluate(ctx),
            GPUFloatTexture::Ptex(t) => t.evaluate(ctx),
            GPUFloatTexture::Image(t) => t.evaluate(ctx),
            GPUFloatTexture::Mix(t) => t.evaluate(ctx),
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub enum SpectrumType {
    Illuminant,
    Albedo,
    Unbounded,
}

#[repr(C)]
#[enum_dispatch]
#[derive(Clone, Copy, Debug)]
pub enum GPUSpectrumTexture {
    Constant(SpectrumConstantTexture),
    Bilerp(SpectrumBilerpTexture),
    Checkerboard(SpectrumCheckerboardTexture),
    Marble(MarbleTexture),
    DirectionMix(GPUSpectrumDirectionMixTexture),
    Dots(SpectrumDotsTexture),
    Scaled(GPUSpectrumScaledTexture),
    Image(GPUSpectrumImageTexture),
    Ptex(GPUSpectrumPtexTexture),
    Mix(GPUSpectrumMixTexture),
}

impl GPUSpectrumTexture {
    pub fn evaluate(
        &self,
        ctx: &TextureEvalContext,
        lambda: &SampledWavelengths,
    ) -> SampledSpectrum {
        match self {
            GPUSpectrumTexture::Constant(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::Bilerp(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::Checkerboard(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::Marble(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::DirectionMix(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::Dots(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::Scaled(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::Ptex(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::Image(t) => t.evaluate(ctx, lambda),
            GPUSpectrumTexture::Mix(t) => t.evaluate(ctx, lambda),
        }
    }
}

pub trait TextureEvaluator: Send + Sync {
    fn evaluate_float(&self, tex: &GPUFloatTexture, ctx: &TextureEvalContext) -> Float;
    fn evaluate_spectrum(
        &self,
        tex: &GPUSpectrumTexture,
        ctx: &TextureEvalContext,
        lambda: &SampledWavelengths,
    ) -> SampledSpectrum;

    fn can_evaluate(
        &self,
        _ftex: &[Ptr<GPUFloatTexture>],
        _stex: &[Ptr<GPUSpectrumTexture>],
    ) -> bool;
}

#[repr(C)]
#[derive(Copy, Clone, Default)]
pub struct UniversalTextureEvaluator;

impl TextureEvaluator for UniversalTextureEvaluator {
    fn evaluate_float(&self, tex: &GPUFloatTexture, ctx: &TextureEvalContext) -> Float {
        tex.evaluate(ctx)
    }

    fn evaluate_spectrum(
        &self,
        tex: &GPUSpectrumTexture,
        ctx: &TextureEvalContext,
        lambda: &SampledWavelengths,
    ) -> SampledSpectrum {
        tex.evaluate(ctx, lambda)
    }

    fn can_evaluate(
        &self,
        _float_textures: &[Ptr<GPUFloatTexture>],
        _spectrum_textures: &[Ptr<GPUSpectrumTexture>],
    ) -> bool {
        true
    }
}