pbrt/src/utils/upload.rs

use crate::core::image::Image;
use crate::spectra::DenselySampledSpectrumBuffer;
use crate::Arena;
use shared::core::color::RGBToSpectrumTable;
use shared::core::image::DeviceImage;
use shared::core::light::Light;
use shared::core::material::Material;
use shared::core::shape::Shape;
use shared::core::spectrum::Spectrum;
use shared::spectra::{DenselySampledSpectrum, DeviceStandardColorSpaces, RGBColorSpace};
use shared::Ptr;
use std::slice::from_raw_parts;

pub trait DeviceRepr {
    /// The `#[repr(C)] Copy` device-side struct.
    type Target: Copy;

    /// Upload into the arena and return the device struct by value.
    /// Use this when embedding the result inline in another device struct.
    fn upload_value(&self, arena: &Arena) -> Self::Target;

    /// Upload into the arena and return a Ptr to the device struct.
    /// This is the common entry point — allocates the Target in the arena.
    fn upload(&self, arena: &Arena) -> Ptr<Self::Target> {
        let value = self.upload_value(arena);
        arena.alloc(value)
    }
}

impl<T: DeviceRepr> DeviceRepr for Option<T> {
    type Target = T::Target;

    fn upload_value(&self, arena: &Arena) -> Self::Target {
        match self {
            Some(val) => val.upload_value(arena),
            None => panic!("Cannot upload_value on None — use upload() which returns Ptr::null()"),
        }
    }

    fn upload(&self, arena: &Arena) -> Ptr<Self::Target> {
        match self {
            Some(val) => val.upload(arena),
            None => Ptr::null(),
        }
    }
}

impl<T: DeviceRepr> DeviceRepr for std::sync::Arc<T> {
    type Target = T::Target;

    fn upload_value(&self, arena: &Arena) -> Self::Target {
        (**self).upload_value(arena)
    }

    fn upload(&self, arena: &Arena) -> Ptr<Self::Target> {
        (**self).upload(arena)
    }
}

impl<T: DeviceRepr> DeviceRepr for Box<T> {
    type Target = T::Target;

    fn upload_value(&self, arena: &Arena) -> Self::Target {
        (**self).upload_value(arena)
    }

    fn upload(&self, arena: &Arena) -> Ptr<Self::Target> {
        (**self).upload(arena)
    }
}

impl DeviceRepr for Shape {
    type Target = Shape;
    fn upload_value(&self, _arena: &Arena) -> Shape {
        self.clone()
    }
}

impl DeviceRepr for Light {
    type Target = Light;
    fn upload_value(&self, _arena: &Arena) -> Light {
        self.clone()
    }
}

impl DeviceRepr for Spectrum {
    type Target = Spectrum;
    fn upload_value(&self, _arena: &Arena) -> Spectrum {
        self.clone()
    }
}

impl DeviceRepr for Material {
    type Target = Material;
    fn upload_value(&self, _arena: &Arena) -> Material {
        self.clone()
    }
}

impl DeviceRepr for Image {
    type Target = DeviceImage;
    fn upload_value(&self, _arena: &Arena) -> DeviceImage {
        *self.device()
    }
}

impl DeviceRepr for DenselySampledSpectrumBuffer {
    type Target = DenselySampledSpectrum;
    fn upload_value(&self, _arena: &Arena) -> DenselySampledSpectrum {
        self.device()
    }
}

impl DeviceRepr for RGBToSpectrumTable {
    type Target = RGBToSpectrumTable;

    fn upload_value(&self, arena: &Arena) -> RGBToSpectrumTable {
        let n_nodes = self.n_nodes as usize;

        // Safety: these Ptrs point into static or previously-uploaded data;
        // we're copying the contents into the arena for a new lifetime.
        let z_slice = unsafe { from_raw_parts(self.z_nodes.as_raw(), n_nodes) };
        let (z_ptr, _) = arena.alloc_slice(z_slice);

        let n_coeffs = 3 * n_nodes.pow(3);
        let coeffs_slice = unsafe { from_raw_parts(self.coeffs.as_raw(), n_coeffs) };
        let (c_ptr, _) = arena.alloc_slice(coeffs_slice);

        RGBToSpectrumTable {
            z_nodes: z_ptr,
            coeffs: c_ptr,
            n_nodes: self.n_nodes,
        }
    }
}

impl DeviceRepr for RGBColorSpace {
    type Target = RGBColorSpace;

    fn upload_value(&self, arena: &Arena) -> RGBColorSpace {
        let table_ptr = self.rgb_to_spectrum_table.upload(arena);

        RGBColorSpace {
            r: self.r,
            g: self.g,
            b: self.b,
            w: self.w,
            illuminant: self.illuminant.clone(),
            rgb_to_spectrum_table: table_ptr,
            xyz_from_rgb: self.xyz_from_rgb,
            rgb_from_xyz: self.rgb_from_xyz,
        }
    }
}

impl DeviceRepr for DeviceStandardColorSpaces {
    type Target = DeviceStandardColorSpaces;

    fn upload_value(&self, arena: &Arena) -> DeviceStandardColorSpaces {
        DeviceStandardColorSpaces {
            srgb: self.srgb.upload(arena),
            dci_p3: self.dci_p3.upload(arena),
            rec2020: self.rec2020.upload(arena),
            aces2065_1: self.aces2065_1.upload(arena),
        }
    }
}