Skip to main content

pfx_trace/
matte.rs

1use crate::gpu::Trace;
2
3#[derive(Clone, Copy, Debug, Default, PartialEq)]
4pub struct Matte {
5    pub transparent: bool,
6    pub catcher: Option<u32>,
7}
8
9impl Matte {
10    pub fn transparent() -> Self {
11        Self {
12            transparent: true,
13            catcher: None,
14        }
15    }
16
17    pub fn catcher(material: u32) -> Self {
18        Self {
19            transparent: true,
20            catcher: Some(material),
21        }
22    }
23
24    pub fn lane(self) -> f32 {
25        match self.catcher {
26            Some(material) => material as f32 + 2.0,
27            None if self.transparent => 1.0,
28            None => 0.0,
29        }
30    }
31}
32
33impl Trace {
34    pub fn set_matte(&mut self, matte: Matte) -> Result<(), String> {
35        if let Some(material) = matte.catcher {
36            if material >= self.materials {
37                return Err(format!(
38                    "the shadow catcher names material {material}, the scene has {}",
39                    self.materials
40                ));
41            }
42            if material >= 1 << 23 {
43                return Err("the shadow catcher's material is past 2^23".into());
44            }
45            if !matte.transparent {
46                return Err("a shadow catcher needs the transparent film".into());
47            }
48        }
49        self.base.forward[3] = matte.lane();
50        self.samples = 0;
51        Ok(())
52    }
53
54    pub fn matte(&self) -> Matte {
55        let lane = self.base.forward[3];
56        Matte {
57            transparent: lane > 0.0,
58            catcher: (lane > 1.5).then(|| lane as u32 - 2),
59        }
60    }
61}
62
63#[derive(Clone, Debug, PartialEq)]
64pub struct Rgba {
65    pub width: u32,
66    pub height: u32,
67    pub texels: Vec<[f32; 4]>,
68}
69
70impl Rgba {
71    pub fn premultiplied(width: u32, height: u32, color: &[u8]) -> Result<Self, String> {
72        let texels = crate::probe::texels(color);
73        if texels.len() != width as usize * height as usize {
74            return Err("the colour readback does not match the size".into());
75        }
76        Ok(Self {
77            width,
78            height,
79            texels,
80        })
81    }
82
83    pub fn straight(&self) -> Vec<[f32; 4]> {
84        self.texels.iter().map(|&texel| straight(texel)).collect()
85    }
86}
87
88pub fn straight(texel: [f32; 4]) -> [f32; 4] {
89    let alpha = texel[3].clamp(0.0, 1.0);
90    if alpha <= 1.0 / 65535.0 {
91        return [0.0, 0.0, 0.0, 0.0];
92    }
93    [texel[0] / alpha, texel[1] / alpha, texel[2] / alpha, alpha]
94}
95
96pub fn over(texel: [f32; 4], background: [f32; 3]) -> [f32; 4] {
97    let alpha = texel[3].clamp(0.0, 1.0);
98    [
99        texel[0] + background[0] * (1.0 - alpha),
100        texel[1] + background[1] * (1.0 - alpha),
101        texel[2] + background[2] * (1.0 - alpha),
102        1.0,
103    ]
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn lanes_round_trip() {
112        assert_eq!(Matte::default().lane(), 0.0);
113        assert_eq!(Matte::transparent().lane(), 1.0);
114        assert_eq!(Matte::catcher(0).lane(), 2.0);
115        assert_eq!(Matte::catcher(7).lane(), 9.0);
116    }
117
118    #[test]
119    fn straight_alpha_undoes_premultiplying() {
120        assert_eq!(straight([0.2, 0.1, 0.05, 0.5]), [0.4, 0.2, 0.1, 0.5]);
121        assert_eq!(straight([0.0, 0.0, 0.0, 0.0]), [0.0; 4]);
122        assert_eq!(straight([0.3, 0.3, 0.3, 1.0]), [0.3, 0.3, 0.3, 1.0]);
123        assert_eq!(
124            over([0.0, 0.0, 0.0, 0.25], [1.0, 0.5, 0.0]),
125            [0.75, 0.375, 0.0, 1.0]
126        );
127        let rgba =
128            Rgba::premultiplied(1, 1, bytemuck::cast_slice(&[0.1f32, 0.2, 0.3, 0.5])).unwrap();
129        assert!((rgba.straight()[0][2] - 0.6).abs() < 1e-6);
130        assert!(Rgba::premultiplied(2, 1, &[0u8; 16]).is_err());
131    }
132}