Skip to main content

scena/render/
screen_space_reflections.rs

1use super::RasterTarget;
2
3#[derive(Debug, Clone, Copy, Default, PartialEq)]
4pub(super) struct MaterialReflectionPixel {
5    sample_x: f32,
6    sample_y: f32,
7    weight: f32,
8    roughness: f32,
9}
10
11impl MaterialReflectionPixel {
12    pub(super) fn new(sample_x: f32, sample_y: f32, weight: f32, roughness: f32) -> Option<Self> {
13        if !sample_x.is_finite()
14            || !sample_y.is_finite()
15            || !weight.is_finite()
16            || !roughness.is_finite()
17        {
18            return None;
19        }
20        let weight = weight.clamp(0.0, 0.88);
21        (weight > 0.0).then_some(Self {
22            sample_x,
23            sample_y,
24            weight,
25            roughness: roughness.clamp(0.0, 1.0),
26        })
27    }
28
29    const fn is_active(self) -> bool {
30        self.weight > 0.0
31    }
32}
33
34#[derive(Debug, Clone, Copy, PartialEq)]
35pub struct ScreenSpaceReflectionConfig {
36    strength: f32,
37    roughness: f32,
38    horizon_fraction: f32,
39    fade: f32,
40}
41
42impl ScreenSpaceReflectionConfig {
43    pub const fn studio_floor() -> Self {
44        Self {
45            strength: 0.72,
46            roughness: 0.18,
47            horizon_fraction: 0.58,
48            fade: 0.35,
49        }
50    }
51
52    pub fn new(strength: f32, roughness: f32, horizon_fraction: f32, fade: f32) -> Self {
53        Self {
54            strength: clamp_unit_or(strength, 0.72),
55            roughness: clamp_unit_or(roughness, 0.18),
56            horizon_fraction: clamp_unit_or(horizon_fraction, 0.58),
57            fade: clamp_unit_or(fade, 0.35),
58        }
59    }
60
61    pub const fn strength(self) -> f32 {
62        self.strength
63    }
64
65    pub const fn roughness(self) -> f32 {
66        self.roughness
67    }
68
69    pub const fn horizon_fraction(self) -> f32 {
70        self.horizon_fraction
71    }
72
73    pub const fn fade(self) -> f32 {
74        self.fade
75    }
76
77    pub fn roughness_radius_px(self) -> u32 {
78        (self.roughness * 8.0).round().clamp(0.0, 8.0) as u32
79    }
80}
81
82impl Default for ScreenSpaceReflectionConfig {
83    fn default() -> Self {
84        Self::studio_floor()
85    }
86}
87
88fn clamp_unit_or(value: f32, fallback: f32) -> f32 {
89    if value.is_finite() {
90        value.clamp(0.0, 1.0)
91    } else {
92        fallback
93    }
94}
95
96pub(super) fn apply_rgba8(
97    target: RasterTarget,
98    frame: &mut [u8],
99    scratch: &mut [u8],
100    config: ScreenSpaceReflectionConfig,
101) -> u64 {
102    if target.width < 3 || target.height < 3 || config.strength() <= 0.0 {
103        return 0;
104    }
105    debug_assert_eq!(frame.len(), target.byte_len());
106    debug_assert_eq!(scratch.len(), target.byte_len());
107    scratch.copy_from_slice(frame);
108
109    let horizon_y = ((target.height as f32) * config.horizon_fraction())
110        .round()
111        .clamp(1.0, (target.height.saturating_sub(2)) as f32) as u32;
112    let floor_height = target.height.saturating_sub(horizon_y).max(1);
113    let radius = config.roughness_radius_px();
114    for y in horizon_y..target.height {
115        let distance = (y.saturating_sub(horizon_y) as f32 / floor_height as f32).clamp(0.0, 1.0);
116        let fade = (1.0 - distance * config.fade()).clamp(0.0, 1.0);
117        let mirrored_y = horizon_y
118            .saturating_sub(y.saturating_sub(horizon_y))
119            .min(target.height - 1);
120        for x in 0..target.width {
121            let offset = pixel_offset(target, x, y);
122            let source_luma = luma_from_srgb8(&scratch[offset..offset + 4]) / 255.0;
123            let floor_mask = (1.0 - source_luma).clamp(0.0, 1.0);
124            let alpha = (config.strength() * fade * floor_mask).clamp(0.0, 1.0);
125            if alpha <= 0.0 {
126                continue;
127            }
128            let reflected = blurred_reflection_sample(target, scratch, x, mirrored_y, radius);
129            for channel in 0..3 {
130                let base = f32::from(scratch[offset + channel]);
131                let value = base * (1.0 - alpha) + reflected[channel] * alpha;
132                frame[offset + channel] = value.round().clamp(0.0, 255.0) as u8;
133            }
134            frame[offset + 3] = scratch[offset + 3];
135        }
136    }
137
138    1
139}
140
141pub(super) fn apply_material_rgba8(
142    target: RasterTarget,
143    frame: &mut [u8],
144    scratch: &mut [u8],
145    material_reflections: &[MaterialReflectionPixel],
146    config: ScreenSpaceReflectionConfig,
147) -> u64 {
148    if target.width < 3 || target.height < 3 || config.strength() <= 0.0 {
149        return 0;
150    }
151    debug_assert_eq!(frame.len(), target.byte_len());
152    debug_assert_eq!(scratch.len(), target.byte_len());
153    debug_assert_eq!(material_reflections.len(), target.pixel_len());
154    scratch.copy_from_slice(frame);
155
156    let mut touched = false;
157    for y in 0..target.height {
158        for x in 0..target.width {
159            let pixel_index = target.pixel_index(x, y);
160            let reflection = material_reflections[pixel_index];
161            if !reflection.is_active() {
162                continue;
163            }
164            touched = true;
165            let offset = pixel_index * 4;
166            let radius = ((reflection.roughness * reflection.roughness * 14.0
167                + config.roughness() * 5.0)
168                .round()
169                .clamp(0.0, 8.0)) as u32;
170            let reflected = blurred_reflection_sample_f32(
171                target,
172                scratch,
173                reflection.sample_x,
174                reflection.sample_y,
175                radius,
176            );
177            for channel in 0..3 {
178                let base = f32::from(scratch[offset + channel]);
179                let value =
180                    base * (1.0 - reflection.weight) + reflected[channel] * reflection.weight;
181                frame[offset + channel] = value.round().clamp(0.0, 255.0) as u8;
182            }
183            frame[offset + 3] = scratch[offset + 3];
184        }
185    }
186
187    u64::from(touched)
188}
189
190fn blurred_reflection_sample(
191    target: RasterTarget,
192    frame: &[u8],
193    x: u32,
194    y: u32,
195    radius: u32,
196) -> [f32; 3] {
197    if radius == 0 {
198        let offset = pixel_offset(target, x, y);
199        return [
200            f32::from(frame[offset]),
201            f32::from(frame[offset + 1]),
202            f32::from(frame[offset + 2]),
203        ];
204    }
205    let min_x = x.saturating_sub(radius);
206    let max_x = x.saturating_add(radius).min(target.width - 1);
207    let min_y = y.saturating_sub(radius);
208    let max_y = y.saturating_add(radius).min(target.height - 1);
209    let mut sum = [0.0_f32; 3];
210    let mut weight_sum = 0.0_f32;
211    for sample_y in min_y..=max_y {
212        for sample_x in min_x..=max_x {
213            let dx = sample_x.abs_diff(x) as f32;
214            let dy = sample_y.abs_diff(y) as f32;
215            let distance = (dx * dx + dy * dy).sqrt();
216            let weight = (1.0 - distance / (radius as f32 + 1.0)).max(0.0);
217            if weight <= 0.0 {
218                continue;
219            }
220            let offset = pixel_offset(target, sample_x, sample_y);
221            sum[0] += f32::from(frame[offset]) * weight;
222            sum[1] += f32::from(frame[offset + 1]) * weight;
223            sum[2] += f32::from(frame[offset + 2]) * weight;
224            weight_sum += weight;
225        }
226    }
227    if weight_sum <= 0.0 {
228        return blurred_reflection_sample(target, frame, x, y, 0);
229    }
230    [
231        sum[0] / weight_sum,
232        sum[1] / weight_sum,
233        sum[2] / weight_sum,
234    ]
235}
236
237fn blurred_reflection_sample_f32(
238    target: RasterTarget,
239    frame: &[u8],
240    x: f32,
241    y: f32,
242    radius: u32,
243) -> [f32; 3] {
244    if radius == 0 {
245        return bilinear_sample_srgb8(target, frame, x, y);
246    }
247    let center_x = x.clamp(0.0, target.width.saturating_sub(1) as f32);
248    let center_y = y.clamp(0.0, target.height.saturating_sub(1) as f32);
249    let min_x = (center_x.floor() as u32).saturating_sub(radius);
250    let max_x = (center_x.ceil() as u32)
251        .saturating_add(radius)
252        .min(target.width - 1);
253    let min_y = (center_y.floor() as u32).saturating_sub(radius);
254    let max_y = (center_y.ceil() as u32)
255        .saturating_add(radius)
256        .min(target.height - 1);
257    let mut sum = [0.0_f32; 3];
258    let mut weight_sum = 0.0_f32;
259    for sample_y in min_y..=max_y {
260        for sample_x in min_x..=max_x {
261            let dx = sample_x as f32 - center_x;
262            let dy = sample_y as f32 - center_y;
263            let distance = (dx * dx + dy * dy).sqrt();
264            let weight = (1.0 - distance / (radius as f32 + 1.0)).max(0.0);
265            if weight <= 0.0 {
266                continue;
267            }
268            let offset = pixel_offset(target, sample_x, sample_y);
269            sum[0] += f32::from(frame[offset]) * weight;
270            sum[1] += f32::from(frame[offset + 1]) * weight;
271            sum[2] += f32::from(frame[offset + 2]) * weight;
272            weight_sum += weight;
273        }
274    }
275    if weight_sum <= 0.0 {
276        return bilinear_sample_srgb8(target, frame, x, y);
277    }
278    [
279        sum[0] / weight_sum,
280        sum[1] / weight_sum,
281        sum[2] / weight_sum,
282    ]
283}
284
285fn bilinear_sample_srgb8(target: RasterTarget, frame: &[u8], x: f32, y: f32) -> [f32; 3] {
286    let x = x.clamp(0.0, target.width.saturating_sub(1) as f32);
287    let y = y.clamp(0.0, target.height.saturating_sub(1) as f32);
288    let x0 = x.floor() as u32;
289    let y0 = y.floor() as u32;
290    let x1 = x0.saturating_add(1).min(target.width - 1);
291    let y1 = y0.saturating_add(1).min(target.height - 1);
292    let tx = x - x0 as f32;
293    let ty = y - y0 as f32;
294    let c00 = sample_pixel_rgb(target, frame, x0, y0);
295    let c10 = sample_pixel_rgb(target, frame, x1, y0);
296    let c01 = sample_pixel_rgb(target, frame, x0, y1);
297    let c11 = sample_pixel_rgb(target, frame, x1, y1);
298    [
299        bilinear_channel(c00[0], c10[0], c01[0], c11[0], tx, ty),
300        bilinear_channel(c00[1], c10[1], c01[1], c11[1], tx, ty),
301        bilinear_channel(c00[2], c10[2], c01[2], c11[2], tx, ty),
302    ]
303}
304
305fn sample_pixel_rgb(target: RasterTarget, frame: &[u8], x: u32, y: u32) -> [f32; 3] {
306    let offset = pixel_offset(target, x, y);
307    [
308        f32::from(frame[offset]),
309        f32::from(frame[offset + 1]),
310        f32::from(frame[offset + 2]),
311    ]
312}
313
314fn bilinear_channel(c00: f32, c10: f32, c01: f32, c11: f32, tx: f32, ty: f32) -> f32 {
315    let top = c00 * (1.0 - tx) + c10 * tx;
316    let bottom = c01 * (1.0 - tx) + c11 * tx;
317    top * (1.0 - ty) + bottom * ty
318}
319
320fn pixel_offset(target: RasterTarget, x: u32, y: u32) -> usize {
321    target.pixel_index(x, y) * 4
322}
323
324fn luma_from_srgb8(pixel: &[u8]) -> f32 {
325    f32::from(pixel[0]) * 0.299 + f32::from(pixel[1]) * 0.587 + f32::from(pixel[2]) * 0.114
326}