More fixes to allow for spirv compilation

.gitignore

@@ -4,9 +4,11 @@ target/
 *.bak
 flip.rs
 .vscode
-rust-analyzer.json
 rust-analyzer.toml
 data/
 src/gpu/
 src/tests/
 tests/
+*.spv
+*.json
+*.txt

shared/src/bxdfs/complex.rs

@@ -9,6 +9,7 @@ use crate::core::scattering::{
     refract,
 };
 use crate::spectra::{DeviceStandardColorSpaces, RGBUnboundedSpectrum, SampledSpectrum};
+use crate::utils::gpu_array_from_fn;
 use crate::utils::math::{
     clamp, fast_exp, i0, lerp, log_i0, radians, safe_acos, safe_asin, safe_sqrt, sample_discrete,
     square, trimmed_logistic,
@@ -82,7 +83,7 @@ impl HairBxDF {
         let ap0 = SampledSpectrum::new(f);
         let ap1 = t * (1.0 - f).powi(2);
         let tf = t * f;
-        core::array::from_fn(|p| match p {
+        gpu_array_from_fn(|p| match p {
             0 => ap0,
             1 => ap1,
             _ if p < P_MAX => ap1 * tf.pow_int(p - 1),
@@ -134,7 +135,7 @@ impl HairBxDF {
         let t = t_value.exp();
         let ap = Self::ap(cos_theta_o, self.eta, self.h, t);
         let sum_y: Float = ap.iter().map(|s| s.average()).sum();
-        core::array::from_fn(|i| ap[i].average() / sum_y)
+        gpu_array_from_fn(|i| ap[i].average() / sum_y)
     }
 
     pub fn sigma_a_from_concentration(

shared/src/bxdfs/diffuse.rs

@@ -68,9 +68,7 @@ impl BxDFTrait for DiffuseBxDF {
         self
     }
 
-    fn regularize(&mut self) {
-        return;
-    }
+    fn regularize(&mut self) {}
 }
 
 #[repr(C)]

shared/src/bxdfs/layered.rs

@@ -32,9 +32,10 @@ use core::any::Any;
 use num_traits::Float as NumFloat;
 
 #[derive(Copy, Clone)]
-pub enum TopOrBottom<'a, T, B> {
-    Top(&'a T),
-    Bottom(&'a B),
+pub struct TopOrBottom<'a, T, B> {
+    top: &'a T,
+    bottom: &'a B,
+    is_top: bool,
 }
 
 impl<'a, T, B> TopOrBottom<'a, T, B>
@@ -42,13 +43,34 @@ where
     T: BxDFTrait,
     B: BxDFTrait,
 {
-    pub fn f(&self, wo: Vector3f, wi: Vector3f, mode: TransportMode) -> SampledSpectrum {
-        match self {
-            Self::Top(t) => t.f(wo, wi, mode),
-            Self::Bottom(b) => b.f(wo, wi, mode),
+    #[inline]
+    pub fn new_top(top: &'a T, bottom: &'a B) -> Self {
+        Self {
+            top,
+            bottom,
+            is_top: true,
         }
     }
 
+    #[inline]
+    pub fn new_bottom(top: &'a T, bottom: &'a B) -> Self {
+        Self {
+            top,
+            bottom,
+            is_top: false,
+        }
+    }
+
+    #[inline]
+    pub fn f(&self, wo: Vector3f, wi: Vector3f, mode: TransportMode) -> SampledSpectrum {
+        if self.is_top {
+            self.top.f(wo, wi, mode)
+        } else {
+            self.bottom.f(wo, wi, mode)
+        }
+    }
+
+    #[inline]
     pub fn sample_f(
         &self,
         wo: Vector3f,
@@ -56,23 +78,28 @@ where
         u: Point2f,
         f_args: FArgs,
     ) -> Option<BSDFSample> {
-        match self {
-            Self::Top(t) => t.sample_f(wo, uc, u, f_args),
-            Self::Bottom(b) => b.sample_f(wo, uc, u, f_args),
+        if self.is_top {
+            self.top.sample_f(wo, uc, u, f_args)
+        } else {
+            self.bottom.sample_f(wo, uc, u, f_args)
         }
     }
 
+    #[inline]
     pub fn pdf(&self, wo: Vector3f, wi: Vector3f, f_args: FArgs) -> Float {
-        match self {
-            Self::Top(t) => t.pdf(wo, wi, f_args),
-            Self::Bottom(b) => b.pdf(wo, wi, f_args),
+        if self.is_top {
+            self.top.pdf(wo, wi, f_args)
+        } else {
+            self.bottom.pdf(wo, wi, f_args)
         }
     }
 
+    #[inline]
     pub fn flags(&self) -> BxDFFlags {
-        match self {
-            Self::Top(t) => t.flags(),
-            Self::Bottom(b) => b.flags(),
+        if self.is_top {
+            self.top.flags()
+        } else {
+            self.bottom.flags()
         }
     }
 }
@@ -99,6 +126,7 @@ where
     T: BxDFTrait,
     B: BxDFTrait,
 {
+    #[allow(clippy::too_many_arguments)]
     pub fn new(
         top: T,
         bottom: B,
@@ -136,11 +164,11 @@ where
         mode: TransportMode,
         entered_top: bool,
         exit_z: Float,
-        interfaces: (TopOrBottom<T, B>, TopOrBottom<T, B>, TopOrBottom<T, B>),
+        enter_interface: TopOrBottom<T, B>,
+        exit_interface: TopOrBottom<T, B>,
+        non_exit_interface: TopOrBottom<T, B>,
         rng: &mut Rng,
     ) -> SampledSpectrum {
-        let (enter_interface, exit_interface, non_exit_interface) = interfaces;
-
         let trans_args = FArgs {
             mode,
             sample_flags: BxDFReflTransFlags::TRANSMISSION,
@@ -151,7 +179,7 @@ where
         };
         let mut r = || rng.uniform::<Float>().min(ONE_MINUS_EPSILON);
 
-        // 1. Sample Initial Directions (Standard NEE-like logic)
+        // Sample Initial Directions
         let Some(wos) = enter_interface
             .sample_f(wo, r(), Point2f::new(r(), r()), trans_args)
             .filter(|s| !s.f.is_black() && s.pdf > 0.0 && s.wi.z() != 0.0)
@@ -242,8 +270,7 @@ where
             }
 
             if z == exit_z {
-                // Account for reflection at exitInterface
-                // Hitting the exit surface -> Transmission
+                // Hitting the exit surface -> Reflection off exit interface
                 let Some(bs) = exit_interface
                     .sample_f(-w, r(), Point2f::new(r(), r()), refl_args)
                     .filter(|s| !s.f.is_black() && s.pdf > 0.0 && s.wi.z() != 0.0)
@@ -345,20 +372,20 @@ where
 
         let entered_top = TWO_SIDED || wo.z() > 0.;
         let enter_interface = if entered_top {
-            TopOrBottom::Top(&self.top)
+            TopOrBottom::new_top(&self.top, &self.bottom)
         } else {
-            TopOrBottom::Bottom(&self.bottom)
+            TopOrBottom::new_bottom(&self.top, &self.bottom)
         };
 
         let (exit_interface, non_exit_interface) = if same_hemisphere(wo, wi) ^ entered_top {
             (
-                TopOrBottom::Bottom(&self.bottom),
-                TopOrBottom::Top(&self.top),
+                TopOrBottom::new_bottom(&self.top, &self.bottom),
+                TopOrBottom::new_top(&self.top, &self.bottom),
             )
         } else {
             (
-                TopOrBottom::Top(&self.top),
-                TopOrBottom::Bottom(&self.bottom),
+                TopOrBottom::new_top(&self.top, &self.bottom),
+                TopOrBottom::new_bottom(&self.top, &self.bottom),
             )
         };
 
@@ -376,9 +403,18 @@ where
         let hash1 = hash_buffer(&[wi.x(), wi.y(), wi.z()], 0);
         let mut rng = Rng::new_with_offset(hash0, hash1);
 
-        let inters = (enter_interface, exit_interface, non_exit_interface);
         for _ in 0..self.n_samples {
-            f += self.evaluate_sample(wo, wi, mode, entered_top, exit_z, inters.clone(), &mut rng)
+            f += self.evaluate_sample(
+                wo,
+                wi,
+                mode,
+                entered_top,
+                exit_z,
+                enter_interface.clone(),
+                exit_interface.clone(),
+                non_exit_interface.clone(),
+                &mut rng,
+            );
         }
 
         f / self.n_samples as Float
@@ -474,9 +510,9 @@ where
             }
 
             let interface = if z == 0. {
-                TopOrBottom::Bottom(&self.bottom)
+                TopOrBottom::new_bottom(&self.top, &self.bottom)
             } else {
-                TopOrBottom::Top(&self.top)
+                TopOrBottom::new_top(&self.top, &self.bottom)
             };
 
             // Sample interface BSDF to determine new path direction
@@ -553,13 +589,13 @@ where
                 // Evaluate TRT term for PDF estimate
                 let (r_interface, t_interface) = if entered_top {
                     (
-                        TopOrBottom::Bottom(&self.bottom),
-                        TopOrBottom::Top(&self.top),
+                        TopOrBottom::new_bottom(&self.top, &self.bottom),
+                        TopOrBottom::new_top(&self.top, &self.bottom),
                     )
                 } else {
                     (
-                        TopOrBottom::Top(&self.top),
-                        TopOrBottom::Bottom(&self.bottom),
+                        TopOrBottom::new_top(&self.top, &self.bottom),
+                        TopOrBottom::new_bottom(&self.top, &self.bottom),
                     )
                 };
 
@@ -589,19 +625,19 @@ where
                     }
                 }
             } else {
-                // Evaluate TT term for PDF estimate>
+                // Evaluate TT term for PDF estimate
                 let valid = |s: &BSDFSample| {
                     !s.f.is_black() && s.pdf > 0.0 && s.wi.z() > 0. || s.is_reflective()
                 };
                 let (to_interface, ti_interface) = if entered_top {
                     (
-                        TopOrBottom::Top(&self.top),
-                        TopOrBottom::Bottom(&self.bottom),
+                        TopOrBottom::new_top(&self.top, &self.bottom),
+                        TopOrBottom::new_bottom(&self.top, &self.bottom),
                     )
                 } else {
                     (
-                        TopOrBottom::Bottom(&self.bottom),
-                        TopOrBottom::Top(&self.top),
+                        TopOrBottom::new_bottom(&self.top, &self.bottom),
+                        TopOrBottom::new_top(&self.top, &self.bottom),
                     )
                 };
 

shared/src/core/film.rs

@@ -18,7 +18,7 @@ use crate::utils::math::linear_least_squares;
 use crate::utils::math::{SquareMatrix, wrap_equal_area_square};
 use crate::utils::sampling::VarianceEstimator;
 use crate::utils::transform::AnimatedTransform;
-use crate::utils::{AtomicFloat, Ptr};
+use crate::utils::{AtomicFloat, Ptr, gpu_array_from_fn};
 use num_traits::Float as NumFloat;
 
 #[repr(C)]
@@ -33,13 +33,23 @@ pub struct RGBFilm {
 }
 
 #[repr(C)]
-#[derive(Debug, Clone, Default)]
+#[derive(Debug, Clone)]
 pub struct RGBPixel {
     rgb_sum: [AtomicFloat; 3],
     weight_sum: AtomicFloat,
     rgb_splat: [AtomicFloat; 3],
 }
 
+impl Default for RGBPixel {
+    fn default() -> Self {
+        Self {
+            rgb_sum: gpu_array_from_fn(|_| AtomicFloat::default()),
+            weight_sum: AtomicFloat::default(),
+            rgb_splat: gpu_array_from_fn(|_| AtomicFloat::default()),
+        }
+    }
+}
+
 impl RGBFilm {
     pub fn base(&self) -> &FilmBase {
         &self.base
@@ -153,7 +163,7 @@ impl RGBFilm {
 }
 
 #[repr(C)]
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Clone)]
 #[cfg_attr(target_os = "cuda", derive(Copy))]
 pub struct GBufferPixel {
     pub rgb_sum: [AtomicFloat; 3],
@@ -162,7 +172,7 @@ pub struct GBufferPixel {
     pub rgb_splat: [AtomicFloat; 3],
     pub p_sum: Point3f,
     pub dz_dx_sum: AtomicFloat,
-    pub dz_dy_sum: Float,
+    pub dz_dy_sum: AtomicFloat,
     pub n_sum: Normal3f,
     pub ns_sum: Normal3f,
     pub uv_sum: Point2f,
@@ -170,6 +180,25 @@ pub struct GBufferPixel {
     pub rgb_variance: VarianceEstimator,
 }
 
+impl Default for GBufferPixel {
+    fn default() -> Self {
+        Self {
+            rgb_sum: gpu_array_from_fn(|_| AtomicFloat::default()),
+            weight_sum: AtomicFloat::default(),
+            rgb_splat: gpu_array_from_fn(|_| AtomicFloat::default()),
+            g_buffer_weight_sum: AtomicFloat::default(),
+            p_sum: Point3f::default(),
+            dz_dx_sum: AtomicFloat::default(),
+            dz_dy_sum: AtomicFloat::default(),
+            n_sum: Normal3f::default(),
+            ns_sum: Normal3f::default(),
+            uv_sum: Point2f::default(),
+            rgb_albedo_sum: gpu_array_from_fn(|_| AtomicFloat::default()),
+            rgb_variance: VarianceEstimator::default(),
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Debug, Clone)]
 #[cfg_attr(target_os = "cuda", derive(Copy))]
@@ -299,9 +328,9 @@ pub struct SpectralPixel {
 impl Clone for SpectralPixel {
     fn clone(&self) -> Self {
         Self {
-            rgb_sum: core::array::from_fn(|i| AtomicFloat::new(self.rgb_sum[i].get())),
+            rgb_sum: gpu_array_from_fn(|i| AtomicFloat::new(self.rgb_sum[i].get())),
             rgb_weight_sum: AtomicFloat::new(self.rgb_weight_sum.get()),
-            rgb_splat: core::array::from_fn(|i| AtomicFloat::new(self.rgb_splat[i].get())),
+            rgb_splat: gpu_array_from_fn(|i| AtomicFloat::new(self.rgb_splat[i].get())),
             bucket_offset: self.bucket_offset,
         }
     }
@@ -310,9 +339,9 @@ impl Clone for SpectralPixel {
 impl Default for SpectralPixel {
     fn default() -> Self {
         Self {
-            rgb_sum: core::array::from_fn(|_| AtomicFloat::new(0.0)),
+            rgb_sum: gpu_array_from_fn(|_| AtomicFloat::new(0.0)),
             rgb_weight_sum: AtomicFloat::new(0.0),
-            rgb_splat: core::array::from_fn(|_| AtomicFloat::new(0.0)),
+            rgb_splat: gpu_array_from_fn(|_| AtomicFloat::new(0.0)),
             bucket_offset: 0,
         }
     }

shared/src/core/geometry/bounds.rs

@@ -2,6 +2,7 @@ use super::{Float, NumFloat};
 use super::{Point, Point2f, Point3, Point3f, Vector, Vector2, Vector2f, Vector3, Vector3f};
 use crate::core::geometry::traits::{SqrtExt, VectorLike};
 use crate::core::geometry::{max, min};
+use crate::utils::gpu_array_from_fn;
 use crate::utils::interval::Interval;
 use crate::utils::math::lerp;
 use core::mem;
@@ -137,7 +138,7 @@ where
     }
 
     pub fn corner(&self, corner_index: usize) -> Point<T, N> {
-        Point(core::array::from_fn(|i| {
+        Point(gpu_array_from_fn(|i| {
             if (corner_index >> i) & 1 == 1 {
                 self.p_max[i]
             } else {

shared/src/core/geometry/traits.rs

@@ -1,4 +1,5 @@
 use crate::core::pbrt::Float;
+use crate::utils::gpu_array_from_fn;
 use crate::utils::interval::Interval;
 use crate::utils::math::{next_float_down, next_float_up};
 use core::ops::{Add, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub};
@@ -18,7 +19,7 @@ pub trait Tuple<T, const N: usize>:
     where
         T: Copy,
     {
-        let new_data = p.map(|index| self[index]);
+        let new_data = gpu_array_from_fn(|i| self[p[i]]);
         Self::from_array(new_data)
     }
 

shared/src/core/image.rs

@@ -69,10 +69,55 @@ impl PixelFormat {
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug)]
-pub enum Pixels {
-    U8(Ptr<u8>),
-    F16(Ptr<u16>),
-    F32(Ptr<f32>),
+pub struct Pixels {
+    ptr: Ptr<u8>,
+    format: PixelFormat,
+}
+
+impl Pixels {
+    pub fn new_u8(ptr: Ptr<u8>) -> Self {
+        Self {
+            ptr,
+            format: PixelFormat::U8,
+        }
+    }
+
+    pub fn new_f16(ptr: Ptr<u16>) -> Self {
+        Self {
+            ptr: Ptr::from_raw(ptr.as_raw() as *const u8),
+            format: PixelFormat::F16,
+        }
+    }
+
+    pub fn new_f32(ptr: Ptr<f32>) -> Self {
+        Self {
+            ptr: Ptr::from_raw(ptr.as_raw() as *const u8),
+            format: PixelFormat::F32,
+        }
+    }
+
+    unsafe fn read_u8(&self, byte_offset: usize) -> u8 {
+        unsafe { *self.ptr.as_raw().add(byte_offset) }
+    }
+
+    unsafe fn read_f16(&self, elem_offset: usize) -> u16 {
+        let byte_offset = elem_offset * 2;
+        unsafe { *(self.ptr.as_raw().add(byte_offset) as *const u16) }
+    }
+
+    unsafe fn read_f32(&self, elem_offset: usize) -> f32 {
+        let byte_offset = elem_offset * 4;
+        unsafe { *(self.ptr.as_raw().add(byte_offset) as *const f32) }
+    }
+
+    // Unsure if ill need this
+    // pub unsafe fn read(&self, offset: usize) -> Float {
+    //     match self.format {
+    //         PixelFormat::U8 => unsafe { self.read_u8(offset) as Float / 255.0 },
+    //         PixelFormat::F16 => unsafe { f16_to_f32_software(self.read_f16(offset)) },
+    //         PixelFormat::F32 => unsafe { self.read_f32(offset) },
+    //     }
+    // }
 }
 
 #[repr(C)]
@@ -170,18 +215,18 @@ impl ImageAccess for DeviceImage {
             return 0.;
         }
 
-        let offset = self.pixel_offset(p) + c as u32;
+        let offset = (self.pixel_offset(p) + c as u32) as usize;
         unsafe {
-            match self.pixels {
-                Pixels::U8(ptr) => {
-                    let raw_val = *ptr.add(offset as usize);
-                    self.base().encoding.to_linear_scalar(raw_val)
+            match self.pixels.format {
+                PixelFormat::U8 => {
+                    let raw_val = self.pixels.read_u8(offset);
+                    self.base.encoding.to_linear_scalar(raw_val)
                 }
-                Pixels::F16(ptr) => {
-                    let raw_val = *ptr.add(offset as usize);
+                PixelFormat::F16 => {
+                    let raw_val = self.pixels.read_f16(offset);
                     f16_to_f32_software(raw_val)
                 }
-                Pixels::F32(ptr) => *ptr.add(offset as usize),
+                PixelFormat::F32 => self.pixels.read_f32(offset),
             }
         }
     }

shared/src/core/material.rs

@@ -74,7 +74,7 @@ impl From<&SurfaceInteraction> for NormalBumpEvalContext {
             p: si.p(),
             uv: si.common.uv,
             n: si.n(),
-            shading: si.shading.clone(),
+            shading: si.shading,
             dudx: si.dudx,
             dudy: si.dudy,
             dvdx: si.dvdx,

shared/src/core/sampler.rs

@@ -91,7 +91,7 @@ pub enum RandomizeStrategy {
 }
 
 #[repr(C)]
-#[derive(Default, Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy)]
 pub struct HaltonSampler {
     pub samples_per_pixel: i32,
     pub randomize: RandomizeStrategy,
@@ -103,6 +103,22 @@ pub struct HaltonSampler {
     pub digit_permutations: Ptr<DeviceDigitPermutation>,
 }
 
+#[allow(clippy::derivable_impls)]
+impl Default for HaltonSampler {
+    fn default() -> Self {
+        Self {
+            samples_per_pixel: 0,
+            randomize: RandomizeStrategy::default(),
+            base_scales: [0; 2],
+            base_exponents: [0; 2],
+            mult_inverse: [0; 2],
+            halton_index: 0,
+            dim: 0,
+            digit_permutations: Ptr::default(),
+        }
+    }
+}
+
 impl HaltonSampler {
     pub fn sample_dimension(&self, dimension: u32) -> Float {
         if self.randomize == RandomizeStrategy::None {

shared/src/lib.rs

@@ -5,6 +5,7 @@
 pub mod bxdfs;
 pub mod cameras;
 pub mod core;
+#[cfg(not(target_arch = "spirv"))]
 pub mod data;
 pub mod filters;
 pub mod lights;

shared/src/lights/sampler.rs

@@ -12,8 +12,8 @@ use crate::{Float, ONE_MINUS_EPSILON, PI};
 use enum_dispatch::enum_dispatch;
 use num_traits::Float as NumFloat;
 
-#[derive(Clone, Copy, Debug, Default)]
 #[repr(C)]
+#[derive(Clone, Copy, Debug)]
 pub struct CompactLightBounds {
     pub w: OctahedralVector,
     pub phi: Float,
@@ -27,6 +27,18 @@ pub struct CompactLightBounds {
     pub qb: [[u16; 3]; 2],
 }
 
+#[allow(clippy::derivable_impls)]
+impl Default for CompactLightBounds {
+    fn default() -> Self {
+        Self {
+            w: OctahedralVector::default(),
+            phi: Float::default(),
+            packed_info: u32::default(),
+            qb: [[u16::default(); 3]; 2],
+        }
+    }
+}
+
 const _: () = assert!(core::mem::size_of::<CompactLightBounds>() == 24);
 
 impl CompactLightBounds {
@@ -269,17 +281,13 @@ impl LightSamplerTrait for PowerLightSampler {
     }
 
     fn pmf(&self, light: &Light) -> Float {
-        let array_start = self.lights.as_raw();
-        let target = light as *const Light as *mut Light;
-
-        unsafe {
-            let index = target.offset_from(array_start);
-
-            if index >= 0 && index < self.lights_len as isize {
-                return self.alias_table.pmf(index as u32);
+        let target = light as *const Light;
+        for i in 0..self.lights_len as usize {
+            if unsafe { self.lights.add(i) }.as_raw() == target {
+                return self.alias_table.pmf(i as u32);
             }
         }
-        0.
+        0.0
     }
 }
 
@@ -353,12 +361,12 @@ impl LightBVHNode {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Copy)]
 pub struct BVHLightSampler {
-    pub nodes: *const LightBVHNode,
-    pub lights: *const Light,
-    pub infinite_lights: *const Light,
-    pub bit_trails: *const u64,
+    pub nodes: Ptr<LightBVHNode>,
+    pub lights: Ptr<Light>,
+    pub infinite_lights: Ptr<Light>,
+    pub bit_trails: Ptr<u64>,
     pub nodes_len: u32,
     pub lights_len: u32,
     pub infinite_lights_len: u32,
@@ -371,22 +379,33 @@ unsafe impl Sync for BVHLightSampler {}
 impl BVHLightSampler {
     #[inline(always)]
     fn node(&self, idx: usize) -> &LightBVHNode {
-        unsafe { &*self.nodes.add(idx) }
+        unsafe { self.nodes.at(idx) }
     }
 
     #[inline(always)]
     fn light(&self, idx: usize) -> Light {
-        unsafe { *self.lights.add(idx) }
+        unsafe { *self.lights.at(idx) }
     }
 
     #[inline(always)]
     fn infinite_light(&self, idx: usize) -> Light {
-        unsafe { *self.infinite_lights.add(idx) }
+        unsafe { *self.infinite_lights.at(idx) }
     }
 
     #[inline(always)]
     fn bit_trail(&self, idx: usize) -> u64 {
-        unsafe { *self.bit_trails.add(idx) }
+        unsafe { *self.bit_trails.at(idx) }
+    }
+
+    #[inline(always)]
+    fn light_index_in(&self, base: Ptr<Light>, len: u32, light: &Light) -> Option<usize> {
+        let target = light as *const Light;
+        for i in 0..len as usize {
+            if unsafe { base.add(i) }.as_raw() == target {
+                return Some(i);
+            }
+        }
+        None
     }
 
     fn evaluate_cost(&self, b: &LightBounds, bounds: &Bounds3f, dim: usize) -> Float {
@@ -459,27 +478,20 @@ impl LightSamplerTrait for BVHLightSampler {
     }
 
     fn pmf_with_context(&self, ctx: &LightSampleContext, light: &Light) -> Float {
-        let light_ptr = light as *const Light;
         let empty_nodes = if self.nodes_len == 0 { 0. } else { 1. };
         let n_infinite = self.infinite_lights_len as Float;
-
-        let inf_start = self.infinite_lights;
-        let inf_end = unsafe { self.infinite_lights.add(self.infinite_lights_len as usize) };
-        if light_ptr >= inf_start && light_ptr < inf_end {
+        if self
+            .light_index_in(self.infinite_lights, self.infinite_lights_len, light)
+            .is_some()
+        {
             return 1.0 / (n_infinite + empty_nodes);
         }
 
-        let finite_start = self.lights;
-        let finite_end = unsafe { self.lights.add(self.lights_len as usize) };
-
-        if light_ptr < finite_start || light_ptr >= finite_end {
+        let Some(light_index) = self.light_index_in(self.lights, self.lights_len, light) else {
             return 0.0;
-        }
-
-        let light_index = unsafe { light_ptr.offset_from(finite_start) as usize };
+        };
 
         let mut bit_trail = self.bit_trail(light_index);
-
         let p_inf = n_infinite / (n_infinite + empty_nodes);
         let mut pmf = 1.0 - p_inf;
         let mut node_ind = 0;

shared/src/shapes/curves.rs

@@ -5,6 +5,7 @@ use crate::core::geometry::{
 };
 use crate::core::interaction::{Interaction, InteractionTrait, SurfaceInteraction};
 use crate::core::shape::{ShapeIntersection, ShapeSample, ShapeSampleContext, ShapeTrait};
+use crate::utils::gpu_array_from_fn;
 use crate::utils::math::{clamp, lerp, square};
 use crate::utils::splines::{
     bound_cubic_bezier, cubic_bezier_control_points, evaluate_cubic_bezier, subdivide_cubic_bezier,
@@ -52,7 +53,7 @@ impl CurveCommon {
         assert_eq!(c.len(), 4);
         let cp_obj: [Point3f; 4] = c[..4].try_into().unwrap();
 
-        let mut n = [Normal3f::default(); 2];
+        let mut n: [Normal3f; 2] = gpu_array_from_fn(|_| Normal3f::default());
         let mut normal_angle: Float = 0.;
         let mut inv_sin_normal_angle: Float = 0.;
         if norm.len() == 2 {
@@ -115,7 +116,7 @@ impl CurveShape {
         }
 
         let ray_from_object = look_at(ray.o, ray.o + ray.d, dx).expect("Inversion error");
-        let cp = [0; 4].map(|i| ray_from_object.apply_to_point(cp_obj[i]));
+        let cp: [Point3f; 4] = gpu_array_from_fn(|i| ray_from_object.apply_to_point(cp_obj[i]));
 
         let max_width = lerp(self.u_min, self.common.width[0], self.common.width[1]).max(lerp(
             self.u_max,

shared/src/spectra/sampled.rs

@@ -1,5 +1,6 @@
 use crate::core::pbrt::Float;
 use crate::core::spectrum::{SpectrumTrait, StandardSpectra};
+use crate::utils::gpu_array_from_fn;
 use crate::utils::math::{clamp, lerp};
 use core::ops::{
     Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
@@ -43,7 +44,7 @@ impl SampledSpectrum {
         F: FnMut(usize) -> Float,
     {
         Self {
-            values: core::array::from_fn(cb),
+            values: gpu_array_from_fn(cb),
         }
     }
 

shared/src/utils/math.rs

@@ -4,7 +4,7 @@ use crate::core::pbrt::{Float, FloatBitOps, FloatBits, ONE_MINUS_EPSILON, PI, PI
 use crate::utils::hash::{hash_buffer, mix_bits};
 use crate::utils::sobol::{SOBOL_MATRICES_32, VDC_SOBOL_MATRICES, VDC_SOBOL_MATRICES_INV};
 
-use crate::utils::Ptr;
+use crate::utils::{Ptr, gpu_array_from_fn};
 use core::fmt::{self, Display, Write};
 use core::iter::{Product, Sum};
 use core::mem;
@@ -1090,7 +1090,7 @@ impl<T, const R: usize, const C: usize> Matrix<T, R, C> {
     where
         T: Clone + Zero,
     {
-        let m: [[T; C]; R] = core::array::from_fn(|_| core::array::from_fn(|_| T::zero()));
+        let m: [[T; C]; R] = gpu_array_from_fn(|_| gpu_array_from_fn(|_| T::zero()));
         Self { m }
     }
 
@@ -1258,8 +1258,8 @@ impl<T, const N: usize> SquareMatrix<T, N> {
         T: Copy + Zero + One,
     {
         Self {
-            m: core::array::from_fn(|i| {
-                core::array::from_fn(|j| if i == j { T::one() } else { T::zero() })
+            m: gpu_array_from_fn(|i| {
+                gpu_array_from_fn(|j| if i == j { T::one() } else { T::zero() })
             }),
         }
     }
@@ -1469,8 +1469,7 @@ where
 {
     type Output = Vector<T, N>;
     fn mul(self, rhs: Vector<T, N>) -> Self::Output {
-        let arr =
-            core::array::from_fn(|i| self.m[i].iter().zip(&rhs.0).map(|(m, v)| *m * *v).sum());
+        let arr = gpu_array_from_fn(|i| self.m[i].iter().zip(&rhs.0).map(|(m, v)| *m * *v).sum());
         Vector(arr)
     }
 }
@@ -1481,8 +1480,7 @@ where
 {
     type Output = Point<T, N>;
     fn mul(self, rhs: Point<T, N>) -> Self::Output {
-        let arr =
-            core::array::from_fn(|i| self.m[i].iter().zip(&rhs.0).map(|(m, v)| *m * *v).sum());
+        let arr = gpu_array_from_fn(|i| self.m[i].iter().zip(&rhs.0).map(|(m, v)| *m * *v).sum());
         Point(arr)
     }
 }

shared/src/utils/mod.rs

@@ -115,6 +115,20 @@ impl AtomicFloat {
     }
 }
 
+#[inline(always)]
+pub fn gpu_array_from_fn<T, const N: usize>(mut f: impl FnMut(usize) -> T) -> [T; N] {
+    unsafe {
+        let mut arr: core::mem::MaybeUninit<[T; N]> = core::mem::MaybeUninit::uninit();
+        let ptr = arr.as_mut_ptr() as *mut T;
+        let mut i = 0;
+        while i < N {
+            ptr.add(i).write(f(i));
+            i += 1;
+        }
+        arr.assume_init()
+    }
+}
+
 // pub struct AtomicDouble {
 //     bits: AtomicU64,
 // }

shared/src/utils/splines.rs

@@ -1,5 +1,6 @@
 use crate::core::geometry::{Bounds3f, Lerp, Point3f, Vector3f, VectorLike};
 use crate::core::pbrt::Float;
+use crate::utils::gpu_array_from_fn;
 use crate::utils::math::lerp;
 use core::ops::Sub;
 use num_traits::Num;
@@ -40,7 +41,7 @@ where
 }
 
 pub fn subdivide_cubic_bezier(cp: &[Point3f]) -> [Point3f; 7] {
-    let v: [Vector3f; 4] = core::array::from_fn(|i| Vector3f::from(cp[i]));
+    let v: [Vector3f; 4] = gpu_array_from_fn(|i| Vector3f::from(cp[i]));
     let v01 = (v[0] + v[1]) / 2.0;
     let v12 = (v[1] + v[2]) / 2.0;
     let v23 = (v[2] + v[3]) / 2.0;

shared/src/utils/transform.rs

@@ -12,6 +12,7 @@ use crate::core::geometry::{
 use crate::core::interaction::{
     Interaction, InteractionBase, InteractionTrait, MediumInteraction, SurfaceInteraction,
 };
+use crate::utils::gpu_array_from_fn;
 use crate::{Float, gamma};
 
 #[repr(C)]
@@ -793,7 +794,7 @@ impl DerivativeTerm {
 }
 
 #[repr(C)]
-#[derive(Debug, Copy, Default, Clone)]
+#[derive(Debug, Copy, Clone)]
 pub struct AnimatedTransform {
     pub start_transform: Transform,
     pub end_transform: Transform,
@@ -811,6 +812,27 @@ pub struct AnimatedTransform {
     c5: [DerivativeTerm; 3],
 }
 
+impl Default for AnimatedTransform {
+    fn default() -> Self {
+        Self {
+            start_transform: Transform::default(),
+            end_transform: Transform::default(),
+            start_time: 0.0,
+            end_time: 0.0,
+            actually_animated: false,
+            t: gpu_array_from_fn(|_| Vector3f::default()),
+            r: gpu_array_from_fn(|_| Quaternion::default()),
+            s: gpu_array_from_fn(|_| SquareMatrix::default()),
+            has_rotation: false,
+            c1: gpu_array_from_fn(|_| DerivativeTerm::default()),
+            c2: gpu_array_from_fn(|_| DerivativeTerm::default()),
+            c3: gpu_array_from_fn(|_| DerivativeTerm::default()),
+            c4: gpu_array_from_fn(|_| DerivativeTerm::default()),
+            c5: gpu_array_from_fn(|_| DerivativeTerm::default()),
+        }
+    }
+}
+
 impl AnimatedTransform {
     pub fn from_transform(t: &Transform) -> Self {
         Self::new(t, 0., t, 1.)
@@ -825,22 +847,7 @@ impl AnimatedTransform {
         let actually_animated = start_transform != end_transform;
 
         if !actually_animated {
-            return Self {
-                start_transform: *start_transform,
-                end_transform: *end_transform,
-                start_time,
-                end_time,
-                actually_animated: false,
-                t: [Vector3f::default(); 2],
-                r: [Quaternion::default(); 2],
-                s: core::array::from_fn(|_| SquareMatrix::default()),
-                has_rotation: false,
-                c1: [DerivativeTerm::default(); 3],
-                c2: [DerivativeTerm::default(); 3],
-                c3: [DerivativeTerm::default(); 3],
-                c4: [DerivativeTerm::default(); 3],
-                c5: [DerivativeTerm::default(); 3],
-            };
+            return Self::default();
         }
 
         let (t0, r_temp, s0) = start_transform.decompose();
@@ -865,11 +872,11 @@ impl AnimatedTransform {
         let s = [s0, s1];
 
         let (c1, c2, c3, c4, c5) = if has_rotation {
-            let mut c1: [DerivativeTerm; 3] = [Default::default(); 3];
-            let mut c2: [DerivativeTerm; 3] = [Default::default(); 3];
-            let mut c3: [DerivativeTerm; 3] = [Default::default(); 3];
-            let mut c4: [DerivativeTerm; 3] = [Default::default(); 3];
-            let mut c5: [DerivativeTerm; 3] = [Default::default(); 3];
+            let mut c1: [DerivativeTerm; 3] = gpu_array_from_fn(|_| DerivativeTerm::default());
+            let mut c2: [DerivativeTerm; 3] = gpu_array_from_fn(|_| DerivativeTerm::default());
+            let mut c3: [DerivativeTerm; 3] = gpu_array_from_fn(|_| DerivativeTerm::default());
+            let mut c4: [DerivativeTerm; 3] = gpu_array_from_fn(|_| DerivativeTerm::default());
+            let mut c5: [DerivativeTerm; 3] = gpu_array_from_fn(|_| DerivativeTerm::default());
             let cos_theta = r0.dot(r1);
             let theta = safe_acos(cos_theta);
             let qperp: Quaternion = (r0 - r1 * cos_theta).normalize();
@@ -1965,11 +1972,11 @@ impl AnimatedTransform {
             (c1, c2, c3, c4, c5)
         } else {
             (
-                core::array::from_fn(|_| DerivativeTerm::default()),
-                core::array::from_fn(|_| DerivativeTerm::default()),
-                core::array::from_fn(|_| DerivativeTerm::default()),
-                core::array::from_fn(|_| DerivativeTerm::default()),
-                core::array::from_fn(|_| DerivativeTerm::default()),
+                gpu_array_from_fn(|_| DerivativeTerm::default()),
+                gpu_array_from_fn(|_| DerivativeTerm::default()),
+                gpu_array_from_fn(|_| DerivativeTerm::default()),
+                gpu_array_from_fn(|_| DerivativeTerm::default()),
+                gpu_array_from_fn(|_| DerivativeTerm::default()),
             )
         };
         AnimatedTransform {

src/utils/sampling.rs

@@ -4,11 +4,11 @@ use crate::utils::backend::GpuAllocator;
 use crate::utils::containers::Array2D;
 use shared::Float;
 use shared::core::geometry::{Bounds2f, Point2i, Vector2f, Vector2i};
-use shared::utils::Ptr;
 use shared::utils::sampling::{
     AliasTable, Bin, DevicePiecewiseConstant1D, DevicePiecewiseConstant2D, DeviceSummedAreaTable,
     DeviceWindowedPiecewiseConstant2D, PiecewiseLinear2D,
 };
+use shared::utils::{Ptr, gpu_array_from_fn};
 use std::sync::Arc;
 
 #[derive(Debug, Clone)]
@@ -240,7 +240,7 @@ impl<const N: usize> PiecewiseLinear2DHost<N> {
 
         let mut param_size = [0u32; N];
         let mut param_strides = [0u32; N];
-        let param_values = std::array::from_fn(|i| param_values[i].to_vec());
+        let param_values = gpu_array_from_fn(|i| param_values[i].to_vec());
 
         let mut slices: u32 = 1;
         for i in (0..N).rev() {