Skip to main content

pfx_trace/
lights.rs

1use bytemuck::{Pod, Zeroable};
2
3use crate::math::{add, cross, dot, length, scale, sub};
4use crate::shapes::Shape;
5
6pub const MAX_EMITTERS: usize = 1 << 24;
7pub const MAX_LIGHTS: usize = 1 << 16;
8const SHAPE_MARGIN: f32 = 0.003;
9
10#[derive(Clone, Copy, Debug, PartialEq)]
11pub enum LightShape {
12    Point,
13    Spot {
14        direction: [f32; 3],
15        inner_deg: f32,
16        outer_deg: f32,
17    },
18    Rect {
19        half_u: [f32; 3],
20        half_v: [f32; 3],
21        two_sided: bool,
22    },
23    Disc {
24        normal: [f32; 3],
25        radius: f32,
26        two_sided: bool,
27    },
28}
29
30#[derive(Clone, Copy, Debug, PartialEq)]
31pub struct LocalLight {
32    pub position: [f32; 3],
33    pub colour: [f32; 3],
34    pub intensity: f32,
35    pub radius: f32,
36    pub range: f32,
37    pub shadow: bool,
38    pub shape: LightShape,
39}
40
41impl LocalLight {
42    pub fn point(position: [f32; 3], colour: [f32; 3], intensity: f32, range: f32) -> Self {
43        Self {
44            position,
45            colour,
46            intensity,
47            radius: 0.0,
48            range,
49            shadow: true,
50            shape: LightShape::Point,
51        }
52    }
53
54    pub fn rect(
55        centre: [f32; 3],
56        half_u: [f32; 3],
57        half_v: [f32; 3],
58        colour: [f32; 3],
59        intensity: f32,
60    ) -> Self {
61        Self {
62            position: centre,
63            colour,
64            intensity,
65            radius: 0.0,
66            range: f32::MAX,
67            shadow: true,
68            shape: LightShape::Rect {
69                half_u,
70                half_v,
71                two_sided: false,
72            },
73        }
74    }
75
76    pub fn disc(
77        centre: [f32; 3],
78        normal: [f32; 3],
79        radius: f32,
80        colour: [f32; 3],
81        intensity: f32,
82    ) -> Self {
83        Self {
84            position: centre,
85            colour,
86            intensity,
87            radius: 0.0,
88            range: f32::MAX,
89            shadow: true,
90            shape: LightShape::Disc {
91                normal,
92                radius,
93                two_sided: false,
94            },
95        }
96    }
97
98    pub fn lit(&self) -> bool {
99        self.intensity > 0.0 && self.range > 0.0 && self.colour.iter().any(|&value| value > 0.0)
100    }
101
102    pub fn valid(&self) -> bool {
103        let scalars = self
104            .position
105            .iter()
106            .chain(&self.colour)
107            .chain([&self.intensity, &self.radius, &self.range])
108            .all(|value| value.is_finite())
109            && self.colour.iter().all(|&value| value >= 0.0)
110            && self.intensity >= 0.0
111            && self.radius >= 0.0
112            && self.range >= 0.0;
113        scalars
114            && match self.shape {
115                LightShape::Point => true,
116                LightShape::Spot {
117                    direction,
118                    inner_deg,
119                    outer_deg,
120                } => {
121                    direction.iter().all(|value| value.is_finite())
122                        && length(direction) > 1e-6
123                        && inner_deg.is_finite()
124                        && outer_deg.is_finite()
125                        && 0.0 <= inner_deg
126                        && inner_deg <= outer_deg
127                        && outer_deg > 0.0
128                        && outer_deg <= 90.0
129                }
130                LightShape::Rect { half_u, half_v, .. } => {
131                    let lu = length(half_u);
132                    let lv = length(half_v);
133                    half_u.iter().chain(&half_v).all(|value| value.is_finite())
134                        && lu > 1e-6
135                        && lv > 1e-6
136                        && dot(half_u, half_v).abs() <= 1e-4 * lu * lv
137                }
138                LightShape::Disc { normal, radius, .. } => {
139                    normal.iter().all(|value| value.is_finite())
140                        && length(normal) > 1e-6
141                        && radius.is_finite()
142                        && radius > 1e-6
143                }
144            }
145    }
146
147    pub fn area(&self) -> f32 {
148        match self.shape {
149            LightShape::Rect { half_u, half_v, .. } => 4.0 * length(cross(half_u, half_v)),
150            LightShape::Disc { radius, .. } => std::f32::consts::PI * radius * radius,
151            _ => 0.0,
152        }
153    }
154
155    pub fn falloff(&self, distance: f32) -> f32 {
156        if distance >= self.range || distance <= 1e-6 {
157            return 0.0;
158        }
159        let ratio = distance / self.range;
160        let window = (1.0 - ratio.powi(4)).clamp(0.0, 1.0);
161        let reach = distance.max(self.radius);
162        self.intensity * window * window / (reach * reach)
163    }
164
165    pub fn cone(&self, toward_surface: [f32; 3]) -> f32 {
166        match self.shape {
167            LightShape::Spot {
168                direction,
169                inner_deg,
170                outer_deg,
171            } => {
172                let (scale, offset) = spot_terms(inner_deg, outer_deg);
173                let axis = crate::math::normalize(direction);
174                let cd = dot(crate::math::normalize(toward_surface), axis);
175                let t = (cd * scale + offset).clamp(0.0, 1.0);
176                t * t
177            }
178            _ => 1.0,
179        }
180    }
181
182    pub(crate) fn pack(&self) -> GpuLight {
183        let colour = self.colour.map(|channel| channel * self.intensity);
184        let shadow = f32::from(u8::from(self.shadow));
185        match self.shape {
186            LightShape::Point => GpuLight {
187                a: vec4(self.position, 0.0),
188                b: vec4(colour, self.range),
189                c: [0.0, 0.0, 0.0, self.radius],
190                d: [0.0, 0.0, 0.0, shadow],
191            },
192            LightShape::Spot {
193                direction,
194                inner_deg,
195                outer_deg,
196            } => {
197                let (scale, offset) = spot_terms(inner_deg, outer_deg);
198                GpuLight {
199                    a: vec4(self.position, 1.0),
200                    b: vec4(colour, self.range),
201                    c: vec4(crate::math::normalize(direction), self.radius),
202                    d: [scale, offset, 0.0, shadow],
203                }
204            }
205            LightShape::Rect {
206                half_u,
207                half_v,
208                two_sided,
209            } => {
210                let area = self.area();
211                GpuLight {
212                    a: vec4(self.position, 2.0),
213                    b: vec4(colour.map(|channel| channel / area), f32::MAX),
214                    c: vec4(half_u, f32::from(u8::from(two_sided))),
215                    d: vec4(half_v, shadow),
216                }
217            }
218            LightShape::Disc {
219                normal,
220                radius,
221                two_sided,
222            } => {
223                let area = self.area();
224                GpuLight {
225                    a: vec4(self.position, 3.0),
226                    b: vec4(colour.map(|channel| channel / area), f32::MAX),
227                    c: vec4(
228                        crate::math::normalize(normal),
229                        f32::from(u8::from(two_sided)),
230                    ),
231                    d: [radius, 0.0, 0.0, shadow],
232                }
233            }
234        }
235    }
236}
237
238fn spot_terms(inner_deg: f32, outer_deg: f32) -> (f32, f32) {
239    let inner = inner_deg.to_radians().cos();
240    let outer = outer_deg.to_radians().cos();
241    let scale = 1.0 / (inner - outer).max(0.001);
242    (scale, -outer * scale)
243}
244
245#[repr(C)]
246#[derive(Clone, Copy, Debug, Pod, Zeroable)]
247pub(crate) struct GpuLight {
248    pub a: [f32; 4],
249    pub b: [f32; 4],
250    pub c: [f32; 4],
251    pub d: [f32; 4],
252}
253
254#[repr(C)]
255#[derive(Clone, Copy, Debug, Pod, Zeroable, PartialEq)]
256pub(crate) struct GpuEmitter {
257    pub a: [f32; 4],
258    pub b: [f32; 4],
259    pub c: [f32; 4],
260    pub emission: [f32; 4],
261}
262
263#[derive(Clone, Copy, Debug, PartialEq)]
264pub(crate) enum Source {
265    Triangle {
266        vertices: [[f32; 3]; 3],
267        slot: u32,
268    },
269    Shape {
270        center: [f32; 3],
271        radius: f32,
272        index: u32,
273        area: f32,
274    },
275}
276
277#[derive(Clone, Copy, Debug, PartialEq)]
278pub(crate) struct Candidate {
279    pub source: Source,
280    pub emission: [f32; 3],
281}
282
283pub(crate) struct EmitterList {
284    pub entries: Vec<GpuEmitter>,
285    pub cdf: Vec<f32>,
286}
287
288pub(crate) fn luminance(v: [f32; 3]) -> f32 {
289    0.2126 * v[0] + 0.7152 * v[1] + 0.0722 * v[2]
290}
291
292pub(crate) fn emits(emission: [f32; 3]) -> bool {
293    emission
294        .iter()
295        .all(|value| value.is_finite() && *value >= 0.0)
296        && luminance(emission) > 0.0
297}
298
299pub(crate) fn triangle_area(vertices: [[f32; 3]; 3]) -> f32 {
300    0.5 * length(cross(
301        sub(vertices[1], vertices[0]),
302        sub(vertices[2], vertices[0]),
303    ))
304}
305
306impl Candidate {
307    fn area(&self) -> f32 {
308        match self.source {
309            Source::Triangle { vertices, .. } => triangle_area(vertices),
310            Source::Shape { area, .. } => area,
311        }
312    }
313}
314
315pub(crate) fn shape_area(shape: Shape, shapes: &[Shape]) -> f32 {
316    use std::f32::consts::PI;
317    match shape {
318        Shape::RoundedBox { half, radius, .. } => {
319            let core = half.map(|value| (value - radius).max(0.0));
320            8.0 * (core[0] * core[1] + core[1] * core[2] + core[0] * core[2])
321                + 2.0 * PI * radius * 2.0 * (core[0] + core[1] + core[2])
322                + 4.0 * PI * radius * radius
323        }
324        Shape::RoundCone {
325            a,
326            b,
327            radius_a,
328            radius_b,
329            ..
330        } => {
331            let axis = length(sub(b, a));
332            let slant = (axis * axis + (radius_a - radius_b).powi(2)).sqrt();
333            PI * (radius_a + radius_b) * slant
334                + 2.0 * PI * (radius_a * radius_a + radius_b * radius_b)
335        }
336        Shape::Ellipsoid { radii, .. } => {
337            let p = 1.6075;
338            let [a, b, c] = radii.map(|value| value.powf(p));
339            4.0 * PI * ((a * b + a * c + b * c) / 3.0).powf(1.0 / p)
340        }
341        Shape::Capsule { a, b, radius, .. } => {
342            2.0 * PI * radius * length(sub(b, a)) + 4.0 * PI * radius * radius
343        }
344        Shape::SmoothUnion { left, right, .. } => {
345            shape_area(shapes[left], shapes) + shape_area(shapes[right], shapes)
346        }
347    }
348}
349
350pub(crate) fn bounding_sphere(shape: Shape, shapes: &[Shape]) -> ([f32; 3], f32) {
351    let (center, radius) = match shape {
352        Shape::RoundedBox { center, half, .. } => (center, length(half)),
353        Shape::RoundCone {
354            a,
355            b,
356            radius_a,
357            radius_b,
358            ..
359        } => (
360            scale(add(a, b), 0.5),
361            0.5 * length(sub(b, a)) + radius_a.max(radius_b),
362        ),
363        Shape::Ellipsoid { center, radii, .. } => (center, radii[0].max(radii[1]).max(radii[2])),
364        Shape::Capsule { a, b, radius, .. } => {
365            (scale(add(a, b), 0.5), 0.5 * length(sub(b, a)) + radius)
366        }
367        Shape::SmoothUnion {
368            left,
369            right,
370            radius,
371            ..
372        } => {
373            let (c0, r0) = bounding_sphere(shapes[left], shapes);
374            let (c1, r1) = bounding_sphere(shapes[right], shapes);
375            let gap = length(sub(c1, c0));
376            if gap + r1 <= r0 {
377                (c0, r0 + radius)
378            } else if gap + r0 <= r1 {
379                (c1, r1 + radius)
380            } else {
381                let r = 0.5 * (gap + r0 + r1);
382                let center = add(c0, scale(sub(c1, c0), (r - r0) / gap.max(1e-12)));
383                (center, r + radius)
384            }
385        }
386    };
387    (center, radius + SHAPE_MARGIN)
388}
389
390pub(crate) fn build_emitters(candidates: &[Candidate], lights: usize) -> EmitterList {
391    let powers: Vec<f64> = candidates
392        .iter()
393        .map(|candidate| f64::from(candidate.area()) * f64::from(luminance(candidate.emission)))
394        .collect();
395    let total: f64 = powers.iter().sum();
396    let mut entries = vec![GpuEmitter {
397        a: [candidates.len() as f32, lights as f32, total as f32, 0.0],
398        b: [0.0; 4],
399        c: [0.0; 4],
400        emission: [0.0; 4],
401    }];
402    let mut cdf = Vec::with_capacity(candidates.len());
403    let mut running = 0.0_f64;
404    for (candidate, power) in candidates.iter().zip(&powers) {
405        running += power;
406        cdf.push((running / total) as f32);
407        let probability = (power / total) as f32;
408        entries.push(match candidate.source {
409            Source::Triangle { vertices, slot } => GpuEmitter {
410                a: vec4(vertices[0], 0.0),
411                b: vec4(vertices[1], slot as f32),
412                c: vec4(vertices[2], triangle_area(vertices)),
413                emission: vec4(candidate.emission, probability),
414            },
415            Source::Shape {
416                center,
417                radius,
418                index,
419                ..
420            } => GpuEmitter {
421                a: vec4(center, 1.0),
422                b: vec4([0.0; 3], index as f32),
423                c: [radius, 0.0, 0.0, 0.0],
424                emission: vec4(candidate.emission, probability),
425            },
426        });
427    }
428    if let Some(last) = cdf.last_mut() {
429        *last = 1.0;
430    }
431    EmitterList { entries, cdf }
432}
433
434fn vec4(v: [f32; 3], w: f32) -> [f32; 4] {
435    [v[0], v[1], v[2], w]
436}
437
438#[cfg(test)]
439mod tests {
440    use super::*;
441
442    fn quad(size: f32, height: f32) -> [[[f32; 3]; 3]; 2] {
443        [
444            [
445                [0.0, height, 0.0],
446                [size, height, 0.0],
447                [size, height, size],
448            ],
449            [
450                [0.0, height, 0.0],
451                [size, height, size],
452                [0.0, height, size],
453            ],
454        ]
455    }
456
457    #[test]
458    fn emitters_are_weighted_by_power() {
459        let mut candidates = Vec::new();
460        for (slot, vertices) in quad(1.0, 2.0).into_iter().enumerate() {
461            candidates.push(Candidate {
462                source: Source::Triangle {
463                    vertices,
464                    slot: slot as u32,
465                },
466                emission: [1.0; 3],
467            });
468        }
469        for (slot, vertices) in quad(2.0, 2.0).into_iter().enumerate() {
470            candidates.push(Candidate {
471                source: Source::Triangle {
472                    vertices,
473                    slot: 2 + slot as u32,
474                },
475                emission: [0.5; 3],
476            });
477        }
478        let list = build_emitters(&candidates, 1);
479        assert_eq!(list.entries.len(), 5);
480        assert_eq!(list.entries[0].a[0], 4.0);
481        assert_eq!(list.entries[0].a[1], 1.0);
482        assert!((list.entries[0].a[2] - 3.0).abs() < 1e-5);
483        let probabilities: Vec<f32> = list.entries[1..].iter().map(|e| e.emission[3]).collect();
484        let expected = [1.0 / 6.0, 1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0];
485        for (got, want) in probabilities.iter().zip(expected) {
486            assert!((got - want).abs() < 1e-6, "{got} against {want}");
487        }
488        let mut running = 0.0;
489        for (value, want) in list.cdf.iter().zip(expected) {
490            running += want;
491            assert!((value - running).abs() < 1e-6);
492        }
493        assert_eq!(list.cdf.len(), 4);
494        assert_eq!(list.cdf[3], 1.0);
495        assert_eq!(list.entries[3].b[3], 2.0);
496        assert!((list.entries[3].c[3] - 2.0).abs() < 1e-6);
497    }
498
499    #[test]
500    fn empty_list_keeps_its_header() {
501        let list = build_emitters(&[], 2);
502        assert_eq!(list.entries.len(), 1);
503        assert_eq!(list.entries[0].a[..2], [0.0, 2.0]);
504        assert!(list.cdf.is_empty());
505    }
506
507    #[test]
508    fn only_positive_emission_emits() {
509        assert!(emits([0.0, 0.0, 1.0]));
510        assert!(!emits([0.0; 3]));
511        assert!(!emits([1.0, -0.5, 1.0]));
512        assert!(!emits([f32::NAN, 1.0, 1.0]));
513    }
514
515    #[test]
516    fn bounding_spheres_hold_their_shapes() {
517        let shapes = [
518            Shape::RoundedBox {
519                center: [1.0, 2.0, 3.0],
520                half: [0.5, 0.2, 0.1],
521                radius: 0.05,
522                material: 0,
523            },
524            Shape::Capsule {
525                a: [0.0, 0.0, 0.0],
526                b: [0.0, 2.0, 0.0],
527                radius: 0.25,
528                material: 0,
529            },
530            Shape::RoundCone {
531                a: [0.0, 0.0, 0.0],
532                b: [1.0, 0.0, 0.0],
533                radius_a: 0.4,
534                radius_b: 0.1,
535                material: 0,
536            },
537            Shape::Ellipsoid {
538                center: [0.5, 0.0, 0.0],
539                radii: [0.3, 0.6, 0.2],
540                material: 0,
541            },
542            Shape::SmoothUnion {
543                left: 1,
544                right: 3,
545                radius: 0.1,
546                material: 0,
547            },
548        ];
549        for shape in shapes {
550            let (center, radius) = bounding_sphere(shape, &shapes);
551            for i in 0..4096 {
552                let z = 1.0 - 2.0 * (i as f32 + 0.5) / 4096.0;
553                let phi = i as f32 * 2.399_963;
554                let r = (1.0 - z * z).sqrt();
555                let p = add(center, scale([r * phi.cos(), z, r * phi.sin()], radius));
556                assert!(shape.distance(p, &shapes) > 0.0, "{shape:?} at {p:?}");
557            }
558        }
559        let sphere = Shape::Ellipsoid {
560            center: [0.0; 3],
561            radii: [0.5; 3],
562            material: 0,
563        };
564        let area = shape_area(sphere, &[]);
565        assert!((area - std::f32::consts::PI).abs() < 1e-4);
566    }
567
568    #[test]
569    fn lights_validate_and_pack() {
570        let point = LocalLight::point([0.0, 2.0, 0.0], [1.0, 0.5, 0.25], 4.0, 10.0);
571        assert!(point.valid() && point.lit());
572        let packed = point.pack();
573        assert_eq!(packed.b, [4.0, 2.0, 1.0, 10.0]);
574        assert_eq!(packed.d[3], 1.0);
575        assert!((point.falloff(2.0) - 4.0 * (1.0 - 0.2_f32.powi(4)).powi(2) / 4.0).abs() < 1e-6);
576        assert_eq!(point.falloff(10.0), 0.0);
577        let rect = LocalLight::rect(
578            [0.0, 2.0, 0.0],
579            [0.5, 0.0, 0.0],
580            [0.0, 0.0, 0.25],
581            [1.0; 3],
582            2.0,
583        );
584        assert!(rect.valid());
585        assert!((rect.area() - 0.5).abs() < 1e-6);
586        assert_eq!(rect.pack().b[..3], [4.0; 3]);
587        let skew = LocalLight {
588            shape: LightShape::Rect {
589                half_u: [0.5, 0.0, 0.0],
590                half_v: [0.5, 0.0, 0.25],
591                two_sided: false,
592            },
593            ..rect
594        };
595        assert!(!skew.valid());
596        let spot = LocalLight {
597            shape: LightShape::Spot {
598                direction: [0.0, -1.0, 0.0],
599                inner_deg: 20.0,
600                outer_deg: 30.0,
601            },
602            ..point
603        };
604        assert!(spot.valid());
605        assert_eq!(spot.cone([0.0, -1.0, 0.0]), 1.0);
606        assert_eq!(spot.cone([1.0, -1.0, 0.0]), 0.0);
607        let edge = spot.cone([25.0_f32.to_radians().tan(), -1.0, 0.0]);
608        assert!(edge > 0.0 && edge < 1.0);
609        let wide = LocalLight {
610            shape: LightShape::Spot {
611                direction: [0.0, -1.0, 0.0],
612                inner_deg: 40.0,
613                outer_deg: 30.0,
614            },
615            ..point
616        };
617        assert!(!wide.valid());
618    }
619}