Skip to main content

pfx_live/
reflect.rs

1use pfx_gpu::{GpuProfiler, wgpu};
2
3pub const REFLECTION_FORMAT: wgpu::TextureFormat = wgpu::TextureFormat::Rgba16Float;
4pub const GUIDE_BINDING: u32 = 26;
5pub const REFINE_STEPS: u32 = 7;
6pub const REACH: f32 = 64.0;
7pub const EDGE_MARGIN: f32 = 0.08;
8pub const DEPTH_TOLERANCE: f32 = 0.03;
9pub const THICKNESS: f32 = 0.03;
10pub const ROUGH_FADE: [f32; 2] = [0.6, 0.75];
11pub const CONSISTENT: f32 = 0.8;
12pub const JOIN: f32 = 0.005;
13
14pub type Matrix = [[f32; 4]; 4];
15
16pub fn half_size(width: u32, height: u32) -> (u32, u32) {
17    (width.div_ceil(2), height.div_ceil(2))
18}
19
20pub fn march_steps(roughness: f32) -> u32 {
21    if roughness > 0.45 {
22        8
23    } else if roughness > 0.18 {
24        12
25    } else {
26        16
27    }
28}
29
30pub fn rough_fade(roughness: f32) -> f32 {
31    1.0 - smoothstep(ROUGH_FADE[0], ROUGH_FADE[1], roughness)
32}
33
34pub fn trace_offset(frame_index: u32) -> [u32; 2] {
35    match frame_index % 4 {
36        0 => [0, 0],
37        1 => [1, 1],
38        2 => [1, 0],
39        _ => [0, 1],
40    }
41}
42
43pub fn edge_confidence(uv: [f32; 2], margin: f32) -> f32 {
44    if margin <= 0.0 {
45        return 0.0;
46    }
47    let distance = uv[0].min(1.0 - uv[0]).min(uv[1]).min(1.0 - uv[1]);
48    (distance / margin).clamp(0.0, 1.0)
49}
50
51fn smoothstep(low: f32, high: f32, x: f32) -> f32 {
52    let t = ((x - low) / (high - low)).clamp(0.0, 1.0);
53    t * t * (3.0 - 2.0 * t)
54}
55
56fn mix(a: f32, b: f32, t: f32) -> f32 {
57    a * (1.0 - t) + b * t
58}
59
60fn dot(a: [f32; 3], b: [f32; 3]) -> f32 {
61    a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
62}
63
64pub struct DepthImage<'a> {
65    pub width: u32,
66    pub height: u32,
67    pub depth: &'a [f32],
68}
69
70impl DepthImage<'_> {
71    pub fn at(&self, pixel: [f32; 2]) -> f32 {
72        let x = pixel[0].clamp(0.0, self.width as f32 - 1.0) as usize;
73        let y = pixel[1].clamp(0.0, self.height as f32 - 1.0) as usize;
74        self.depth[y * self.width as usize + x]
75    }
76}
77
78pub fn to_pixel(projection: &Matrix, size: [f32; 2], point: [f32; 3]) -> [f32; 3] {
79    let p = [point[0], point[1], point[2], 1.0];
80    let clip: [f32; 4] = std::array::from_fn(|row| {
81        (0..4)
82            .map(|column| projection[column][row] * p[column])
83            .sum()
84    });
85    [
86        (clip[0] / clip[3] * 0.5 + 0.5) * size[0],
87        (clip[1] / clip[3] * -0.5 + 0.5) * size[1],
88        clip[3],
89    ]
90}
91
92#[derive(Clone, Copy, Debug, PartialEq)]
93pub struct Hit {
94    pub pixel: [f32; 2],
95    pub confidence: f32,
96}
97
98pub fn march(
99    depth: &DepthImage<'_>,
100    projection: &Matrix,
101    origin: [f32; 3],
102    direction: [f32; 3],
103    steps: u32,
104    thickness: f32,
105) -> Option<Hit> {
106    let size = [depth.width as f32, depth.height as f32];
107    let near = -origin[2];
108    let mut reach = near * REACH;
109    if direction[2] > 0.0 {
110        reach = reach.min(0.9 * near / direction[2]);
111    }
112    let start = to_pixel(projection, size, origin);
113    let end = to_pixel(
114        projection,
115        size,
116        std::array::from_fn(|axis| origin[axis] + direction[axis] * reach),
117    );
118    let delta = [end[0] - start[0], end[1] - start[1]];
119    let mut limit = 1.0f32;
120    for axis in 0..2 {
121        if delta[axis] > 0.0 {
122            limit = limit.min((size[axis] - 0.5 - start[axis]) / delta[axis]);
123        }
124        if delta[axis] < 0.0 {
125            limit = limit.min((0.5 - start[axis]) / delta[axis]);
126        }
127    }
128    let span = (delta[0] * delta[0] + delta[1] * delta[1]).sqrt();
129    let pixels = span * limit;
130    if pixels.is_nan() || pixels < 1.0 {
131        return None;
132    }
133    let k = [1.0 / start[2], 1.0 / end[2]];
134    let ray_depth = |s: f32| 1.0 / mix(k[0], k[1], s);
135    let place = |s: f32| [start[0] + delta[0] * s, start[1] + delta[1] * s];
136    let mut previous = 0.0f32;
137    let mut previous_depth = start[2];
138    let mut front = start[2] <= depth.at([start[0], start[1]]);
139    for i in 1..=steps {
140        let u = i as f32 / steps as f32;
141        let travelled = (pixels * u * u).max(i as f32).min(pixels);
142        let s = travelled / span;
143        let ray = ray_depth(s);
144        let scene = depth.at(place(s));
145        if ray > scene && front && previous_depth.min(ray) <= scene * (1.0 + thickness) {
146            let mut lo = previous;
147            let mut hi = s;
148            for _ in 0..REFINE_STEPS {
149                let mid = 0.5 * (lo + hi);
150                if ray_depth(mid) > depth.at(place(mid)) {
151                    hi = mid;
152                } else {
153                    lo = mid;
154                }
155            }
156            let pixel = place(hi);
157            let hit_ray = ray_depth(hi);
158            let hit_scene = depth.at(pixel);
159            let span_depth = (hit_ray - ray_depth(lo)).abs();
160            let beyond = (depth.at(place(hi + 1.0 / span)) - hit_scene).abs();
161            if hit_ray - hit_scene <= hit_scene * thickness + span_depth
162                && (depth.at(place(lo)) - hit_scene).abs()
163                    <= hit_scene * JOIN + 2.0 * beyond + span_depth
164            {
165                let fade = 1.0 - smoothstep(0.8, 1.0, hi);
166                let uv = [pixel[0] / size[0], pixel[1] / size[1]];
167                return Some(Hit {
168                    pixel,
169                    confidence: edge_confidence(uv, EDGE_MARGIN) * fade,
170                });
171            }
172        }
173        previous = s;
174        previous_depth = ray;
175        front = ray <= scene;
176        if travelled >= pixels {
177            break;
178        }
179    }
180    None
181}
182
183pub fn similarity(
184    depth: f32,
185    normal: [f32; 3],
186    other_depth: f32,
187    other_normal: [f32; 3],
188    tolerance: f32,
189) -> f32 {
190    if depth <= 0.0 || other_depth <= 0.0 {
191        return 0.0;
192    }
193    let near = (1.0 - (depth - other_depth).abs() / tolerance.max(1e-6)).clamp(0.0, 1.0);
194    let facing = smoothstep(0.6, 0.85, dot(normal, other_normal));
195    near * facing
196}
197
198#[derive(Clone, Copy, Debug, PartialEq)]
199pub struct Texel {
200    pub colour: [f32; 4],
201    pub depth: f32,
202    pub normal: [f32; 3],
203}
204
205pub fn consistent(width: u32, height: u32, texels: &[Texel], x: u32, y: u32) -> bool {
206    let at = |dx: u32, dy: u32| {
207        let cx = (x + dx).min(width - 1) as usize;
208        let cy = (y + dy).min(height - 1) as usize;
209        texels[cy * width as usize + cx]
210    };
211    let centre = at(0, 0);
212    if centre.depth <= 0.0 {
213        return false;
214    }
215    let right = at(1, 0);
216    let down = at(0, 1);
217    let slope = (right.depth - centre.depth)
218        .abs()
219        .max((down.depth - centre.depth).abs())
220        .min(centre.depth * DEPTH_TOLERANCE);
221    let tolerance = centre.depth * DEPTH_TOLERANCE + slope;
222    [
223        (right, tolerance),
224        (down, tolerance),
225        (at(1, 1), tolerance + slope),
226    ]
227    .iter()
228    .all(|(other, tolerance)| {
229        similarity(
230            centre.depth,
231            centre.normal,
232            other.depth,
233            other.normal,
234            *tolerance,
235        ) > CONSISTENT
236    })
237}
238
239pub fn upsample(
240    width: u32,
241    height: u32,
242    texels: &[Texel],
243    uv: [f32; 2],
244    depth: f32,
245    normal: [f32; 3],
246) -> [f32; 4] {
247    let f = [uv[0] * width as f32 - 0.5, uv[1] * height as f32 - 0.5];
248    let base = [f[0].floor(), f[1].floor()];
249    let t = [f[0] - base[0], f[1] - base[1]];
250    let tolerance = depth * DEPTH_TOLERANCE;
251    let cell = |i: usize| {
252        let offset = [(i & 1) as i32, (i >> 1) as i32];
253        let x = (base[0] as i32 + offset[0]).clamp(0, width as i32 - 1) as u32;
254        let y = (base[1] as i32 + offset[1]).clamp(0, height as i32 - 1) as u32;
255        let bilinear = if offset[0] == 1 { t[0] } else { 1.0 - t[0] }
256            * if offset[1] == 1 { t[1] } else { 1.0 - t[1] };
257        (x, y, bilinear)
258    };
259    let (x0, y0, _) = cell(0);
260    let first = texels[(y0 * width + x0) as usize];
261    let fast = consistent(width, height, texels, x0, y0)
262        && similarity(depth, normal, first.depth, first.normal, tolerance) > CONSISTENT;
263    let mut sum = [0.0f32; 4];
264    let mut total = 0.0f32;
265    for i in 0..4 {
266        let (x, y, bilinear) = cell(i);
267        let texel = texels[(y * width + x) as usize];
268        let weight = if fast {
269            bilinear
270        } else {
271            (bilinear + 0.01) * similarity(depth, normal, texel.depth, texel.normal, tolerance)
272        };
273        for (channel, value) in sum.iter_mut().zip(texel.colour) {
274            *channel += value * weight;
275        }
276        total += weight;
277    }
278    if total < 0.001 {
279        return [0.0; 4];
280    }
281    sum.map(|value| value / total)
282}
283
284pub fn accumulate(
285    current: [f32; 4],
286    low: [f32; 4],
287    high: [f32; 4],
288    history: [f32; 4],
289    coverage: f32,
290    roughness: f32,
291    history_valid: bool,
292) -> [f32; 4] {
293    if !history_valid || coverage <= 0.001 {
294        return current;
295    }
296    let weight = mix(0.6, 0.9, roughness.clamp(0.0, 1.0)) * (coverage * 2.0).clamp(0.0, 1.0);
297    std::array::from_fn(|channel| {
298        let old = history[channel].clamp(low[channel], high[channel]);
299        mix(current[channel], old, weight)
300    })
301}
302
303macro_rules! guide_wgsl {
304    () => {
305        r#"
306fn ssr_encode_normal(n: vec3f) -> vec2f {
307    let p = n.xy / (abs(n.x) + abs(n.y) + abs(n.z));
308    if (n.z >= 0.0) { return p; }
309    return (vec2f(1.0) - abs(p.yx)) * select(vec2f(-1.0), vec2f(1.0), p >= vec2f(0.0));
310}
311fn ssr_decode_normal(e: vec2f) -> vec3f {
312    var n = vec3f(e, 1.0 - abs(e.x) - abs(e.y));
313    let t = max(-n.z, 0.0);
314    n.x += select(t, -t, n.x >= 0.0);
315    n.y += select(t, -t, n.y >= 0.0);
316    return normalize(n);
317}
318fn ssr_similarity(depth: f32, normal: vec3f, other_depth: f32, other_normal: vec3f, tolerance: f32) -> f32 {
319    if (depth <= 0.0 || other_depth <= 0.0) { return 0.0; }
320    let near = clamp(1.0 - abs(depth - other_depth) / max(tolerance, 0.000001), 0.0, 1.0);
321    let facing = smoothstep(0.6, 0.85, dot(normal, other_normal));
322    return near * facing;
323}
324const SSR_DEPTH_TOLERANCE: f32 = 0.03;
325"#
326    };
327}
328
329macro_rules! params_wgsl {
330    () => {
331        r#"
332struct ReflectParams {
333    size: vec2u,
334    half_size: vec2u,
335    projection: mat4x4f,
336    inverse_projection: mat4x4f,
337    history_valid: u32,
338    frame_index: u32,
339    refine: u32,
340    thickness: f32,
341    history_half: vec2u,
342}
343@group(0) @binding(0) var<uniform> params: ReflectParams;
344fn inside(uv: vec2f) -> bool {
345    return all(uv >= vec2f(0.0)) && all(uv < vec2f(1.0));
346}
347fn trace_pixel(cell: vec2u) -> vec2u {
348    let phase = params.frame_index % 4u;
349    let offset = vec2u(select(0u, 1u, phase == 1u || phase == 2u), select(0u, 1u, phase == 1u || phase == 3u));
350    return min(cell * 2u + offset, params.size - 1u);
351}
352"#
353    };
354}
355
356pub const TRACE_WGSL: &str = concat!(
357    guide_wgsl!(),
358    params_wgsl!(),
359    r#"
360@group(0) @binding(1) var depth_tex: texture_2d<f32>;
361@group(0) @binding(2) var normal_roughness: texture_2d<f32>;
362@group(0) @binding(3) var source_colour: texture_2d<f32>;
363@group(0) @binding(4) var colour_sampler: sampler;
364@group(0) @binding(5) var raw_out: texture_storage_2d<rgba16float, write>;
365@group(0) @binding(6) var guide_out: texture_storage_2d<rgba16float, write>;
366fn edge_weight(uv: vec2f) -> f32 {
367    let d = min(min(uv.x, 1.0 - uv.x), min(uv.y, 1.0 - uv.y));
368    return clamp(d / 0.08, 0.0, 1.0);
369}
370fn to_pixel(p: vec3f) -> vec3f {
371    let clip = params.projection * vec4f(p, 1.0);
372    let ndc = clip.xy / clip.w;
373    return vec3f((ndc * vec2f(0.5, -0.5) + vec2f(0.5)) * vec2f(params.size), clip.w);
374}
375fn scene_depth(pixel: vec2f) -> f32 {
376    let at = vec2i(clamp(pixel, vec2f(0.0), vec2f(params.size) - vec2f(1.0)));
377    return textureLoad(depth_tex, at, 0).r;
378}
379fn ray_depth(k: vec2f, s: f32) -> f32 {
380    return 1.0 / mix(k.x, k.y, s);
381}
382fn march_steps(roughness: f32) -> u32 {
383    return select(select(16u, 12u, roughness > 0.18), 8u, roughness > 0.45);
384}
385fn march(origin: vec3f, direction: vec3f, steps: u32) -> vec3f {
386    let near = -origin.z;
387    var reach = near * 64.0;
388    if (direction.z > 0.0) {
389        reach = min(reach, 0.9 * near / direction.z);
390    }
391    let start = to_pixel(origin);
392    let end = to_pixel(origin + direction * reach);
393    let delta = end.xy - start.xy;
394    let size = vec2f(params.size);
395    var limit = 1.0;
396    if (delta.x > 0.0) { limit = min(limit, (size.x - 0.5 - start.x) / delta.x); }
397    if (delta.x < 0.0) { limit = min(limit, (0.5 - start.x) / delta.x); }
398    if (delta.y > 0.0) { limit = min(limit, (size.y - 0.5 - start.y) / delta.y); }
399    if (delta.y < 0.0) { limit = min(limit, (0.5 - start.y) / delta.y); }
400    let span = length(delta);
401    let pixels = span * limit;
402    if (!(pixels >= 1.0)) { return vec3f(0.0); }
403    let k = vec2f(1.0 / start.z, 1.0 / end.z);
404    var previous = 0.0;
405    var previous_depth = start.z;
406    var front = start.z <= scene_depth(start.xy);
407    for (var i = 1u; i <= steps; i++) {
408        let u = f32(i) / f32(steps);
409        let travelled = min(max(pixels * u * u, f32(i)), pixels);
410        let s = travelled / span;
411        let ray = ray_depth(k, s);
412        let scene = scene_depth(start.xy + delta * s);
413        if (ray > scene && front && min(previous_depth, ray) <= scene * (1.0 + params.thickness)) {
414            var lo = previous;
415            var hi = s;
416            for (var j = 0u; j < params.refine; j++) {
417                let mid = 0.5 * (lo + hi);
418                if (ray_depth(k, mid) > scene_depth(start.xy + delta * mid)) {
419                    hi = mid;
420                } else {
421                    lo = mid;
422                }
423            }
424            let pixel = start.xy + delta * hi;
425            let hit_ray = ray_depth(k, hi);
426            let hit_scene = scene_depth(pixel);
427            let span_depth = abs(hit_ray - ray_depth(k, lo));
428            let beyond = abs(scene_depth(start.xy + delta * (hi + 1.0 / span)) - hit_scene);
429            let joined = abs(scene_depth(start.xy + delta * lo) - hit_scene) <= hit_scene * 0.005 + 2.0 * beyond + span_depth;
430            if (hit_ray - hit_scene <= hit_scene * params.thickness + span_depth && joined) {
431                let fade = 1.0 - smoothstep(0.8, 1.0, hi);
432                return vec3f(pixel, edge_weight(pixel / size) * fade);
433            }
434        }
435        previous = s;
436        previous_depth = ray;
437        front = ray <= scene;
438        if (travelled >= pixels) { break; }
439    }
440    return vec3f(0.0);
441}
442fn view_position(pixel: vec2i) -> vec3f {
443    let at = clamp(pixel, vec2i(0), vec2i(params.size) - vec2i(1));
444    let depth = textureLoad(depth_tex, at, 0).r;
445    let uv = (vec2f(at) + vec2f(0.5)) / vec2f(params.size);
446    let clip = vec4f(uv * vec2f(2.0, -2.0) + vec2f(-1.0, 1.0), 1.0, 1.0);
447    let unprojected = params.inverse_projection * clip;
448    return unprojected.xyz / unprojected.w * (depth / max(-unprojected.z / unprojected.w, 0.0001));
449}
450fn nearer(origin: vec3f, a: vec3f, b: vec3f) -> vec3f {
451    return select(origin - b, a - origin, abs(a.z - origin.z) < abs(b.z - origin.z));
452}
453@compute @workgroup_size(8, 8)
454fn trace_pass(@builtin(global_invocation_id) id: vec3u) {
455    if (any(id.xy >= params.half_size)) { return; }
456    let pixel = trace_pixel(id.xy);
457    let depth = textureLoad(depth_tex, vec2i(pixel), 0).r;
458    let nr = textureLoad(normal_roughness, vec2i(pixel), 0);
459    if (depth <= 0.0 || length(nr.xyz) < 0.1) {
460        textureStore(raw_out, vec2i(id.xy), vec4f(0.0));
461        textureStore(guide_out, vec2i(id.xy), vec4f(0.0));
462        return;
463    }
464    let p = vec2i(pixel);
465    let origin = view_position(p);
466    let across = nearer(origin, view_position(p + vec2i(1, 0)), view_position(p - vec2i(1, 0)));
467    let along = nearer(origin, view_position(p + vec2i(0, 1)), view_position(p - vec2i(0, 1)));
468    let normal = normalize(nr.xyz);
469    var surface = cross(along, across);
470    surface = select(normal, normalize(surface), dot(surface, surface) > 0.0);
471    surface = select(surface, -surface, dot(surface, origin) > 0.0);
472    let roughness = clamp(nr.w, 0.0, 1.0);
473    let view_ray = normalize(origin);
474    let angle = 6.2831853 * f32(params.frame_index % 8u) * 0.125;
475    let jitter = vec3f(cos(angle), sin(angle), 0.0) * roughness * roughness * 0.15;
476    var direction = normalize(reflect(view_ray, normal) + jitter);
477    let lift = 0.05 - dot(direction, surface);
478    if (lift > 0.0) {
479        direction = normalize(direction + surface * lift);
480    }
481    var result = vec4f(0.0);
482    var hit = vec3f(0.0);
483    let fade = 1.0 - smoothstep(0.6, 0.75, roughness);
484    if (fade > 0.0) {
485        hit = march(origin + surface * depth * 0.002, direction, march_steps(roughness));
486    }
487    if (hit.z > 0.0) {
488        let at = vec2i(clamp(hit.xy, vec2f(0.0), vec2f(params.size) - vec2f(1.0)));
489        let facing = dot(textureLoad(normal_roughness, at, 0).xyz, direction);
490        let confidence = hit.z * fade * clamp(1.0 - 4.0 * (facing - 0.1), 0.0, 1.0);
491        let lod = clamp(roughness * 4.0 - 0.5, 0.0, f32(textureNumLevels(source_colour) - 1u));
492        let colour = textureSampleLevel(source_colour, colour_sampler, hit.xy / vec2f(textureDimensions(source_colour)), lod).rgb;
493        result = vec4f(colour * confidence, confidence);
494    }
495    textureStore(raw_out, vec2i(id.xy), result);
496    textureStore(guide_out, vec2i(id.xy), vec4f(depth, ssr_encode_normal(normal), roughness));
497}
498"#
499);
500
501pub const RESOLVE_WGSL: &str = concat!(
502    guide_wgsl!(),
503    params_wgsl!(),
504    r#"
505@group(0) @binding(1) var raw_tex: texture_2d<f32>;
506@group(0) @binding(2) var guide_tex: texture_2d<f32>;
507@group(0) @binding(3) var history_tex: texture_2d<f32>;
508@group(0) @binding(4) var history_guide: texture_2d<f32>;
509@group(0) @binding(5) var motion_tex: texture_2d<f32>;
510@group(0) @binding(6) var output_tex: texture_storage_2d<rgba16float, write>;
511@group(0) @binding(7) var guide_out: texture_storage_2d<rgba16float, write>;
512@compute @workgroup_size(8, 8)
513fn resolve_pass(@builtin(global_invocation_id) id: vec3u) {
514    if (any(id.xy >= params.half_size)) { return; }
515    let cell = vec2i(id.xy);
516    let last = vec2i(params.half_size) - vec2i(1);
517    let current = textureLoad(raw_tex, cell, 0);
518    let guide = textureLoad(guide_tex, cell, 0);
519    if (guide.x <= 0.0) {
520        textureStore(output_tex, cell, vec4f(0.0));
521        textureStore(guide_out, cell, vec4f(0.0));
522        return;
523    }
524    let normal = ssr_decode_normal(guide.yz);
525    let right = textureLoad(guide_tex, clamp(cell + vec2i(1, 0), vec2i(0), last), 0);
526    let down = textureLoad(guide_tex, clamp(cell + vec2i(0, 1), vec2i(0), last), 0);
527    let diagonal = textureLoad(guide_tex, clamp(cell + vec2i(1, 1), vec2i(0), last), 0);
528    let slope = min(max(abs(right.x - guide.x), abs(down.x - guide.x)), guide.x * SSR_DEPTH_TOLERANCE);
529    let tolerance = guide.x * SSR_DEPTH_TOLERANCE + slope;
530    let consistent = ssr_similarity(guide.x, normal, right.x, ssr_decode_normal(right.yz), tolerance) > 0.8
531        && ssr_similarity(guide.x, normal, down.x, ssr_decode_normal(down.yz), tolerance) > 0.8
532        && ssr_similarity(guide.x, normal, diagonal.x, ssr_decode_normal(diagonal.yz), tolerance + slope) > 0.8;
533    textureStore(guide_out, cell, vec4f(guide.xyz, select(-1.0 - guide.w, guide.w, consistent)));
534    var low = current;
535    var high = current;
536    for (var i = 0; i < 4; i++) {
537        let offset = select(vec2i(0, select(-1, 1, i == 3)), vec2i(select(-1, 1, i == 1), 0), i < 2);
538        let near = textureLoad(raw_tex, clamp(cell + offset, vec2i(0), last), 0);
539        low = min(low, near);
540        high = max(high, near);
541    }
542    var result = current;
543    let pixel = trace_pixel(id.xy);
544    let prev = (vec2f(pixel) + vec2f(0.5)) / vec2f(params.size) + textureLoad(motion_tex, vec2i(pixel), 0).rg;
545    if (params.history_valid != 0u && inside(prev)) {
546        let history_last = vec2i(params.history_half) - vec2i(1);
547        let f = prev * vec2f(params.history_half) - vec2f(0.5);
548        let base = vec2i(floor(f));
549        let t = f - floor(f);
550        var sum = vec4f(0.0);
551        var total = 0.0;
552        var coverage = 0.0;
553        for (var i = 0; i < 4; i++) {
554            let offset = vec2i(i & 1, i >> 1u);
555            let at = clamp(base + offset, vec2i(0), history_last);
556            let bilinear = select(1.0 - t.x, t.x, offset.x == 1) * select(1.0 - t.y, t.y, offset.y == 1);
557            let g = textureLoad(history_guide, at, 0);
558            let similar = ssr_similarity(guide.x, normal, g.x, ssr_decode_normal(g.yz), tolerance);
559            let weight = (bilinear + 0.01) * similar;
560            sum += textureLoad(history_tex, at, 0) * weight;
561            total += weight;
562            coverage += bilinear * similar;
563        }
564        if (total > 0.001 && coverage > 0.001) {
565            let old = clamp(sum / total, low, high);
566            let weight = mix(0.6, 0.9, clamp(guide.w, 0.0, 1.0)) * clamp(coverage * 2.0, 0.0, 1.0);
567            result = mix(current, old, weight);
568        }
569    }
570    textureStore(output_tex, cell, result);
571}
572"#
573);
574
575pub const DFG_FIT: [[f32; 3]; 24] = [
576    [0.8148172, 0.1125372, -0.01374308],
577    [-0.3511763, -0.1810853, 0.1673936],
578    [-0.02032731, -0.4825735, 0.04973246],
579    [0.462142, -0.411372, -0.06177693],
580    [0.02489464, 0.1060682, 0.07428504],
581    [-0.4534661, 0.6691472, -0.1939928],
582    [0.05596158, -0.1952609, -0.01097813],
583    [0.5484627, 0.5771128, -0.3130064],
584    [0.2092212, 2.050402, 1.003509],
585    [-2.763494, -3.524146, -0.6253516],
586    [-0.6803499, -2.683833, -1.808813],
587    [2.698063, 3.846977, 1.655182],
588    [0.03557167, -0.07356264, 0.03812399],
589    [0.1204449, -0.2777732, 0.1592967],
590    [-0.1640982, 0.3660604, -0.2035734],
591    [-0.3198297, 0.7459178, -0.4316107],
592    [0.2148219, -0.494718, 0.2827545],
593    [0.1391939, -0.3242128, 0.1875587],
594    [0.1873088, 0.3210023, 0.1747016],
595    [-1.316489, -1.468577, -1.091297],
596    [1.369269, -0.2278983, -0.9679774],
597    [1.550273, 4.079052, 4.483775],
598    [-1.484991, -0.5113827, 0.5854544],
599    [-0.3898746, -2.342672, -3.369676],
600];
601
602pub fn environment_dfg(roughness: f32, ndv: f32) -> [f32; 2] {
603    let u = 2.0 * roughness.clamp(0.05, 1.0) - 1.0;
604    let n = ndv.clamp(0.02, 1.0);
605    let s = 2.0 * n.sqrt() - 1.0;
606    let p = [1.0, s, s * s];
607    let f5 = (1.0 - n).powi(5);
608    let row = |k: usize| {
609        (0..3)
610            .map(|j| p[j] * (DFG_FIT[k][j] + f5 * DFG_FIT[k + 6][j]))
611            .sum::<f32>()
612    };
613    let mut dfg = [0.0f32; 2];
614    for i in (0..6).rev() {
615        dfg = [dfg[0] * u + row(i), dfg[1] * u + row(12 + i)];
616    }
617    dfg.map(|v| v.clamp(0.0, 1.0))
618}
619
620pub const PROBE_WGSL: &str = concat!(
621    guide_wgsl!(),
622    r#"
623override use_local_reflections: bool = true;
624const LOCAL_DEPTH_LEVELS: f32 = 6.0;
625const LOCAL_DEPTH_START: f32 = 0.02;
626const LOCAL_DEPTH_GROWTH: f32 = 2.5;
627const LOCAL_DEPTH_STEPS: u32 = 9u;
628const LOCAL_DEPTH_REFINE: u32 = 4u;
629const LOCAL_DEPTH_TOLERANCE: f32 = 0.02;
630const LOCAL_DEPTH_THICKNESS: f32 = 0.03;
631struct ReflectionProbe {
632    center: vec3f,
633    layer: i32,
634    box_min: vec3f,
635    priority: f32,
636    box_max: vec3f,
637    fade: f32,
638}
639fn local_probe_direction(probe: ReflectionProbe, world: vec3f, direction: vec3f) -> vec3f {
640    let safe = select(vec3f(-1.0), vec3f(1.0), direction >= vec3f(0.0)) * max(abs(direction), vec3f(0.00001));
641    let a = (probe.box_min - world) / safe;
642    let b = (probe.box_max - world) / safe;
643    let reach = min(max(a.x, b.x), min(max(a.y, b.y), max(a.z, b.z)));
644    return normalize(world + direction * max(reach, 0.0) - probe.center);
645}
646fn ssr_resolve(colour: texture_2d<f32>, guide: texture_2d<f32>, samp: sampler, uv: vec2f, written: vec2f, depth: f32, normal: vec3f) -> vec4f {
647    let dimensions = vec2i(textureDimensions(colour));
648    if (any(dimensions != vec2i(textureDimensions(guide)))) { return vec4f(0.0); }
649    let size = min(dimensions, vec2i(written));
650    let f = uv * vec2f(size) - vec2f(0.5);
651    let base = vec2i(floor(f));
652    let t = f - floor(f);
653    let tolerance = depth * SSR_DEPTH_TOLERANCE;
654    let first = textureLoad(guide, clamp(base, vec2i(0), size - vec2i(1)), 0);
655    if (first.w >= 0.0 && ssr_similarity(depth, normal, first.x, ssr_decode_normal(first.yz), tolerance) > 0.8) {
656        let inner = any(size < dimensions);
657        let texel = select(uv * vec2f(size), min(uv * vec2f(size), vec2f(size) - vec2f(0.5)), inner);
658        return textureSampleLevel(colour, samp, select(uv, texel / vec2f(dimensions), inner), 0.0);
659    }
660    var sum = vec4f(0.0);
661    var total = 0.0;
662    for (var i = 0; i < 4; i++) {
663        let offset = vec2i(i & 1, i >> 1u);
664        let at = clamp(base + offset, vec2i(0), size - vec2i(1));
665        let bilinear = select(1.0 - t.x, t.x, offset.x == 1) * select(1.0 - t.y, t.y, offset.y == 1);
666        let g = textureLoad(guide, at, 0);
667        let weight = (bilinear + 0.01) * ssr_similarity(depth, normal, g.x, ssr_decode_normal(g.yz), tolerance);
668        sum += textureLoad(colour, at, 0) * weight;
669        total += weight;
670    }
671    if (total < 0.001) { return vec4f(0.0); }
672    return sum / total;
673}
674fn local_probe_weight(probe: ReflectionProbe, world: vec3f) -> vec2f {
675    if (any(world < probe.box_min) || any(world > probe.box_max)) {
676        return vec2f(0.0);
677    }
678    let low = world - probe.box_min;
679    let high = probe.box_max - world;
680    let edge = min(min(min(low.x, low.y), low.z), min(min(high.x, high.y), high.z));
681    let seam = select(clamp(edge / probe.fade, 0.0, 1.0), 1.0, probe.fade == 0.0);
682    return vec2f(seam * (1.0 + max(probe.priority, 0.0)) / (length(world - probe.center) + 0.001), seam);
683}
684fn local_probe_weights(probe_a: ReflectionProbe, probe_b: ReflectionProbe, present: vec2f, world: vec3f) -> vec2f {
685    let a = local_probe_weight(probe_a, world) * present.x;
686    let b = local_probe_weight(probe_b, world) * present.y;
687    let total = a.x + b.x;
688    if (total <= 0.0) {
689        return vec2f(0.0);
690    }
691    let coverage = max(a.y, b.y);
692    return vec2f(a.x, b.x) * coverage / total;
693}
694fn local_probe_gap(probe: ReflectionProbe, reached: vec3f, local: texture_cube_array<f32>, samp: sampler, level: f32) -> f32 {
695    let distance = max(length(reached), 0.000001);
696    let stored = 1.0 / max(-1.0 - textureSampleLevel(local, samp, reached / distance, probe.layer, level).a, 0.00001);
697    return (distance - stored) / distance;
698}
699fn local_probe_lookup(probe: ReflectionProbe, world: vec3f, direction: vec3f, local: texture_cube_array<f32>, samp: sampler, local_mip: f32) -> vec3f {
700    let level = max(local_mip - LOCAL_DEPTH_LEVELS, 0.0);
701    if (textureSampleLevel(local, samp, direction, probe.layer, level).a >= 0.0) {
702        return local_probe_direction(probe, world, direction);
703    }
704    let offset = world - probe.center;
705    var ahead = local_probe_gap(probe, offset, local, samp, level) < LOCAL_DEPTH_THICKNESS;
706    var near = 0.0;
707    var t = LOCAL_DEPTH_START;
708    for (var step = 0u; step < LOCAL_DEPTH_STEPS; step++) {
709        let outer = local_probe_gap(probe, offset + direction * t, local, samp, level);
710        if (outer <= LOCAL_DEPTH_TOLERANCE) {
711            ahead = true;
712        } else {
713            if (ahead) {
714                var low = near;
715                var high = t;
716                var gap = outer;
717                for (var refine = 0u; refine < LOCAL_DEPTH_REFINE; refine++) {
718                    let middle = (low + high) * 0.5;
719                    let found = local_probe_gap(probe, offset + direction * middle, local, samp, level);
720                    if (found > LOCAL_DEPTH_TOLERANCE) {
721                        high = middle;
722                        gap = found;
723                    } else {
724                        low = middle;
725                    }
726                }
727                let reached = offset + direction * high;
728                if (gap < LOCAL_DEPTH_THICKNESS + (high - low) / max(length(reached), 0.000001)) {
729                    return select(direction, normalize(reached), dot(reached, reached) > 0.0000001);
730                }
731            }
732            ahead = false;
733        }
734        near = t;
735        t *= LOCAL_DEPTH_GROWTH;
736    }
737    return direction;
738}
739fn local_probe_reference(probe: ReflectionProbe, world: vec3f, n: vec3f, local: texture_cube_array<f32>, samp: sampler, local_mip: f32) -> vec3f {
740    let along = textureSampleLevel(local, samp, n, probe.layer, local_mip);
741    if (along.a < 0.0) {
742        return along.rgb;
743    }
744    return textureSampleLevel(local, samp, local_probe_direction(probe, world, n), probe.layer, local_mip).rgb;
745}
746struct LocalLookup {
747    weights: vec2f,
748    first: vec3f,
749    second: vec3f,
750}
751fn local_lookup(probe_a: ReflectionProbe, probe_b: ReflectionProbe, present: vec2f, world: vec3f, direction: vec3f, local: texture_cube_array<f32>, samp: sampler, local_mip: f32) -> LocalLookup {
752    var lookup = LocalLookup(vec2f(0.0), direction, direction);
753    if (!use_local_reflections) {
754        return lookup;
755    }
756    let weights = local_probe_weights(probe_a, probe_b, present, world);
757    let a = clamp(weights.x, 0.0, 1.0);
758    lookup.weights = vec2f(a, clamp(weights.y, 0.0, 1.0 - a));
759    if (lookup.weights.x > 0.0) {
760        lookup.first = local_probe_lookup(probe_a, world, direction, local, samp, local_mip);
761    }
762    if (lookup.weights.y > 0.0) {
763        lookup.second = local_probe_lookup(probe_b, world, direction, local, samp, local_mip);
764    }
765    return lookup;
766}
767fn reflection_dfg(roughness: f32, ndv: f32) -> vec2f {
768    let u = 2.0 * clamp(roughness, 0.05, 1.0) - 1.0;
769    let n = clamp(ndv, 0.02, 1.0);
770    let s = 2.0 * sqrt(n) - 1.0;
771    let p = vec3f(1.0, s, s * s);
772    let f5 = pow(1.0 - n, 5.0);
773    var dfg = vec2f(0.0);
774    dfg = dfg * u + vec2f(dot(p, vec3f(-0.4534661, 0.6691472, -0.1939928) + f5 * vec3f(2.698063, 3.846977, 1.655182)), dot(p, vec3f(0.1391939, -0.3242128, 0.1875587) + f5 * vec3f(-0.3898746, -2.342672, -3.369676)));
775    dfg = dfg * u + vec2f(dot(p, vec3f(0.02489464, 0.1060682, 0.07428504) + f5 * vec3f(-0.6803499, -2.683833, -1.808813)), dot(p, vec3f(0.2148219, -0.494718, 0.2827545) + f5 * vec3f(-1.484991, -0.5113827, 0.5854544)));
776    dfg = dfg * u + vec2f(dot(p, vec3f(0.462142, -0.411372, -0.06177693) + f5 * vec3f(-2.763494, -3.524146, -0.6253516)), dot(p, vec3f(-0.3198297, 0.7459178, -0.4316107) + f5 * vec3f(1.550273, 4.079052, 4.483775)));
777    dfg = dfg * u + vec2f(dot(p, vec3f(-0.02032731, -0.4825735, 0.04973246) + f5 * vec3f(0.2092212, 2.050402, 1.003509)), dot(p, vec3f(-0.1640982, 0.3660604, -0.2035734) + f5 * vec3f(1.369269, -0.2278983, -0.9679774)));
778    dfg = dfg * u + vec2f(dot(p, vec3f(-0.3511763, -0.1810853, 0.1673936) + f5 * vec3f(0.5484627, 0.5771128, -0.3130064)), dot(p, vec3f(0.1204449, -0.2777732, 0.1592967) + f5 * vec3f(-1.316489, -1.468577, -1.091297)));
779    dfg = dfg * u + vec2f(dot(p, vec3f(0.8148172, 0.1125372, -0.01374308) + f5 * vec3f(0.05596158, -0.1952609, -0.01097813)), dot(p, vec3f(0.03557167, -0.07356264, 0.03812399) + f5 * vec3f(0.1873088, 0.3210023, 0.1747016)));
780    return clamp(dfg, vec2f(0.0), vec2f(1.0));
781}
782fn reflection_environment(surface: Shaded, ndv: f32, film: vec3f, film_normal: vec3f) -> vec3f {
783    let f0 = reflection_filmed(surface, 1.0, film_normal);
784    if (all(f0 <= vec3f(0.0)) && surface.thin_film_amount <= 0.0) {
785        return vec3f(0.0);
786    }
787    let dfg = reflection_dfg(surface.roughness, ndv);
788    let fitted = clamp(f0 * (dfg.x - dfg.y) + vec3f(dfg.y), vec3f(0.0), vec3f(1.0));
789    let filmed = reflection_filmed(surface, ndv, film) - (f0 + (vec3f(1.0) - f0) * pow(1.0 - ndv, 5.0));
790    return clamp(fitted + select(vec3f(0.0), filmed, surface.thin_film_amount > 0.0), vec3f(0.0), vec3f(1.0));
791}
792fn reflection_diffuse_weight(surface: Shaded, view: vec3f, normal: vec3f, film: vec3f, film_normal: vec3f) -> vec3f {
793    let ndv = clamp(dot(normal, view), 0.0, 1.0);
794    let fresnel = reflection_environment(surface, ndv, film, film_normal);
795    let coat = clamp(surface.clearcoat, 0.0, 1.0);
796    let coat_f = schlick(vec3f(0.04), ndv, 5.0).r;
797    return (vec3f(1.0) - fresnel) * (1.0 - clamp(surface.metalness, 0.0, 1.0)) * (1.0 - clamp(surface.transmission, 0.0, 1.0)) * (1.0 - coat * coat_f);
798}
799fn reflection_fallback(
800    direction: vec3f,
801    lookup: LocalLookup,
802    rough: f32,
803    open: f32,
804    sky: texture_cube<f32>,
805    local: texture_cube_array<f32>,
806    probe_a: ReflectionProbe,
807    probe_b: ReflectionProbe,
808    samp: sampler,
809    local_mip: f32,
810    sky_mip: f32,
811) -> vec3f {
812    var fallback = vec3f(0.0);
813    if (open != 0.0) {
814        fallback = textureSampleLevel(sky, samp, direction, rough * sky_mip).rgb * open;
815    }
816    if (lookup.weights.x != 0.0) {
817        fallback += textureSampleLevel(local, samp, lookup.first, probe_a.layer, rough * local_mip).rgb * lookup.weights.x;
818    }
819    if (lookup.weights.y != 0.0) {
820        fallback += textureSampleLevel(local, samp, lookup.second, probe_b.layer, rough * local_mip).rgb * lookup.weights.y;
821    }
822    return fallback;
823}
824fn reflection_radiance(
825    surface: Shaded,
826    view: vec3f,
827    normal: vec3f,
828    hit: vec4f,
829    sky: texture_cube<f32>,
830    local: texture_cube_array<f32>,
831    probe_a: ReflectionProbe,
832    probe_b: ReflectionProbe,
833    lookup: LocalLookup,
834    samp: sampler,
835    local_mip: f32,
836    sky_mip: f32,
837    film: vec3f,
838    film_normal: vec3f,
839    sky_visible: f32,
840) -> vec3f {
841    let ndv = clamp(dot(normal, view), 0.0, 1.0);
842    let direction = reflect(-view, normal);
843    let rough = clamp(surface.roughness, 0.0, 1.0);
844    let open = (1.0 - lookup.weights.x - lookup.weights.y) * sky_visible;
845    let confidence = clamp(hit.a, 0.0, 1.0);
846    let caught = max(hit.rgb, vec3f(0.0));
847    let radiance = reflection_fallback(direction, lookup, rough, open, sky, local, probe_a, probe_b, samp, local_mip, sky_mip) * (1.0 - confidence) + caught;
848    let fresnel = reflection_environment(surface, ndv, film, film_normal);
849    let coat = clamp(surface.clearcoat, 0.0, 1.0);
850    let coat_f = schlick(vec3f(0.04), ndv, 5.0).r;
851    let coat_spec = coat * coat_f;
852    let substrate = min(fresnel * (1.0 - coat_spec), vec3f(1.0 - coat_spec));
853    if (coat_spec <= 0.0) {
854        return radiance * substrate;
855    }
856    let coat_rough = clamp(surface.clearcoat_roughness, 0.0, 1.0);
857    var coated = radiance;
858    if (coat_rough != rough) {
859        coated = reflection_fallback(direction, lookup, coat_rough, open, sky, local, probe_a, probe_b, samp, local_mip, sky_mip) * (1.0 - confidence) + caught;
860    }
861    return radiance * substrate + coated * coat_spec;
862}
863"#
864);
865
866#[repr(C)]
867#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
868struct TraceUniform {
869    size: [u32; 2],
870    half_size: [u32; 2],
871    projection: [[f32; 4]; 4],
872    inverse_projection: [[f32; 4]; 4],
873    history_valid: u32,
874    frame_index: u32,
875    refine: u32,
876    thickness: f32,
877    history_half: [u32; 2],
878    pad: [u32; 2],
879}
880
881fn sampled(binding: u32, filterable: bool) -> wgpu::BindGroupLayoutEntry {
882    wgpu::BindGroupLayoutEntry {
883        binding,
884        visibility: wgpu::ShaderStages::COMPUTE,
885        ty: wgpu::BindingType::Texture {
886            sample_type: wgpu::TextureSampleType::Float { filterable },
887            view_dimension: wgpu::TextureViewDimension::D2,
888            multisampled: false,
889        },
890        count: None,
891    }
892}
893
894fn storage(binding: u32) -> wgpu::BindGroupLayoutEntry {
895    wgpu::BindGroupLayoutEntry {
896        binding,
897        visibility: wgpu::ShaderStages::COMPUTE,
898        ty: wgpu::BindingType::StorageTexture {
899            access: wgpu::StorageTextureAccess::WriteOnly,
900            format: REFLECTION_FORMAT,
901            view_dimension: wgpu::TextureViewDimension::D2,
902        },
903        count: None,
904    }
905}
906
907fn uniform(binding: u32) -> wgpu::BindGroupLayoutEntry {
908    wgpu::BindGroupLayoutEntry {
909        binding,
910        visibility: wgpu::ShaderStages::COMPUTE,
911        ty: wgpu::BindingType::Buffer {
912            ty: wgpu::BufferBindingType::Uniform,
913            has_dynamic_offset: false,
914            min_binding_size: None,
915        },
916        count: None,
917    }
918}
919
920fn sampler_entry(binding: u32) -> wgpu::BindGroupLayoutEntry {
921    wgpu::BindGroupLayoutEntry {
922        binding,
923        visibility: wgpu::ShaderStages::COMPUTE,
924        ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
925        count: None,
926    }
927}
928
929fn texture(device: &wgpu::Device, label: &'static str, width: u32, height: u32) -> wgpu::Texture {
930    device.create_texture(&wgpu::TextureDescriptor {
931        label: Some(label),
932        size: wgpu::Extent3d {
933            width,
934            height,
935            depth_or_array_layers: 1,
936        },
937        mip_level_count: 1,
938        sample_count: 1,
939        dimension: wgpu::TextureDimension::D2,
940        format: REFLECTION_FORMAT,
941        usage: wgpu::TextureUsages::TEXTURE_BINDING
942            | wgpu::TextureUsages::STORAGE_BINDING
943            | wgpu::TextureUsages::COPY_SRC,
944        view_formats: &[],
945    })
946}
947
948pub struct Inputs<'a> {
949    pub depth: &'a wgpu::TextureView,
950    pub normal_roughness: &'a wgpu::TextureView,
951    pub colour: &'a wgpu::TextureView,
952    pub motion: &'a wgpu::TextureView,
953    pub projection: [[f32; 4]; 4],
954    pub inverse_projection: [[f32; 4]; 4],
955    pub frame_index: u32,
956    pub thickness: f32,
957    pub history_valid: bool,
958}
959
960#[derive(Clone, Debug)]
961pub struct Reflections {
962    pub colour: wgpu::TextureView,
963    pub guide: wgpu::TextureView,
964    pub size: [u32; 2],
965}
966
967pub struct Pass {
968    width: u32,
969    height: u32,
970    viewport: (u32, u32),
971    history: [u32; 2],
972    _raw: [wgpu::Texture; 2],
973    raw_view: wgpu::TextureView,
974    raw_guide: wgpu::TextureView,
975    _results: [wgpu::Texture; 2],
976    result_views: [wgpu::TextureView; 2],
977    _guides: [wgpu::Texture; 2],
978    guide_views: [wgpu::TextureView; 2],
979    sampler: wgpu::Sampler,
980    next: usize,
981    valid: bool,
982    trace_layout: wgpu::BindGroupLayout,
983    resolve_layout: wgpu::BindGroupLayout,
984    trace_pipeline: wgpu::ComputePipeline,
985    resolve_pipeline: wgpu::ComputePipeline,
986}
987
988impl Pass {
989    pub fn new(device: &wgpu::Device, width: u32, height: u32) -> Self {
990        assert!(width > 0 && height > 0);
991        let trace_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
992            label: Some("reflection trace"),
993            entries: &[
994                uniform(0),
995                sampled(1, false),
996                sampled(2, false),
997                sampled(3, true),
998                sampler_entry(4),
999                storage(5),
1000                storage(6),
1001            ],
1002        });
1003        let resolve_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
1004            label: Some("reflection resolve"),
1005            entries: &[
1006                uniform(0),
1007                sampled(1, false),
1008                sampled(2, false),
1009                sampled(3, false),
1010                sampled(4, false),
1011                sampled(5, false),
1012                storage(6),
1013                storage(7),
1014            ],
1015        });
1016        let pipeline = |label, source: &str, layout: &wgpu::BindGroupLayout, entry| {
1017            let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
1018                label: Some(label),
1019                source: wgpu::ShaderSource::Wgsl(source.into()),
1020            });
1021            let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
1022                label: Some(label),
1023                bind_group_layouts: &[Some(layout)],
1024                immediate_size: 0,
1025            });
1026            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
1027                label: Some(label),
1028                layout: Some(&layout),
1029                module: &module,
1030                entry_point: Some(entry),
1031                compilation_options: Default::default(),
1032                cache: None,
1033            })
1034        };
1035        let trace_pipeline = pipeline("reflection trace", TRACE_WGSL, &trace_layout, "trace_pass");
1036        let resolve_pipeline = pipeline(
1037            "reflection resolve",
1038            RESOLVE_WGSL,
1039            &resolve_layout,
1040            "resolve_pass",
1041        );
1042        let (half_width, half_height) = half_size(width, height);
1043        let raw = [
1044            texture(device, "reflection raw", half_width, half_height),
1045            texture(device, "reflection raw guide", half_width, half_height),
1046        ];
1047        let raw_view = raw[0].create_view(&Default::default());
1048        let raw_guide = raw[1].create_view(&Default::default());
1049        let results = [
1050            texture(device, "reflection result 0", half_width, half_height),
1051            texture(device, "reflection result 1", half_width, half_height),
1052        ];
1053        let result_views = [
1054            results[0].create_view(&Default::default()),
1055            results[1].create_view(&Default::default()),
1056        ];
1057        let guides = [
1058            texture(device, "reflection guide 0", half_width, half_height),
1059            texture(device, "reflection guide 1", half_width, half_height),
1060        ];
1061        let guide_views = [
1062            guides[0].create_view(&Default::default()),
1063            guides[1].create_view(&Default::default()),
1064        ];
1065        let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
1066            label: Some("reflection colour"),
1067            address_mode_u: wgpu::AddressMode::ClampToEdge,
1068            address_mode_v: wgpu::AddressMode::ClampToEdge,
1069            mag_filter: wgpu::FilterMode::Linear,
1070            min_filter: wgpu::FilterMode::Linear,
1071            mipmap_filter: wgpu::MipmapFilterMode::Linear,
1072            ..Default::default()
1073        });
1074        Self {
1075            width,
1076            height,
1077            viewport: (width, height),
1078            history: [half_width, half_height],
1079            _raw: raw,
1080            raw_view,
1081            raw_guide,
1082            _results: results,
1083            result_views,
1084            _guides: guides,
1085            guide_views,
1086            sampler,
1087            next: 0,
1088            valid: false,
1089            trace_layout,
1090            resolve_layout,
1091            trace_pipeline,
1092            resolve_pipeline,
1093        }
1094    }
1095
1096    pub fn size(&self) -> (u32, u32) {
1097        (self.width, self.height)
1098    }
1099
1100    pub fn resize(&mut self, device: &wgpu::Device, width: u32, height: u32) {
1101        if self.size() != (width, height) {
1102            *self = Self::new(device, width, height);
1103        }
1104    }
1105
1106    pub fn viewport(&self) -> (u32, u32) {
1107        self.viewport
1108    }
1109
1110    pub fn textures(&self) -> Vec<wgpu::Texture> {
1111        self._raw
1112            .iter()
1113            .chain(&self._results)
1114            .chain(&self._guides)
1115            .cloned()
1116            .collect()
1117    }
1118
1119    pub fn set_viewport(&mut self, width: u32, height: u32) -> Result<(), String> {
1120        if width == 0 || height == 0 || width > self.width || height > self.height {
1121            return Err(format!(
1122                "the viewport {width}x{height} must be inside the reflections' {}x{}",
1123                self.width, self.height
1124            ));
1125        }
1126        self.viewport = (width, height);
1127        Ok(())
1128    }
1129
1130    pub fn result(&self) -> Option<Reflections> {
1131        let last = 1 - self.next;
1132        self.valid.then(|| Reflections {
1133            colour: self.result_views[last].clone(),
1134            guide: self.guide_views[last].clone(),
1135            size: self.history,
1136        })
1137    }
1138
1139    pub fn encode(
1140        &mut self,
1141        device: &wgpu::Device,
1142        encoder: &mut wgpu::CommandEncoder,
1143        inputs: &Inputs<'_>,
1144        mut profiler: Option<&mut GpuProfiler>,
1145    ) {
1146        use wgpu::util::DeviceExt;
1147        let (width, height) = self.viewport;
1148        let (half_width, half_height) = half_size(width, height);
1149        let params = TraceUniform {
1150            size: [width, height],
1151            half_size: [half_width, half_height],
1152            projection: inputs.projection,
1153            inverse_projection: inputs.inverse_projection,
1154            history_valid: u32::from(self.valid && inputs.history_valid),
1155            frame_index: inputs.frame_index,
1156            refine: REFINE_STEPS,
1157            thickness: inputs.thickness,
1158            history_half: self.history,
1159            pad: [0; 2],
1160        };
1161        let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
1162            label: Some("reflection uniforms"),
1163            contents: bytemuck::bytes_of(&params),
1164            usage: wgpu::BufferUsages::UNIFORM,
1165        });
1166        let current = self.next;
1167        let last = 1 - current;
1168        let view = wgpu::BindingResource::TextureView;
1169        let bind =
1170            |layout: &wgpu::BindGroupLayout, label, resources: Vec<wgpu::BindingResource>| {
1171                let entries = std::iter::once(buffer.as_entire_binding())
1172                    .chain(resources)
1173                    .enumerate()
1174                    .map(|(binding, resource)| wgpu::BindGroupEntry {
1175                        binding: binding as u32,
1176                        resource,
1177                    })
1178                    .collect::<Vec<_>>();
1179                device.create_bind_group(&wgpu::BindGroupDescriptor {
1180                    label: Some(label),
1181                    layout,
1182                    entries: &entries,
1183                })
1184            };
1185        let trace_bind = bind(
1186            &self.trace_layout,
1187            "reflection trace bindings",
1188            vec![
1189                view(inputs.depth),
1190                view(inputs.normal_roughness),
1191                view(inputs.colour),
1192                wgpu::BindingResource::Sampler(&self.sampler),
1193                view(&self.raw_view),
1194                view(&self.raw_guide),
1195            ],
1196        );
1197        let resolve_bind = bind(
1198            &self.resolve_layout,
1199            "reflection resolve bindings",
1200            vec![
1201                view(&self.raw_view),
1202                view(&self.raw_guide),
1203                view(&self.result_views[last]),
1204                view(&self.guide_views[last]),
1205                view(inputs.motion),
1206                view(&self.result_views[current]),
1207                view(&self.guide_views[current]),
1208            ],
1209        );
1210        let groups = [half_width.div_ceil(8), half_height.div_ceil(8)];
1211        for (label, pipeline, group) in [
1212            ("reflections", &self.trace_pipeline, &trace_bind),
1213            ("reflection resolve", &self.resolve_pipeline, &resolve_bind),
1214        ] {
1215            let timing = profiler.as_deref_mut().and_then(|p| p.pass(label));
1216            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
1217                label: Some(label),
1218                timestamp_writes: profiler.as_deref().and_then(|p| p.compute_writes(timing)),
1219            });
1220            pass.set_pipeline(pipeline);
1221            pass.set_bind_group(0, group, &[]);
1222            pass.dispatch_workgroups(groups[0], groups[1], 1);
1223        }
1224        self.next = last;
1225        self.valid = true;
1226        self.history = [half_width, half_height];
1227    }
1228}
1229
1230#[cfg(test)]
1231mod tests {
1232    use super::*;
1233
1234    const FAR: f32 = 80.0;
1235
1236    fn traced_dfg(roughness: f32, ndv: f32) -> [f32; 2] {
1237        let a2 = pfx_materials::alpha(roughness.max(0.05)).powi(2);
1238        let view = [(1.0 - ndv * ndv).sqrt(), 0.0, ndv];
1239        let count = 4096u32;
1240        let mut sum = [0.0f64; 2];
1241        for k in 0..count {
1242            let u1 = (k as f32 + 0.5) / count as f32;
1243            let u2 = k.reverse_bits() as f32 / 4_294_967_296.0;
1244            let theta = (a2.sqrt() * (u1 / (1.0 - u1)).sqrt()).atan();
1245            let phi = std::f32::consts::TAU * u2;
1246            let h = [
1247                theta.sin() * phi.cos(),
1248                theta.sin() * phi.sin(),
1249                theta.cos(),
1250            ];
1251            let vh = view[0] * h[0] + view[2] * h[2];
1252            let ndl = 2.0 * vh * h[2] - view[2];
1253            if ndl <= 0.0 || vh <= 0.0 {
1254                continue;
1255            }
1256            let g = pfx_materials::smith_g(ndv, ndl, a2);
1257            let weight = f64::from(g * vh / (h[2] * ndv));
1258            let fresnel = f64::from((1.0 - vh).powi(5));
1259            sum[0] += weight;
1260            sum[1] += weight * fresnel;
1261        }
1262        sum.map(|v| (v / f64::from(count)) as f32)
1263    }
1264
1265    #[test]
1266    fn the_environment_brdf_follows_the_tracers_ggx() {
1267        let mut worst = [0.0f32; 2];
1268        for i in 0..20 {
1269            for j in 0..20 {
1270                let roughness = 0.05 + 0.95 * (i as f32 + 0.37) / 20.0;
1271                let ndv = 0.1 + 0.9 * (j as f32 + 0.61) / 20.0;
1272                let truth = traced_dfg(roughness, ndv);
1273                let fit = environment_dfg(roughness, ndv);
1274                let metal = (fit[0] - truth[0]).abs();
1275                let dielectric = (0.04 * (fit[0] - fit[1]) + fit[1]
1276                    - (0.04 * (truth[0] - truth[1]) + truth[1]))
1277                    .abs();
1278                worst = [worst[0].max(metal), worst[1].max(dielectric)];
1279            }
1280        }
1281        assert!(worst[0] < 0.02, "metal off by {}", worst[0]);
1282        assert!(worst[1] < 0.012, "dielectric off by {}", worst[1]);
1283    }
1284
1285    #[test]
1286    fn a_rough_metal_keeps_the_light_the_old_fit_lost() {
1287        let truth = traced_dfg(0.35, 0.75);
1288        let fit = environment_dfg(0.35, 0.75);
1289        assert!(
1290            (fit[0] - truth[0]).abs() < 0.01,
1291            "{fit:?} against {truth:?}"
1292        );
1293        assert!(truth[0] > 0.95, "{truth:?}");
1294    }
1295
1296    #[test]
1297    fn the_shader_carries_the_same_fit() {
1298        for row in DFG_FIT {
1299            let text = format!("vec3f({}, {}, {})", row[0], row[1], row[2]);
1300            assert!(PROBE_WGSL.contains(&text), "{text} is not in the shader");
1301        }
1302    }
1303
1304    fn perspective(fovy: f32, aspect: f32, near: f32, far: f32) -> Matrix {
1305        let b = 1.0 / (fovy / 2.0).tan();
1306        let a = b / aspect;
1307        let c = far / (near - far);
1308        let d = far * near / (near - far);
1309        [
1310            [a, 0.0, 0.0, 0.0],
1311            [0.0, b, 0.0, 0.0],
1312            [0.0, 0.0, c, -1.0],
1313            [0.0, 0.0, d, 0.0],
1314        ]
1315    }
1316
1317    fn inverse_perspective(m: &Matrix) -> Matrix {
1318        let (a, b, c, d) = (m[0][0], m[1][1], m[2][2], m[3][2]);
1319        [
1320            [1.0 / a, 0.0, 0.0, 0.0],
1321            [0.0, 1.0 / b, 0.0, 0.0],
1322            [0.0, 0.0, 0.0, 1.0 / d],
1323            [0.0, 0.0, -1.0, c / d],
1324        ]
1325    }
1326
1327    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
1328    enum Surface {
1329        Sky,
1330        Floor,
1331        Card,
1332        Wall,
1333    }
1334
1335    struct Desk {
1336        gap: f32,
1337        width: u32,
1338        height: u32,
1339        projection: Matrix,
1340        pitch: f32,
1341        eye: [f32; 3],
1342        depth: Vec<f32>,
1343        surface: Vec<Surface>,
1344        position: Vec<[f32; 3]>,
1345    }
1346
1347    const CARD_X: [f32; 2] = [-0.25, 0.06];
1348    const CARD_TOP: f32 = 0.4;
1349    const WALL_Z: f32 = -0.6;
1350
1351    fn to_view(pitch: f32, eye: [f32; 3], world: [f32; 3]) -> [f32; 3] {
1352        let p = [world[0] - eye[0], world[1] - eye[1], world[2] - eye[2]];
1353        let (s, c) = pitch.sin_cos();
1354        [p[0], p[1] * c - p[2] * s, p[1] * s + p[2] * c]
1355    }
1356
1357    fn to_world(pitch: f32, v: [f32; 3]) -> [f32; 3] {
1358        let (s, c) = pitch.sin_cos();
1359        [v[0], v[1] * c + v[2] * s, -v[1] * s + v[2] * c]
1360    }
1361
1362    fn cast(gap: f32, origin: [f32; 3], direction: [f32; 3]) -> Option<(f32, Surface)> {
1363        let mut best: Option<(f32, Surface)> = None;
1364        let mut keep = |t: f32, surface| {
1365            if t > 1e-5 && best.is_none_or(|(b, _)| t < b) {
1366                best = Some((t, surface));
1367            }
1368        };
1369        if direction[1] < 0.0 {
1370            keep(-origin[1] / direction[1], Surface::Floor);
1371        }
1372        if direction[2] != 0.0 {
1373            let t = -origin[2] / direction[2];
1374            let x = origin[0] + direction[0] * t;
1375            let y = origin[1] + direction[1] * t;
1376            if (CARD_X[0]..CARD_X[1]).contains(&x) && (gap..CARD_TOP).contains(&y) {
1377                keep(t, Surface::Card);
1378            }
1379            keep((WALL_Z - origin[2]) / direction[2], Surface::Wall);
1380        }
1381        best
1382    }
1383
1384    fn pixel_ray(desk: &Desk, x: f32, y: f32) -> [f32; 3] {
1385        let ndc = [
1386            x / desk.width as f32 * 2.0 - 1.0,
1387            1.0 - y / desk.height as f32 * 2.0,
1388        ];
1389        let v = [
1390            ndc[0] / desk.projection[0][0],
1391            ndc[1] / desk.projection[1][1],
1392            -1.0,
1393        ];
1394        let length = dot(v, v).sqrt();
1395        v.map(|value| value / length)
1396    }
1397
1398    fn desk(width: u32, height: u32) -> Desk {
1399        raised_desk(width, height, 0.03)
1400    }
1401
1402    fn raised_desk(width: u32, height: u32, gap: f32) -> Desk {
1403        let projection = perspective(0.75, width as f32 / height as f32, 0.05, FAR);
1404        let pitch = 0.32f32;
1405        let eye = [-0.04, 0.32, 1.0];
1406        let mut d = Desk {
1407            gap,
1408            width,
1409            height,
1410            projection,
1411            pitch,
1412            eye,
1413            depth: Vec::new(),
1414            surface: Vec::new(),
1415            position: Vec::new(),
1416        };
1417        for y in 0..height {
1418            for x in 0..width {
1419                let ray = pixel_ray(&d, x as f32 + 0.5, y as f32 + 0.5);
1420                match cast(gap, eye, to_world(pitch, ray)) {
1421                    Some((t, surface)) => {
1422                        let view = ray.map(|value| value * t);
1423                        d.depth.push(-view[2]);
1424                        d.surface.push(surface);
1425                        d.position.push(view);
1426                    }
1427                    None => {
1428                        d.depth.push(FAR);
1429                        d.surface.push(Surface::Sky);
1430                        d.position.push([0.0, 0.0, -FAR]);
1431                    }
1432                }
1433            }
1434        }
1435        d
1436    }
1437
1438    fn floor_normal(pitch: f32) -> [f32; 3] {
1439        let n = to_view(pitch, [0.0; 3], [0.0, 1.0, 0.0]);
1440        let length = dot(n, n).sqrt();
1441        n.map(|value| value / length)
1442    }
1443
1444    fn reflect(v: [f32; 3], n: [f32; 3]) -> [f32; 3] {
1445        let d = 2.0 * dot(v, n);
1446        [v[0] - d * n[0], v[1] - d * n[1], v[2] - d * n[2]]
1447    }
1448
1449    fn reflects_card(desk: &Desk, index: usize) -> bool {
1450        let origin = desk.position[index];
1451        let length = dot(origin, origin).sqrt();
1452        let n = floor_normal(desk.pitch);
1453        let direction = reflect(origin.map(|v| v / length), n);
1454        let image = DepthImage {
1455            width: desk.width,
1456            height: desk.height,
1457            depth: &desk.depth,
1458        };
1459        let depth = -origin[2];
1460        let start = std::array::from_fn(|axis| origin[axis] + n[axis] * depth * 0.002);
1461        march(
1462            &image,
1463            &desk.projection,
1464            start,
1465            direction,
1466            march_steps(0.05),
1467            THICKNESS,
1468        )
1469        .is_some_and(|hit| {
1470            let x = hit.pixel[0] as usize;
1471            let y = hit.pixel[1] as usize;
1472            desk.surface[y * desk.width as usize + x] == Surface::Card
1473        })
1474    }
1475
1476    fn truly_reflects_card(desk: &Desk, index: usize) -> bool {
1477        let origin = desk.position[index];
1478        let length = dot(origin, origin).sqrt();
1479        let direction = reflect(origin.map(|v| v / length), floor_normal(desk.pitch));
1480        let world = to_world(desk.pitch, origin);
1481        let world = std::array::from_fn(|axis| world[axis] + desk.eye[axis]);
1482        match cast(desk.gap, world, to_world(desk.pitch, direction)) {
1483            Some((t, Surface::Card)) => {
1484                let hit = std::array::from_fn(|axis| origin[axis] + direction[axis] * t);
1485                let size = [desk.width as f32, desk.height as f32];
1486                let p = to_pixel(&desk.projection, size, hit);
1487                (0.0..size[0]).contains(&p[0]) && (0.0..size[1]).contains(&p[1])
1488            }
1489            _ => false,
1490        }
1491    }
1492
1493    fn boundaries(
1494        desk: &Desk,
1495        card: impl Fn(usize) -> bool,
1496        rows: std::ops::Range<u32>,
1497    ) -> Vec<(u32, f32)> {
1498        let mut found = Vec::new();
1499        for y in rows {
1500            let row = (y * desk.width) as usize;
1501            let mut edge = None;
1502            for x in 1..desk.width as usize {
1503                let a = row + x - 1;
1504                let b = row + x;
1505                if desk.surface[a] != Surface::Floor || desk.surface[b] != Surface::Floor {
1506                    continue;
1507                }
1508                if card(a) && !card(b) {
1509                    edge = Some(x as f32);
1510                }
1511            }
1512            if let Some(x) = edge {
1513                found.push((y, x));
1514            }
1515        }
1516        found
1517    }
1518
1519    #[test]
1520    fn reflected_vertical_edge_is_monotonic_and_step_free() {
1521        let desk = desk(640, 360);
1522        let rows = 190..330;
1523        let traced = boundaries(&desk, |i| reflects_card(&desk, i), rows.clone());
1524        let expected = boundaries(&desk, |i| truly_reflects_card(&desk, i), rows);
1525        assert!(traced.len() > 100, "{} rows reflect the edge", traced.len());
1526        let truth_at = expected
1527            .iter()
1528            .copied()
1529            .collect::<std::collections::BTreeMap<_, _>>();
1530        let direction = (traced.last().unwrap().1 - traced[0].1).signum();
1531        for pair in traced.windows(2) {
1532            let (y0, x0) = pair[0];
1533            let (y1, x1) = pair[1];
1534            assert_eq!(
1535                y1,
1536                y0 + 1,
1537                "the reflected edge breaks between rows {y0} and {y1}"
1538            );
1539            let step = x1 - x0;
1540            assert!(step.abs() <= 1.0, "a stair-step of {step} px at row {y1}");
1541        }
1542        for (i, &(y0, x0)) in traced.iter().enumerate() {
1543            for &(y1, x1) in &traced[i + 1..] {
1544                assert!(
1545                    (x1 - x0) * direction >= -1.0,
1546                    "the reflected edge turns back by more than a pixel: {x0} at row {y0}, {x1} at row {y1}"
1547                );
1548            }
1549        }
1550        for (y, x) in &traced {
1551            let want = truth_at[y];
1552            assert!(
1553                (x - want).abs() <= 1.5,
1554                "row {y}: the reflected edge at {x}, the true edge at {want}"
1555            );
1556        }
1557    }
1558
1559    #[test]
1560    fn floor_behind_a_low_card_never_reflects_it() {
1561        for gap in [0.03, 0.008, 0.004] {
1562            let desk = raised_desk(640, 360, gap);
1563            let mut behind = 0;
1564            for index in 0..desk.surface.len() {
1565                if desk.surface[index] != Surface::Floor {
1566                    continue;
1567                }
1568                let world = to_world(desk.pitch, desk.position[index]);
1569                if world[2] + desk.eye[2] >= 0.0 {
1570                    continue;
1571                }
1572                behind += 1;
1573                assert!(
1574                    !reflects_card(&desk, index),
1575                    "floor pixel {} behind a card {gap} m off the floor reflects the card's face",
1576                    index
1577                );
1578            }
1579            assert!(behind > 10000, "{behind} floor pixels behind the card");
1580        }
1581    }
1582
1583    #[test]
1584    fn reflected_points_land_within_a_pixel_of_the_truth() {
1585        let desk = desk(640, 360);
1586        let size = [desk.width as f32, desk.height as f32];
1587        let n = floor_normal(desk.pitch);
1588        let image = DepthImage {
1589            width: desk.width,
1590            height: desk.height,
1591            depth: &desk.depth,
1592        };
1593        let mut checked = 0;
1594        let mut missed = 0;
1595        let mut worst = 0.0f32;
1596        for y in (180..330).step_by(3) {
1597            for x in (200..400).step_by(3) {
1598                let index = y * desk.width as usize + x;
1599                if desk.surface[index] != Surface::Floor || !truly_reflects_card(&desk, index) {
1600                    continue;
1601                }
1602                let origin = desk.position[index];
1603                let length = dot(origin, origin).sqrt();
1604                let direction = reflect(origin.map(|v| v / length), n);
1605                let world = to_world(desk.pitch, origin);
1606                let world = std::array::from_fn(|axis| world[axis] + desk.eye[axis]);
1607                let (t, _) = cast(desk.gap, world, to_world(desk.pitch, direction)).unwrap();
1608                let want = to_pixel(
1609                    &desk.projection,
1610                    size,
1611                    std::array::from_fn(|axis| origin[axis] + direction[axis] * t),
1612                );
1613                let start = std::array::from_fn(|axis| origin[axis] - origin[2] * 0.002 * n[axis]);
1614                match march(
1615                    &image,
1616                    &desk.projection,
1617                    start,
1618                    direction,
1619                    march_steps(0.05),
1620                    THICKNESS,
1621                ) {
1622                    Some(hit) => {
1623                        let error = (hit.pixel[0] - want[0]).hypot(hit.pixel[1] - want[1]);
1624                        worst = worst.max(error);
1625                        checked += 1;
1626                    }
1627                    None => missed += 1,
1628                }
1629            }
1630        }
1631        assert!(checked > 1000, "{checked} reflected points checked");
1632        assert!(
1633            missed * 50 < checked,
1634            "{missed} of {checked} reflected points missed"
1635        );
1636        assert!(worst <= 1.0, "a reflected point {worst} px from the truth");
1637    }
1638
1639    #[test]
1640    fn reflected_lower_edge_holds_on_every_column() {
1641        let desk = desk(640, 360);
1642        let width = desk.width as usize;
1643        let mut traced = Vec::new();
1644        let rise = |card: &dyn Fn(usize) -> bool, x: usize| {
1645            let mut seen = false;
1646            for y in (0..desk.height as usize).rev() {
1647                let index = y * width + x;
1648                if desk.surface[index] != Surface::Floor {
1649                    return None;
1650                }
1651                if card(index) {
1652                    seen = true;
1653                } else if seen {
1654                    return Some(y as f32);
1655                }
1656            }
1657            None
1658        };
1659        for x in 240..350usize {
1660            let first_wall = rise(&|i| reflects_card(&desk, i), x);
1661            let true_wall = rise(&|i| truly_reflects_card(&desk, i), x);
1662            if let (Some(found), Some(want)) = (first_wall, true_wall) {
1663                assert!(
1664                    (found - want).abs() <= 1.5,
1665                    "column {x}: the reflected lower edge at row {found}, the true edge at {want}"
1666                );
1667                traced.push(found);
1668            }
1669        }
1670        assert!(
1671            traced.len() > 80,
1672            "{} columns reflect the lower edge",
1673            traced.len()
1674        );
1675        for pair in traced.windows(2) {
1676            assert!(
1677                (pair[1] - pair[0]).abs() <= 1.0,
1678                "a stair-step along the lower edge: {} then {}",
1679                pair[0],
1680                pair[1]
1681            );
1682        }
1683    }
1684
1685    #[test]
1686    fn upsample_does_not_bleed_across_a_depth_edge() {
1687        let (width, height) = (8u32, 4u32);
1688        let up = [0.0, 0.0, 1.0];
1689        let near = Texel {
1690            colour: [1.0, 0.0, 0.0, 1.0],
1691            depth: 1.0,
1692            normal: up,
1693        };
1694        let far = Texel {
1695            colour: [0.0, 0.0, 1.0, 1.0],
1696            depth: 4.0,
1697            normal: up,
1698        };
1699        let texels = (0..width * height)
1700            .map(|i| if i % width < 4 { near } else { far })
1701            .collect::<Vec<_>>();
1702        for step in 0..=40 {
1703            let x = 0.25 * step as f32 / 40.0;
1704            for (depth, keep, lose, u) in [(4.0, 2, 0, 0.5 + x), (1.0, 0, 2, 0.4999 - x)] {
1705                let value = upsample(width, height, &texels, [u, 0.5], depth, up);
1706                assert!(
1707                    value[lose] == 0.0 && (value[keep] - 1.0).abs() < 1e-6,
1708                    "at u {u} and depth {depth}: {value:?}"
1709                );
1710            }
1711        }
1712        let side = Texel {
1713            colour: [0.0, 1.0, 0.0, 1.0],
1714            depth: 4.0,
1715            normal: [1.0, 0.0, 0.0],
1716        };
1717        let crease = (0..width * height)
1718            .map(|i| if i % width < 4 { side } else { far })
1719            .collect::<Vec<_>>();
1720        let value = upsample(width, height, &crease, [0.5, 0.5], 4.0, up);
1721        assert_eq!(
1722            value,
1723            [0.0, 0.0, 1.0, 1.0],
1724            "a crease at one depth keeps its faces apart"
1725        );
1726        let lone = upsample(width, height, &texels, [0.3, 0.5], 9.0, up);
1727        assert_eq!(lone, [0.0; 4], "nothing near falls back to the probes");
1728        let bilinear = upsample(width, height, &texels, [0.7, 0.5], 4.0, up);
1729        assert_eq!(bilinear, [0.0, 0.0, 1.0, 1.0]);
1730    }
1731
1732    #[test]
1733    fn history_resets_on_a_cut() {
1734        use crate::renderer::History;
1735        let a = perspective(0.75, 16.0 / 9.0, 0.05, FAR);
1736        let mut b = a;
1737        b[3][0] = 0.4;
1738        let current = [0.2, 0.3, 0.4, 0.5];
1739        let old = [0.25, 0.35, 0.45, 0.55];
1740        let low = [0.1, 0.1, 0.1, 0.1];
1741        let high = [0.6, 0.6, 0.6, 0.6];
1742        let mut history = History::default();
1743        let mut frame = |n: f32, previous, now| {
1744            let valid = history.advance(n / 60.0, previous, now);
1745            (valid, accumulate(current, low, high, old, 1.0, 0.1, valid))
1746        };
1747        assert_eq!(frame(0.0, a, a), (false, current));
1748        let (valid, kept) = frame(1.0, a, a);
1749        assert!(
1750            valid && kept != current,
1751            "a continuing camera keeps history"
1752        );
1753        assert_eq!(frame(2.0, b, b), (false, current), "a cut drops history");
1754        let (valid, kept) = frame(3.0, b, b);
1755        assert!(valid && kept != current);
1756        let ghost = accumulate(current, low, high, [5.0, 5.0, 5.0, 1.0], 1.0, 0.1, true);
1757        assert!(
1758            ghost
1759                .iter()
1760                .zip(high)
1761                .all(|(value, top)| *value <= top + 1e-6),
1762            "history outside this frame's neighbourhood is clamped: {ghost:?}"
1763        );
1764        assert_eq!(
1765            accumulate(current, low, high, old, 0.0, 0.1, true),
1766            current,
1767            "a disoccluded texel has no history"
1768        );
1769    }
1770
1771    #[test]
1772    fn steps_shrink_with_roughness_and_offsets_cover_the_quad() {
1773        assert!(march_steps(0.05) > march_steps(0.3));
1774        assert!(march_steps(0.3) > march_steps(0.8));
1775        assert_eq!(rough_fade(0.3), 1.0);
1776        assert_eq!(rough_fade(0.8), 0.0);
1777        assert!((0.0..1.0).contains(&rough_fade(0.7)));
1778        let mut seen = (0..4).map(trace_offset).collect::<Vec<_>>();
1779        seen.sort();
1780        assert_eq!(seen, vec![[0, 0], [0, 1], [1, 0], [1, 1]]);
1781    }
1782
1783    #[test]
1784    fn shaders_validate() {
1785        for source in [
1786            TRACE_WGSL.to_string(),
1787            RESOLVE_WGSL.to_string(),
1788            format!("{}\n{}", pfx_materials::BRDF, PROBE_WGSL),
1789        ] {
1790            let module = naga::front::wgsl::parse_str(&source).unwrap();
1791            naga::valid::Validator::new(
1792                naga::valid::ValidationFlags::all(),
1793                naga::valid::Capabilities::all(),
1794            )
1795            .validate(&module)
1796            .unwrap();
1797        }
1798    }
1799
1800    #[test]
1801    fn confidence_fades_at_each_edge() {
1802        assert_eq!(edge_confidence([0.0, 0.5], 0.1), 0.0);
1803        assert_eq!(edge_confidence([1.0, 0.5], 0.1), 0.0);
1804        assert_eq!(edge_confidence([0.5, 0.0], 0.1), 0.0);
1805        assert_eq!(edge_confidence([0.5, 1.0], 0.1), 0.0);
1806        assert!((edge_confidence([0.05, 0.5], 0.1) - 0.5).abs() < 1e-6);
1807        assert_eq!(edge_confidence([0.5, 0.5], 0.1), 1.0);
1808    }
1809
1810    struct Views {
1811        depth: wgpu::TextureView,
1812        normal: wgpu::TextureView,
1813        colour: wgpu::TextureView,
1814        motion: wgpu::TextureView,
1815    }
1816
1817    fn upload(
1818        gpu: &pfx_gpu::Gpu,
1819        width: u32,
1820        height: u32,
1821        format: wgpu::TextureFormat,
1822        bytes: &[u8],
1823    ) -> wgpu::TextureView {
1824        let texture = gpu.device.create_texture(&wgpu::TextureDescriptor {
1825            label: Some("reflection test input"),
1826            size: wgpu::Extent3d {
1827                width,
1828                height,
1829                depth_or_array_layers: 1,
1830            },
1831            mip_level_count: 1,
1832            sample_count: 1,
1833            dimension: wgpu::TextureDimension::D2,
1834            format,
1835            usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
1836            view_formats: &[],
1837        });
1838        let stride = bytes.len() as u32 / (width * height);
1839        gpu.queue.write_texture(
1840            texture.as_image_copy(),
1841            bytes,
1842            wgpu::TexelCopyBufferLayout {
1843                offset: 0,
1844                bytes_per_row: Some(width * stride),
1845                rows_per_image: Some(height),
1846            },
1847            wgpu::Extent3d {
1848                width,
1849                height,
1850                depth_or_array_layers: 1,
1851            },
1852        );
1853        texture.create_view(&Default::default())
1854    }
1855
1856    fn halves(values: impl Iterator<Item = f32>) -> Vec<u8> {
1857        values
1858            .flat_map(|v| half::f16::from_f32(v).to_bits().to_le_bytes())
1859            .collect()
1860    }
1861
1862    fn views(gpu: &pfx_gpu::Gpu, desk: &Desk, card: [f32; 3]) -> Views {
1863        let (w, h) = (desk.width, desk.height);
1864        let depth_bytes = desk
1865            .depth
1866            .iter()
1867            .flat_map(|d| d.to_le_bytes())
1868            .collect::<Vec<_>>();
1869        let up = floor_normal(desk.pitch);
1870        let normals = desk.surface.iter().flat_map(|surface| match surface {
1871            Surface::Floor => [up[0], up[1], up[2], 0.05],
1872            Surface::Sky => [0.0; 4],
1873            _ => [0.0, 0.0, 1.0, 0.6],
1874        });
1875        let colours = desk.surface.iter().flat_map(|surface| match surface {
1876            Surface::Card => [card[0], card[1], card[2], 1.0],
1877            Surface::Wall => [0.0, 0.0, 1.0, 1.0],
1878            Surface::Floor | Surface::Sky => [0.0, 0.0, 0.0, 1.0],
1879        });
1880        Views {
1881            depth: upload(gpu, w, h, wgpu::TextureFormat::R32Float, &depth_bytes),
1882            normal: upload(
1883                gpu,
1884                w,
1885                h,
1886                wgpu::TextureFormat::Rgba16Float,
1887                &halves(normals),
1888            ),
1889            colour: upload(
1890                gpu,
1891                w,
1892                h,
1893                wgpu::TextureFormat::Rgba16Float,
1894                &halves(colours),
1895            ),
1896            motion: upload(
1897                gpu,
1898                w,
1899                h,
1900                wgpu::TextureFormat::Rg16Float,
1901                &vec![0u8; (w * h * 4) as usize],
1902            ),
1903        }
1904    }
1905
1906    fn inputs<'a>(
1907        views: &'a Views,
1908        desk: &Desk,
1909        frame_index: u32,
1910        history_valid: bool,
1911    ) -> Inputs<'a> {
1912        Inputs {
1913            depth: &views.depth,
1914            normal_roughness: &views.normal,
1915            colour: &views.colour,
1916            motion: &views.motion,
1917            projection: desk.projection,
1918            inverse_projection: inverse_perspective(&desk.projection),
1919            frame_index,
1920            thickness: THICKNESS,
1921            history_valid,
1922        }
1923    }
1924
1925    fn run(
1926        gpu: &pfx_gpu::Gpu,
1927        pass: &mut Pass,
1928        views: &Views,
1929        desk: &Desk,
1930        frame_index: u32,
1931        history_valid: bool,
1932    ) {
1933        let mut encoder = gpu.device.create_command_encoder(&Default::default());
1934        pass.encode(
1935            &gpu.device,
1936            &mut encoder,
1937            &inputs(views, desk, frame_index, history_valid),
1938            None,
1939        );
1940        gpu.queue.submit(Some(encoder.finish()));
1941    }
1942
1943    fn read(gpu: &pfx_gpu::Gpu, texture: &wgpu::Texture) -> Vec<[f32; 4]> {
1944        let (width, height) = (texture.width(), texture.height());
1945        let row = (width * 8).div_ceil(256) * 256;
1946        let buffer = gpu.device.create_buffer(&wgpu::BufferDescriptor {
1947            label: Some("reflection test readback"),
1948            size: u64::from(row) * u64::from(height),
1949            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
1950            mapped_at_creation: false,
1951        });
1952        let mut encoder = gpu.device.create_command_encoder(&Default::default());
1953        encoder.copy_texture_to_buffer(
1954            texture.as_image_copy(),
1955            wgpu::TexelCopyBufferInfo {
1956                buffer: &buffer,
1957                layout: wgpu::TexelCopyBufferLayout {
1958                    offset: 0,
1959                    bytes_per_row: Some(row),
1960                    rows_per_image: Some(height),
1961                },
1962            },
1963            wgpu::Extent3d {
1964                width,
1965                height,
1966                depth_or_array_layers: 1,
1967            },
1968        );
1969        gpu.queue.submit(Some(encoder.finish()));
1970        let (send, receive) = std::sync::mpsc::channel();
1971        buffer
1972            .slice(..)
1973            .map_async(wgpu::MapMode::Read, move |result| {
1974                let _ = send.send(result);
1975            });
1976        gpu.device
1977            .poll(wgpu::PollType::wait_indefinitely())
1978            .unwrap();
1979        receive.recv().unwrap().unwrap();
1980        let mapped = buffer.slice(..).get_mapped_range().expect("a mapped range");
1981        let mut out = Vec::with_capacity((width * height) as usize);
1982        for y in 0..height {
1983            let start = (y * row) as usize;
1984            for pixel in mapped[start..start + width as usize * 8].chunks_exact(8) {
1985                out.push(std::array::from_fn(|c| {
1986                    half::f16::from_bits(u16::from_le_bytes([pixel[2 * c], pixel[2 * c + 1]]))
1987                        .to_f32()
1988                }));
1989            }
1990        }
1991        out
1992    }
1993
1994    #[test]
1995    #[ignore = "needs a GPU; run under pgpu"]
1996    fn gpu_reflected_edge_holds_and_resets_on_a_cut() {
1997        let gpu = pollster::block_on(pfx_gpu::Gpu::headless()).unwrap();
1998        let desk = desk(640, 360);
1999        let red = views(&gpu, &desk, [1.0, 0.0, 0.0]);
2000        let mut pass = Pass::new(&gpu.device, desk.width, desk.height);
2001        run(&gpu, &mut pass, &red, &desk, 0, false);
2002        let raw = read(&gpu, &pass._raw[0]);
2003        let (hw, hh) = (
2004            pass._raw[0].width() as usize,
2005            pass._raw[0].height() as usize,
2006        );
2007        let width = desk.width as usize;
2008        let mut edges = Vec::new();
2009        for y in 0..hh {
2010            let mut edge = None;
2011            for x in 1..hw {
2012                let floor = desk.surface[2 * y * width + 2 * x] == Surface::Floor
2013                    && desk.surface[2 * y * width + 2 * x - 2] == Surface::Floor;
2014                let card = |texel: [f32; 4]| texel[3] > 0.05 && texel[0] > texel[2];
2015                if floor && card(raw[y * hw + x - 1]) && !card(raw[y * hw + x]) {
2016                    edge = Some(x as f32);
2017                }
2018            }
2019            if let Some(x) = edge {
2020                edges.push((y, x));
2021            }
2022        }
2023        assert!(
2024            edges.len() > 30,
2025            "{} half rows reflect the edge",
2026            edges.len()
2027        );
2028        for pair in edges.windows(2) {
2029            assert!(
2030                pair[1].0 == pair[0].0 + 1 && (pair[1].1 - pair[0].1).abs() <= 1.0,
2031                "a stair-step on the GPU between half rows {} and {}: {} then {}",
2032                pair[0].0,
2033                pair[1].0,
2034                pair[0].1,
2035                pair[1].1
2036            );
2037        }
2038        for frame in 1..8 {
2039            run(&gpu, &mut pass, &red, &desk, frame, true);
2040        }
2041        let kept = read(&gpu, &pass._results[1 - pass.next]);
2042        assert!(
2043            kept.iter().any(|texel| texel[0] > 0.1),
2044            "red history builds up"
2045        );
2046        let green = views(&gpu, &desk, [0.0, 1.0, 0.0]);
2047        run(&gpu, &mut pass, &green, &desk, 8, false);
2048        let after = read(&gpu, &pass._results[1 - pass.next]);
2049        let fresh = read(&gpu, &pass._raw[0]);
2050        assert_eq!(
2051            after, fresh,
2052            "after a cut the result is this frame's trace alone"
2053        );
2054        assert!(
2055            after.iter().all(|texel| texel[0] < 1e-3),
2056            "no red survives the cut"
2057        );
2058    }
2059
2060    #[test]
2061    #[ignore = "needs a GPU; run under pgpu"]
2062    fn gpu_bench_4k() {
2063        let gpu = &pollster::block_on(pfx_gpu::Gpu::headless()).unwrap();
2064        let desk = desk(3840, 2160);
2065        let views = views(gpu, &desk, [1.0, 0.0, 0.0]);
2066        let mut pass = Pass::new(&gpu.device, desk.width, desk.height);
2067        let mut profiler = GpuProfiler::new(&gpu.device, &gpu.queue);
2068        let mut samples: std::collections::BTreeMap<String, Vec<f64>> = Default::default();
2069        let mut turns = pfx_gpu::pace::Turns::default();
2070        for frame_index in 0..14 {
2071            let mut encoder = gpu.device.create_command_encoder(&Default::default());
2072            pass.encode(
2073                &gpu.device,
2074                &mut encoder,
2075                &inputs(&views, &desk, frame_index, frame_index > 0),
2076                Some(&mut profiler),
2077            );
2078            let slot = profiler.finish(&mut encoder);
2079            gpu.queue.submit(Some(encoder.finish()));
2080            if let Some(slot) = slot {
2081                profiler.submitted(slot);
2082                gpu.device
2083                    .poll(wgpu::PollType::wait_indefinitely())
2084                    .unwrap();
2085                let timings: Vec<_> = profiler
2086                    .collect(&gpu.device)
2087                    .into_iter()
2088                    .flatten()
2089                    .collect();
2090                turns.add(timings.iter().map(|timing| timing.milliseconds).sum());
2091                if frame_index >= 4 {
2092                    for timing in timings {
2093                        samples
2094                            .entry(timing.label)
2095                            .or_default()
2096                            .push(timing.milliseconds);
2097                    }
2098                }
2099            }
2100        }
2101        assert!(!samples.is_empty(), "GPU timestamp query unavailable");
2102        let mut total = 0.0;
2103        for (label, values) in &mut samples {
2104            values.sort_by(f64::total_cmp);
2105            let median = values[values.len() / 2];
2106            total += median;
2107            eprintln!("4K {label} median: {median:.3} ms");
2108        }
2109        eprintln!("4K SSR median total: {total:.3} ms; target 1.000 ms");
2110        assert!(
2111            total <= 1.0,
2112            "4K SSR missed the 1.0 ms target: {total:.3} ms"
2113        );
2114        let result = read(gpu, &pass._results[1 - pass.next]);
2115        let hits = result.iter().filter(|texel| texel[3] > 0.0).count();
2116        eprintln!("4K SSR confident texels: {hits}");
2117        assert!(hits > 0, "synthetic scene produced no SSR hits");
2118    }
2119}