Skip to main content

concinnity_core/build/environment_map/
bake.rs

1//! The IBL convolutions themselves: cube sampling, the two kernels, and the
2//! row-sized unit of work they decompose into.
3//!
4//! Every output texel is an independent integral over a read-only source, so a
5//! bake is expressed as [`CubeBake`] plus the rows of its output. The caller
6//! decides whether to run those rows one at a time or all at once; this module
7//! spawns nothing and takes no thread, so the two produce identical bytes.
8
9use crate::math::{cos, floor, sin, sin_cos, sqrt};
10use alloc::vec;
11use alloc::vec::Vec;
12
13use crate::gfx::cubemap;
14use crate::math::vec3::{cross as cross3, dot as dot3, length};
15
16/// Default azimuthal samples per irradiance texel.
17pub const DEFAULT_IRRADIANCE_PHI_SAMPLES: u32 = 64;
18/// Default polar samples per irradiance texel.
19pub const DEFAULT_IRRADIANCE_THETA_SAMPLES: u32 = 16;
20
21// Cube sampling
22
23// Cube-face sampler: project a normalised direction onto the dominant axis
24// to pick a face, then bilinearly sample within that face. Edges are clamped
25// per-face (no seamless filtering).
26fn sample_cube(faces: &[Vec<f32>; 6], face_size: u32, dir: [f32; 3]) -> [f32; 3] {
27    let ax = dir[0].abs();
28    let ay = dir[1].abs();
29    let az = dir[2].abs();
30    let (face, ma, s, t) = if ax >= ay && ax >= az {
31        if dir[0] > 0.0 {
32            (0usize, ax, -dir[2], -dir[1])
33        } else {
34            (1, ax, dir[2], -dir[1])
35        }
36    } else if ay >= az {
37        if dir[1] > 0.0 {
38            (2usize, ay, dir[0], dir[2])
39        } else {
40            (3, ay, dir[0], -dir[2])
41        }
42    } else if dir[2] > 0.0 {
43        (4usize, az, dir[0], -dir[1])
44    } else {
45        (5, az, -dir[0], -dir[1])
46    };
47    let inv = 0.5 / ma.max(1e-20);
48    let fs = face_size as f32;
49    // s, t in [-1, 1] after multiplying by inv*2; map to pixel coords.
50    let fx = (s * inv + 0.5) * fs - 0.5;
51    let fy = (t * inv + 0.5) * fs - 0.5;
52    let x0 = (floor(fx) as i32).clamp(0, face_size as i32 - 1);
53    let y0 = (floor(fy) as i32).clamp(0, face_size as i32 - 1);
54    let x1 = (x0 + 1).clamp(0, face_size as i32 - 1);
55    let y1 = (y0 + 1).clamp(0, face_size as i32 - 1);
56    let dx = (fx - floor(fx)).clamp(0.0, 1.0);
57    let dy = (fy - floor(fy)).clamp(0.0, 1.0);
58    let p = |x: i32, y: i32| -> [f32; 3] {
59        let off = ((y as usize) * face_size as usize + x as usize) * 4;
60        let face_data = &faces[face];
61        [face_data[off], face_data[off + 1], face_data[off + 2]]
62    };
63    let p00 = p(x0, y0);
64    let p10 = p(x1, y0);
65    let p01 = p(x0, y1);
66    let p11 = p(x1, y1);
67    let w00 = (1.0 - dx) * (1.0 - dy);
68    let w10 = dx * (1.0 - dy);
69    let w01 = (1.0 - dx) * dy;
70    let w11 = dx * dy;
71    [
72        p00[0] * w00 + p10[0] * w10 + p01[0] * w01 + p11[0] * w11,
73        p00[1] * w00 + p10[1] * w10 + p01[1] * w01 + p11[1] * w11,
74        p00[2] * w00 + p10[2] * w10 + p01[2] * w01 + p11[2] * w11,
75    ]
76}
77
78fn normalize3(v: [f32; 3]) -> [f32; 3] {
79    let l = length(v).max(1e-20);
80    [v[0] / l, v[1] / l, v[2] / l]
81}
82
83// Build an orthonormal basis around `n` (N = up axis). Returns (tangent, bitangent).
84fn make_tbn(n: [f32; 3]) -> ([f32; 3], [f32; 3]) {
85    let up = if n[2].abs() < 0.999 {
86        [0.0, 0.0, 1.0]
87    } else {
88        [1.0, 0.0, 0.0]
89    };
90    let t = normalize3(cross3(up, n));
91    let b = cross3(n, t);
92    (t, b)
93}
94
95// Hammersley + GGX importance sampling
96
97// Hammersley quasi-random 2D sequence over `n` samples. Used to drive the
98// GGX importance sampler for prefilter convolution.
99fn hammersley(i: u32, n: u32) -> [f32; 2] {
100    let mut bits = i;
101    bits = bits.rotate_right(16);
102    bits = ((bits & 0x5555_5555) << 1) | ((bits & 0xAAAA_AAAA) >> 1);
103    bits = ((bits & 0x3333_3333) << 2) | ((bits & 0xCCCC_CCCC) >> 2);
104    bits = ((bits & 0x0F0F_0F0F) << 4) | ((bits & 0xF0F0_F0F0) >> 4);
105    bits = ((bits & 0x00FF_00FF) << 8) | ((bits & 0xFF00_FF00) >> 8);
106    let radical_inverse = (bits as f32) * 2.328_306_4e-10; // 1 / 2^32
107    [i as f32 / n as f32, radical_inverse]
108}
109
110// Sample the GGX distribution in world space around normal `n`. Returns a
111// half-vector H. The caller supplies the tangent basis and `a2m1` (a^2 - 1),
112// both constant across a texel's whole sample set.
113fn importance_sample_ggx(
114    xi: [f32; 2],
115    n: [f32; 3],
116    basis: ([f32; 3], [f32; 3]),
117    a2m1: f32,
118) -> [f32; 3] {
119    let (t, b) = basis;
120    let phi = 2.0 * core::f32::consts::PI * xi[0];
121    let cos_theta = sqrt((1.0 - xi[1]) / (1.0 + a2m1 * xi[1]));
122    let sin_theta = sqrt((1.0 - cos_theta * cos_theta).max(0.0));
123    let (sin_phi, cos_phi) = sin_cos(phi);
124    let h_local = [sin_theta * cos_phi, sin_theta * sin_phi, cos_theta];
125    normalize3([
126        t[0] * h_local[0] + b[0] * h_local[1] + n[0] * h_local[2],
127        t[1] * h_local[0] + b[1] * h_local[1] + n[1] * h_local[2],
128        t[2] * h_local[0] + b[2] * h_local[1] + n[2] * h_local[2],
129    ])
130}
131
132// Cap a sampled radiance so a single very bright source texel (a sun disk, a
133// blown sky highlight) cannot dominate a glossy reflection mip. Scales RGB
134// uniformly to keep its hue when luminance exceeds `clamp`; `clamp <= 0`
135// disables the cap. This suppresses the lone-hot-texel "bright squares" a
136// clear-sky HDR otherwise smears across reflective floors. The rough mips always
137// pass through here; mip 0 only does for a reflection probe (`clamp_mip0`), never
138// for an imported environment map, so the on-screen skybox keeps its true HDR.
139fn clamp_radiance(rgb: [f32; 3], clamp: f32) -> [f32; 3] {
140    if clamp <= 0.0 {
141        return rgb;
142    }
143    let lum = 0.2126 * rgb[0] + 0.7152 * rgb[1] + 0.0722 * rgb[2];
144    if lum > clamp {
145        let s = clamp / lum;
146        return [rgb[0] * s, rgb[1] * s, rgb[2] * s];
147    }
148    rgb
149}
150
151// The chunked bake
152
153/// One row of a bake's output: the RGBA texels of row `y` on cube face `face`.
154///
155/// Handed out by [`face_rows`] and filled by [`CubeBake::compute_row`]. A row
156/// borrows only its own texels, so the whole set can be worked through in any
157/// order and on any thread.
158pub struct FaceRow<'a> {
159    face: usize,
160    y: u32,
161    texels: &'a mut [f32],
162}
163
164/// Split six output faces of `face_size` edge into their rows, in face-major
165/// order.
166///
167/// This is the unit of work an environment-map bake decomposes into: each row
168/// reads only the immutable source and writes only the texels it holds, so a
169/// caller is free to fan the set across a thread pool.
170pub fn face_rows(faces: &mut [Vec<f32>; 6], face_size: u32) -> Vec<FaceRow<'_>> {
171    let stride = face_size as usize * 4;
172    faces
173        .iter_mut()
174        .enumerate()
175        .flat_map(|(face, data)| {
176            data.chunks_mut(stride)
177                .enumerate()
178                .map(move |(y, texels)| FaceRow {
179                    face,
180                    y: y as u32,
181                    texels,
182                })
183        })
184        .collect()
185}
186
187// The per-texel integral a bake evaluates. Both kernels precompute the terms
188// that are constant across a whole bake so the inner loops do not.
189enum Kernel {
190    Irradiance(IrradianceKernel),
191    Ggx(GgxKernel),
192}
193
194struct IrradianceKernel {
195    phi_samples: u32,
196    theta_samples: u32,
197    inv_n_phi: f32,
198    inv_n_theta: f32,
199    weight: f32,
200}
201
202struct GgxKernel {
203    samples: u32,
204    a2m1: f32,
205    clamp: f32,
206}
207
208/// A convolution of a source cubemap into a new cubemap, decomposed into
209/// independent output rows.
210///
211/// [`compute`](Self::compute) runs the whole thing in one call; a caller that
212/// wants the rows spread over worker threads takes [`output_faces`] and
213/// [`face_rows`] and drives [`compute_row`] itself. Both produce the same
214/// bytes: no row observes another.
215///
216/// [`output_faces`]: Self::output_faces
217/// [`compute_row`]: Self::compute_row
218pub struct CubeBake<'a> {
219    source: &'a [Vec<f32>; 6],
220    source_face_size: u32,
221    output_face_size: u32,
222    kernel: Kernel,
223}
224
225impl<'a> CubeBake<'a> {
226    /// Cosine-weighted hemisphere integral over each output direction, by
227    /// uniform (phi, theta) sampling. The result includes the cosine and
228    /// Jacobian terms, so a shader plugs it straight in as
229    /// `irradiance / π * albedo`.
230    pub fn irradiance(
231        source: &'a [Vec<f32>; 6],
232        source_face_size: u32,
233        output_face_size: u32,
234        phi_samples: u32,
235        theta_samples: u32,
236    ) -> Self {
237        let inv_n_phi = 1.0 / phi_samples as f32;
238        let inv_n_theta = 1.0 / theta_samples as f32;
239        // discrete weight: (Δθ * Δφ) = (π/2 / N_θ) * (2π / N_φ) = π² / (N_θ N_φ)
240        let weight = core::f32::consts::PI * core::f32::consts::PI * inv_n_phi * inv_n_theta;
241        Self {
242            source,
243            source_face_size,
244            output_face_size,
245            kernel: Kernel::Irradiance(IrradianceKernel {
246                phi_samples,
247                theta_samples,
248                inv_n_phi,
249                inv_n_theta,
250                weight,
251            }),
252        }
253    }
254
255    /// GGX convolution at `roughness`, importance sampled with `samples`
256    /// Hammersley points per texel and capped at `clamp` (see the firefly
257    /// suppression in the module prose).
258    pub fn ggx(
259        source: &'a [Vec<f32>; 6],
260        source_face_size: u32,
261        output_face_size: u32,
262        roughness: f32,
263        samples: u32,
264        clamp: f32,
265    ) -> Self {
266        let a = roughness * roughness;
267        Self {
268            source,
269            source_face_size,
270            output_face_size,
271            kernel: Kernel::Ggx(GgxKernel {
272                samples,
273                a2m1: a * a - 1.0,
274                clamp,
275            }),
276        }
277    }
278
279    /// Cube face edge of this bake's output, in pixels.
280    pub fn output_face_size(&self) -> u32 {
281        self.output_face_size
282    }
283
284    /// Six zeroed RGBA32F faces sized for this bake's output.
285    pub fn output_faces(&self) -> [Vec<f32>; 6] {
286        let f = self.output_face_size as usize;
287        core::array::from_fn(|_| vec![0.0; f * f * 4])
288    }
289
290    /// Convolve the source into six output faces, handing the independent rows
291    /// to `scheduler`. A caller with a thread pool gets the fan-out for free;
292    /// one without uses [`Serial`](super::schedule::Serial).
293    pub fn bake<S: super::schedule::RowScheduler>(&self, scheduler: &S) -> [Vec<f32>; 6] {
294        let mut faces = self.output_faces();
295        let mut rows = face_rows(&mut faces, self.output_face_size());
296        scheduler.run(&mut rows, &|row| self.compute_row(row));
297        faces
298    }
299
300    /// Evaluate one output row. Reads only the source, writes only `row`.
301    pub fn compute_row(&self, row: &mut FaceRow<'_>) {
302        match &self.kernel {
303            Kernel::Irradiance(k) => self.irradiance_row(row, k),
304            Kernel::Ggx(k) => self.ggx_row(row, k),
305        }
306    }
307
308    /// Evaluate every row into a fresh output cube.
309    pub fn compute(&self) -> [Vec<f32>; 6] {
310        let mut faces = self.output_faces();
311        for row in &mut face_rows(&mut faces, self.output_face_size) {
312            self.compute_row(row);
313        }
314        faces
315    }
316
317    fn irradiance_row(&self, row: &mut FaceRow<'_>, k: &IrradianceKernel) {
318        for x in 0..self.output_face_size {
319            let n = cubemap::texel_dir(row.face, x, row.y, self.output_face_size);
320            let (tan, bit) = make_tbn(n);
321            let mut sum = [0.0f32; 3];
322            for phi_i in 0..k.phi_samples {
323                let phi = 2.0 * core::f32::consts::PI * (phi_i as f32 + 0.5) * k.inv_n_phi;
324                let sin_phi = sin(phi);
325                let cos_phi = cos(phi);
326                for theta_i in 0..k.theta_samples {
327                    let theta =
328                        0.5 * core::f32::consts::PI * (theta_i as f32 + 0.5) * k.inv_n_theta;
329                    let sin_theta = sin(theta);
330                    let cos_theta = cos(theta);
331                    let l_local = [sin_theta * cos_phi, sin_theta * sin_phi, cos_theta];
332                    let dir = [
333                        tan[0] * l_local[0] + bit[0] * l_local[1] + n[0] * l_local[2],
334                        tan[1] * l_local[0] + bit[1] * l_local[1] + n[1] * l_local[2],
335                        tan[2] * l_local[0] + bit[2] * l_local[1] + n[2] * l_local[2],
336                    ];
337                    let env = sample_cube(self.source, self.source_face_size, normalize3(dir));
338                    // cos(θ) for the Lambert cosine, sin(θ) for the spherical
339                    // area element. Both already in [0, 1] for the hemisphere.
340                    let w = cos_theta * sin_theta;
341                    sum[0] += env[0] * w;
342                    sum[1] += env[1] * w;
343                    sum[2] += env[2] * w;
344                }
345            }
346            let off = x as usize * 4;
347            row.texels[off] = sum[0] * k.weight;
348            row.texels[off + 1] = sum[1] * k.weight;
349            row.texels[off + 2] = sum[2] * k.weight;
350            row.texels[off + 3] = 1.0;
351        }
352    }
353
354    fn ggx_row(&self, row: &mut FaceRow<'_>, k: &GgxKernel) {
355        for x in 0..self.output_face_size {
356            let n = cubemap::texel_dir(row.face, x, row.y, self.output_face_size);
357            // The tangent basis depends only on N, so it is built once per
358            // texel rather than once per sample.
359            let basis = make_tbn(n);
360            // Split-sum approximation: V = R = N. The light direction is
361            // then L = reflect(-V, H) = 2 (N·H) H - N.
362            let mut accum = [0.0f32; 3];
363            let mut total_weight = 0.0f32;
364            for i in 0..k.samples {
365                let xi = hammersley(i, k.samples);
366                let h = importance_sample_ggx(xi, n, basis, k.a2m1);
367                let ndh = dot3(n, h);
368                if ndh <= 0.0 {
369                    continue;
370                }
371                let l = normalize3([
372                    2.0 * ndh * h[0] - n[0],
373                    2.0 * ndh * h[1] - n[1],
374                    2.0 * ndh * h[2] - n[2],
375                ]);
376                let ndl = dot3(n, l).max(0.0);
377                if ndl > 0.0 {
378                    let env =
379                        clamp_radiance(sample_cube(self.source, self.source_face_size, l), k.clamp);
380                    accum[0] += env[0] * ndl;
381                    accum[1] += env[1] * ndl;
382                    accum[2] += env[2] * ndl;
383                    total_weight += ndl;
384                }
385            }
386            let off = x as usize * 4;
387            if total_weight > 0.0 {
388                let inv = 1.0 / total_weight;
389                row.texels[off] = accum[0] * inv;
390                row.texels[off + 1] = accum[1] * inv;
391                row.texels[off + 2] = accum[2] * inv;
392            } else {
393                let n_sample = sample_cube(self.source, self.source_face_size, n);
394                row.texels[off] = n_sample[0];
395                row.texels[off + 1] = n_sample[1];
396                row.texels[off + 2] = n_sample[2];
397            }
398            row.texels[off + 3] = 1.0;
399        }
400    }
401}
402
403// Prefilter chain
404
405/// Mip 0 of a prefilter chain: the source copied through at alpha 1.
406///
407/// `clamp_mip0` decides whether the firefly cap the rough mips always apply
408/// also caps this mirror mip. An imported environment map leaves it OFF -- mip 0
409/// is drawn directly as the on-screen skybox, which must keep its true HDR
410/// sun/sky. A reflection probe turns it ON -- the probe is never a skybox (it is
411/// sampled only by the specular term, as a low-res fallback when SSR/RT miss on
412/// a near-mirror surface), so a lone blown highlight in the capture would
413/// otherwise alias into a bright square there; capping it (at the same `clamp`)
414/// suppresses that without touching any sky.
415pub fn prefilter_mip0(
416    source: &[Vec<f32>; 6],
417    source_face_size: u32,
418    clamp: f32,
419    clamp_mip0: bool,
420) -> [Vec<f32>; 6] {
421    let f = source_face_size as usize;
422    let mut mip0: [Vec<f32>; 6] = core::array::from_fn(|_| vec![0.0; f * f * 4]);
423    for face in 0..6 {
424        for i in 0..f * f {
425            let off = i * 4;
426            let mut rgb = [
427                source[face][off],
428                source[face][off + 1],
429                source[face][off + 2],
430            ];
431            if clamp_mip0 {
432                rgb = clamp_radiance(rgb, clamp);
433            }
434            mip0[face][off] = rgb[0];
435            mip0[face][off + 1] = rgb[1];
436            mip0[face][off + 2] = rgb[2];
437            mip0[face][off + 3] = 1.0;
438        }
439    }
440    mip0
441}
442
443/// GGX roughness for mip `mip` of a `mip_count` chain: 0 at mip 0, 1 at the
444/// last mip.
445pub fn prefilter_roughness(mip: u32, mip_count: u32) -> f32 {
446    mip as f32 / (mip_count - 1) as f32
447}
448
449// Whole-bake entry points
450
451/// Compute a low-resolution irradiance cubemap. The serial form of
452/// [`CubeBake::irradiance`].
453pub fn compute_irradiance(
454    source: &[Vec<f32>; 6],
455    source_face_size: u32,
456    output_face_size: u32,
457    phi_samples: u32,
458    theta_samples: u32,
459) -> [Vec<f32>; 6] {
460    CubeBake::irradiance(
461        source,
462        source_face_size,
463        output_face_size,
464        phi_samples,
465        theta_samples,
466    )
467    .compute()
468}
469
470/// Build a prefiltered radiance cube mip chain. Mip 0 is the unmodified
471/// source (roughness=0 → Dirac lobe). Mip N is the GGX convolution at
472/// roughness = N / (mip_count - 1). The serial form of [`prefilter_mip0`]
473/// followed by one [`CubeBake::ggx`] per remaining mip.
474pub fn compute_prefilter(
475    source: &[Vec<f32>; 6],
476    source_face_size: u32,
477    mip_count: u32,
478    samples_per_texel: u32,
479    clamp: f32,
480    clamp_mip0: bool,
481) -> Vec<[Vec<f32>; 6]> {
482    let mut mips: Vec<[Vec<f32>; 6]> = Vec::with_capacity(mip_count as usize);
483    mips.push(prefilter_mip0(source, source_face_size, clamp, clamp_mip0));
484    for mip in 1..mip_count {
485        mips.push(
486            CubeBake::ggx(
487                source,
488                source_face_size,
489                source_face_size >> mip,
490                prefilter_roughness(mip, mip_count),
491                samples_per_texel,
492                clamp,
493            )
494            .compute(),
495        );
496    }
497    mips
498}
499
500#[cfg(test)]
501mod tests {
502    use super::super::schedule::{RowScheduler, Serial};
503    use super::*;
504
505    // A scheduler that walks the rows backwards. The bake's whole claim is that
506    // the rows are independent, so the order it runs them in must not show up in
507    // the output.
508    struct Reversed;
509
510    impl RowScheduler for Reversed {
511        fn run<T: Send>(&self, items: &mut [T], compute: &(dyn Fn(&mut T) + Send + Sync)) {
512            items.iter_mut().rev().for_each(compute);
513        }
514    }
515
516    #[test]
517    fn a_bake_does_not_depend_on_the_row_order() {
518        let source = solid_cube(8, [0.4, 0.6, 0.9]);
519        let bake = CubeBake::ggx(&source, 8, 8, 0.5, 16, 0.0);
520        assert_eq!(bake.bake(&Serial), bake.bake(&Reversed));
521    }
522
523    #[test]
524    fn a_bake_fills_every_output_face() {
525        let source = solid_cube(8, [1.0, 1.0, 1.0]);
526        let bake = CubeBake::irradiance(&source, 8, 4, 8, 8);
527        let faces = bake.bake(&Serial);
528        for face in &faces {
529            assert_eq!(face.len(), 4 * 4 * 4);
530            assert!(face.chunks_exact(4).all(|px| px[0] > 0.0));
531        }
532    }
533
534    fn solid_cube(face_size: u32, color: [f32; 3]) -> [Vec<f32>; 6] {
535        let f = face_size as usize;
536        core::array::from_fn(|_| {
537            let mut face = Vec::with_capacity(f * f * 4);
538            for _ in 0..f * f {
539                face.extend_from_slice(&[color[0], color[1], color[2], 1.0]);
540            }
541            face
542        })
543    }
544
545    fn face_mean(face: &[f32]) -> [f32; 3] {
546        let n = face.len() / 4;
547        let mut m = [0.0f32; 3];
548        for px in face.chunks_exact(4) {
549            m[0] += px[0];
550            m[1] += px[1];
551            m[2] += px[2];
552        }
553        [m[0] / n as f32, m[1] / n as f32, m[2] / n as f32]
554    }
555
556    fn face_variance_red(face: &[f32]) -> f32 {
557        let n = face.len() / 4;
558        let mean = face.chunks_exact(4).map(|p| p[0]).sum::<f32>() / n as f32;
559
560        face.chunks_exact(4)
561            .map(|p| (p[0] - mean).powi(2))
562            .sum::<f32>()
563            / n as f32
564    }
565
566    // A source cube with one blazing texel on +Z over a uniform background: the
567    // stand-in for a sun disk or a blown highlight in a probe capture.
568    fn firefly_cube(face: usize) -> [Vec<f32>; 6] {
569        let mut s: [Vec<f32>; 6] = core::array::from_fn(|_| vec![0.0; face * face * 4]);
570        for fd in s.iter_mut() {
571            for p in fd.chunks_exact_mut(4) {
572                p[0] = 1.0;
573                p[1] = 1.0;
574                p[2] = 1.0;
575                p[3] = 1.0;
576            }
577        }
578        let off = ((face / 2) * face + face / 2) * 4;
579        s[4][off] = 2000.0;
580        s[4][off + 1] = 2000.0;
581        s[4][off + 2] = 2000.0;
582        s
583    }
584
585    fn peak(fd: &[f32]) -> f32 {
586        fd.chunks_exact(4)
587            .map(|p| p[0].max(p[1]).max(p[2]))
588            .fold(0.0f32, f32::max)
589    }
590
591    // A source with structure in every direction, so a row that read the wrong
592    // face or the wrong y would produce different bytes rather than the same
593    // constant.
594    fn gradient_cube(face_size: u32) -> [Vec<f32>; 6] {
595        let f = face_size as usize;
596        core::array::from_fn(|face| {
597            let mut data = vec![0.0f32; f * f * 4];
598            for y in 0..f {
599                for x in 0..f {
600                    let off = (y * f + x) * 4;
601                    data[off] = face as f32 + x as f32 / f as f32;
602                    data[off + 1] = y as f32 / f as f32;
603                    data[off + 2] = (x + y) as f32 / (2 * f) as f32;
604                    data[off + 3] = 1.0;
605                }
606            }
607            data
608        })
609    }
610
611    // The contract the chunk API exists for: a caller that drives the rows
612    // itself, in an order nothing guarantees, gets the bytes `compute` would
613    // have produced. Reversed here because a row that leaked state into the
614    // next one would still pass in forward order.
615    fn assert_rows_match_whole_image(bake: &CubeBake<'_>) {
616        let whole = bake.compute();
617        let mut chunked = bake.output_faces();
618        let mut rows = face_rows(&mut chunked, bake.output_face_size());
619        assert_eq!(
620            rows.len(),
621            6 * bake.output_face_size() as usize,
622            "one row per face line"
623        );
624        for row in rows.iter_mut().rev() {
625            bake.compute_row(row);
626        }
627        assert_eq!(chunked, whole, "chunked bake diverged from the whole image");
628    }
629
630    #[test]
631    fn chunked_irradiance_matches_the_whole_image() {
632        let source = gradient_cube(8);
633        assert_rows_match_whole_image(&CubeBake::irradiance(&source, 8, 4, 16, 8));
634    }
635
636    #[test]
637    fn chunked_ggx_matches_the_whole_image() {
638        let source = gradient_cube(16);
639        for roughness in [0.25f32, 0.5, 1.0] {
640            assert_rows_match_whole_image(&CubeBake::ggx(&source, 16, 8, roughness, 32, 0.0));
641        }
642    }
643
644    // The firefly cap is part of the kernel, so it has to survive the split too.
645    #[test]
646    fn chunked_ggx_matches_the_whole_image_under_the_firefly_cap() {
647        let source = firefly_cube(16);
648        assert_rows_match_whole_image(&CubeBake::ggx(&source, 16, 8, 0.5, 64, 8.0));
649    }
650
651    #[test]
652    fn face_rows_cover_every_texel_exactly_once() {
653        let mut faces: [Vec<f32>; 6] = core::array::from_fn(|_| vec![0.0; 4 * 4 * 4]);
654        for row in &mut face_rows(&mut faces, 4) {
655            assert_eq!(row.texels.len(), 4 * 4, "a row is one line of RGBA texels");
656            for t in row.texels.iter_mut() {
657                *t += 1.0;
658            }
659        }
660        assert!(
661            faces.iter().flatten().all(|&t| t == 1.0),
662            "every texel written exactly once"
663        );
664    }
665
666    #[test]
667    fn hammersley_first_sample_is_zero() {
668        let s = hammersley(0, 1024);
669        assert!(s[0].abs() < 1e-6, "x was {}", s[0]);
670        assert!(s[1].abs() < 1e-6, "y was {}", s[1]);
671    }
672
673    #[test]
674    fn hammersley_last_sample_is_just_under_one() {
675        let s = hammersley(1023, 1024);
676        assert!(s[0] > 0.99 && s[0] < 1.0, "x was {}", s[0]);
677    }
678
679    #[test]
680    fn importance_sample_ggx_at_xi_zero_returns_n() {
681        let n = [0.0, 0.0, 1.0];
682        let a = 0.5f32 * 0.5;
683        let h = importance_sample_ggx([0.0, 0.0], n, make_tbn(n), a * a - 1.0);
684        // xi=(0,0) → cos_theta = 1 → H aligns with N.
685        assert!((h[0] - 0.0).abs() < 1e-5);
686        assert!((h[1] - 0.0).abs() < 1e-5);
687        assert!((h[2] - 1.0).abs() < 1e-5);
688    }
689
690    #[test]
691    fn irradiance_solid_color_is_pi_times_color() {
692        // Uniform environment L = (1, 0.5, 0.25). The hemispherical integral
693        // of L * cos(θ) over the upper hemisphere is π * L. The discrete
694        // (phi, theta) integration should converge to that.
695        let source = solid_cube(8, [1.0, 0.5, 0.25]);
696        let irr = compute_irradiance(&source, 8, 4, 64, 16);
697        let mean = face_mean(&irr[0]);
698        let expected = [
699            core::f32::consts::PI * 1.0,
700            core::f32::consts::PI * 0.5,
701            core::f32::consts::PI * 0.25,
702        ];
703        // Discrete integration loses a few percent; accept ±5%.
704        for c in 0..3 {
705            let delta = (mean[c] - expected[c]).abs() / expected[c];
706            assert!(
707                delta < 0.05,
708                "channel {} mean {} expected {}",
709                c,
710                mean[c],
711                expected[c]
712            );
713        }
714    }
715
716    #[test]
717    fn prefilter_mip_zero_matches_source_with_alpha_one() {
718        let source = solid_cube(16, [0.7, 0.3, 0.1]);
719        let mips = compute_prefilter(&source, 16, 3, 16, 0.0, false);
720        for face in &mips[0] {
721            for px in 0..16 * 16 {
722                let off = px * 4;
723                assert!((face[off] - 0.7).abs() < 1e-6);
724                assert!((face[off + 1] - 0.3).abs() < 1e-6);
725                assert!((face[off + 2] - 0.1).abs() < 1e-6);
726                assert!((face[off + 3] - 1.0).abs() < 1e-6);
727            }
728        }
729    }
730
731    #[test]
732    fn prefilter_solid_color_stays_solid_at_high_roughness() {
733        let source = solid_cube(16, [0.5, 0.5, 0.5]);
734        let mips = compute_prefilter(&source, 16, 4, 32, 0.0, false);
735        // Last mip should still be ~0.5 grey since input is uniform.
736        let mean = face_mean(&mips[3][0]);
737        for (c, m) in mean.iter().enumerate() {
738            assert!((m - 0.5).abs() < 0.02, "channel {} mean {}", c, m);
739        }
740    }
741
742    #[test]
743    fn prefilter_roughness_spans_zero_to_one() {
744        assert_eq!(prefilter_roughness(0, 5), 0.0);
745        assert_eq!(prefilter_roughness(4, 5), 1.0);
746        assert_eq!(prefilter_roughness(2, 5), 0.5);
747    }
748
749    #[test]
750    fn prefilter_blurs_a_red_seam() {
751        // Place a bright red column on +Z face only; prefilter at roughness=1
752        // should spread it across the face so the variance drops vs the input.
753        let face = 16usize;
754        let mut source: [Vec<f32>; 6] = core::array::from_fn(|_| vec![0.0; face * face * 4]);
755        for face_data in source.iter_mut() {
756            for p in face_data.chunks_exact_mut(4) {
757                p[3] = 1.0;
758            }
759        }
760        // +Z face (index 4): paint x=8 column bright red.
761        for y in 0..face {
762            let off = (y * face + 8) * 4;
763            source[4][off] = 20.0;
764        }
765        let mips = compute_prefilter(&source, face as u32, 3, 256, 0.0, false);
766        // Compare variance of +Z face at mip 0 vs mip 2.
767        let v0 = face_variance_red(&mips[0][4]);
768        let v2 = face_variance_red(&mips[2][4]);
769        assert!(
770            v2 < v0 * 0.5,
771            "prefilter did not blur: mip 0 var={}, mip 2 var={}",
772            v0,
773            v2
774        );
775    }
776
777    #[test]
778    fn clamp_radiance_caps_luminance_and_keeps_hue() {
779        // Below the cap: untouched.
780        let dim = [1.0, 0.5, 0.25];
781        assert_eq!(clamp_radiance(dim, 10.0), dim);
782        // Disabled (clamp <= 0): untouched even when very bright.
783        let hot = [100.0, 50.0, 25.0];
784        assert_eq!(clamp_radiance(hot, 0.0), hot);
785        // Above the cap: luminance is pulled to the cap, hue preserved.
786        let capped = clamp_radiance(hot, 10.0);
787        let lum = 0.2126 * capped[0] + 0.7152 * capped[1] + 0.0722 * capped[2];
788        assert!((lum - 10.0).abs() < 1e-3, "luminance {} != cap 10", lum);
789        assert!((capped[0] / capped[1] - hot[0] / hot[1]).abs() < 1e-4);
790        assert!((capped[1] / capped[2] - hot[1] / hot[2]).abs() < 1e-4);
791    }
792
793    #[test]
794    fn prefilter_clamp_suppresses_a_firefly() {
795        // Unclamped the blazing texel survives the GGX convolution as a hot spot
796        // that smears into bright squares; the cap spreads its energy so the
797        // brightest reflection texel is far dimmer, while the background (below
798        // the cap) is preserved.
799        let face = 16usize;
800        let unclamped = compute_prefilter(&firefly_cube(face), face as u32, 3, 256, 0.0, false);
801        let clamped = compute_prefilter(&firefly_cube(face), face as u32, 3, 256, 8.0, false);
802        let p_unclamped = peak(&unclamped[1][4]);
803        let p_clamped = peak(&clamped[1][4]);
804        assert!(
805            p_clamped < p_unclamped * 0.5,
806            "clamp did not suppress the firefly: unclamped {}, clamped {}",
807            p_unclamped,
808            p_clamped
809        );
810        assert!(
811            p_clamped >= 0.9,
812            "clamp crushed the background: clamped peak {}",
813            p_clamped
814        );
815    }
816
817    #[test]
818    fn prefilter_mip0_clamp_caps_a_mirror_firefly_only_when_requested() {
819        // Mip 0 is the mirror (roughness 0) reflection a near-mirror surface
820        // samples on an SSR/RT miss.
821        let face = 16usize;
822        // clamp_mip0 = false (an imported env map / skybox): the blazing texel survives
823        // mip 0 untouched, so the on-screen sky would keep its true HDR sun.
824        let unclamped_mip0 = compute_prefilter(&firefly_cube(face), face as u32, 2, 16, 8.0, false);
825        assert!(
826            peak(&unclamped_mip0[0][4]) > 1000.0,
827            "mip 0 should be unclamped when clamp_mip0 is false: peak {}",
828            peak(&unclamped_mip0[0][4])
829        );
830        // clamp_mip0 = true (a reflection probe): the same firefly is capped at mip 0,
831        // while the uniform background (below the cap) is preserved.
832        let clamped_mip0 = compute_prefilter(&firefly_cube(face), face as u32, 2, 16, 8.0, true);
833        let p = peak(&clamped_mip0[0][4]);
834        assert!(p <= 8.0 + 1e-3, "mip 0 firefly not capped: peak {p}");
835        assert!(p >= 0.9, "mip 0 clamp crushed the background: peak {p}");
836    }
837}