Skip to main content

TRACE_WGSL

Constant TRACE_WGSL 

Source
pub const TRACE_WGSL: &str = "\nfn ssr_encode_normal(n: vec3f) -> vec2f {\n    let p = n.xy / (abs(n.x) + abs(n.y) + abs(n.z));\n    if (n.z >= 0.0) { return p; }\n    return (vec2f(1.0) - abs(p.yx)) * select(vec2f(-1.0), vec2f(1.0), p >= vec2f(0.0));\n}\nfn ssr_decode_normal(e: vec2f) -> vec3f {\n    var n = vec3f(e, 1.0 - abs(e.x) - abs(e.y));\n    let t = max(-n.z, 0.0);\n    n.x += select(t, -t, n.x >= 0.0);\n    n.y += select(t, -t, n.y >= 0.0);\n    return normalize(n);\n}\nfn ssr_similarity(depth: f32, normal: vec3f, other_depth: f32, other_normal: vec3f, tolerance: f32) -> f32 {\n    if (depth <= 0.0 || other_depth <= 0.0) { return 0.0; }\n    let near = clamp(1.0 - abs(depth - other_depth) / max(tolerance, 0.000001), 0.0, 1.0);\n    let facing = smoothstep(0.6, 0.85, dot(normal, other_normal));\n    return near * facing;\n}\nconst SSR_DEPTH_TOLERANCE: f32 = 0.03;\n\nstruct ReflectParams {\n    size: vec2u,\n    half_size: vec2u,\n    projection: mat4x4f,\n    inverse_projection: mat4x4f,\n    history_valid: u32,\n    frame_index: u32,\n    refine: u32,\n    thickness: f32,\n    history_half: vec2u,\n}\n@group(0) @binding(0) var<uniform> params: ReflectParams;\nfn inside(uv: vec2f) -> bool {\n    return all(uv >= vec2f(0.0)) && all(uv < vec2f(1.0));\n}\nfn trace_pixel(cell: vec2u) -> vec2u {\n    let phase = params.frame_index % 4u;\n    let offset = vec2u(select(0u, 1u, phase == 1u || phase == 2u), select(0u, 1u, phase == 1u || phase == 3u));\n    return min(cell * 2u + offset, params.size - 1u);\n}\n\n@group(0) @binding(1) var depth_tex: texture_2d<f32>;\n@group(0) @binding(2) var normal_roughness: texture_2d<f32>;\n@group(0) @binding(3) var source_colour: texture_2d<f32>;\n@group(0) @binding(4) var colour_sampler: sampler;\n@group(0) @binding(5) var raw_out: texture_storage_2d<rgba16float, write>;\n@group(0) @binding(6) var guide_out: texture_storage_2d<rgba16float, write>;\nfn edge_weight(uv: vec2f) -> f32 {\n    let d = min(min(uv.x, 1.0 - uv.x), min(uv.y, 1.0 - uv.y));\n    return clamp(d / 0.08, 0.0, 1.0);\n}\nfn to_pixel(p: vec3f) -> vec3f {\n    let clip = params.projection * vec4f(p, 1.0);\n    let ndc = clip.xy / clip.w;\n    return vec3f((ndc * vec2f(0.5, -0.5) + vec2f(0.5)) * vec2f(params.size), clip.w);\n}\nfn scene_depth(pixel: vec2f) -> f32 {\n    let at = vec2i(clamp(pixel, vec2f(0.0), vec2f(params.size) - vec2f(1.0)));\n    return textureLoad(depth_tex, at, 0).r;\n}\nfn ray_depth(k: vec2f, s: f32) -> f32 {\n    return 1.0 / mix(k.x, k.y, s);\n}\nfn march_steps(roughness: f32) -> u32 {\n    return select(select(16u, 12u, roughness > 0.18), 8u, roughness > 0.45);\n}\nfn march(origin: vec3f, direction: vec3f, steps: u32) -> vec3f {\n    let near = -origin.z;\n    var reach = near * 64.0;\n    if (direction.z > 0.0) {\n        reach = min(reach, 0.9 * near / direction.z);\n    }\n    let start = to_pixel(origin);\n    let end = to_pixel(origin + direction * reach);\n    let delta = end.xy - start.xy;\n    let size = vec2f(params.size);\n    var limit = 1.0;\n    if (delta.x > 0.0) { limit = min(limit, (size.x - 0.5 - start.x) / delta.x); }\n    if (delta.x < 0.0) { limit = min(limit, (0.5 - start.x) / delta.x); }\n    if (delta.y > 0.0) { limit = min(limit, (size.y - 0.5 - start.y) / delta.y); }\n    if (delta.y < 0.0) { limit = min(limit, (0.5 - start.y) / delta.y); }\n    let span = length(delta);\n    let pixels = span * limit;\n    if (!(pixels >= 1.0)) { return vec3f(0.0); }\n    let k = vec2f(1.0 / start.z, 1.0 / end.z);\n    var previous = 0.0;\n    var previous_depth = start.z;\n    var front = start.z <= scene_depth(start.xy);\n    for (var i = 1u; i <= steps; i++) {\n        let u = f32(i) / f32(steps);\n        let travelled = min(max(pixels * u * u, f32(i)), pixels);\n        let s = travelled / span;\n        let ray = ray_depth(k, s);\n        let scene = scene_depth(start.xy + delta * s);\n        if (ray > scene && front && min(previous_depth, ray) <= scene * (1.0 + params.thickness)) {\n            var lo = previous;\n            var hi = s;\n            for (var j = 0u; j < params.refine; j++) {\n                let mid = 0.5 * (lo + hi);\n                if (ray_depth(k, mid) > scene_depth(start.xy + delta * mid)) {\n                    hi = mid;\n                } else {\n                    lo = mid;\n                }\n            }\n            let pixel = start.xy + delta * hi;\n            let hit_ray = ray_depth(k, hi);\n            let hit_scene = scene_depth(pixel);\n            let span_depth = abs(hit_ray - ray_depth(k, lo));\n            let beyond = abs(scene_depth(start.xy + delta * (hi + 1.0 / span)) - hit_scene);\n            let joined = abs(scene_depth(start.xy + delta * lo) - hit_scene) <= hit_scene * 0.005 + 2.0 * beyond + span_depth;\n            if (hit_ray - hit_scene <= hit_scene * params.thickness + span_depth && joined) {\n                let fade = 1.0 - smoothstep(0.8, 1.0, hi);\n                return vec3f(pixel, edge_weight(pixel / size) * fade);\n            }\n        }\n        previous = s;\n        previous_depth = ray;\n        front = ray <= scene;\n        if (travelled >= pixels) { break; }\n    }\n    return vec3f(0.0);\n}\nfn view_position(pixel: vec2i) -> vec3f {\n    let at = clamp(pixel, vec2i(0), vec2i(params.size) - vec2i(1));\n    let depth = textureLoad(depth_tex, at, 0).r;\n    let uv = (vec2f(at) + vec2f(0.5)) / vec2f(params.size);\n    let clip = vec4f(uv * vec2f(2.0, -2.0) + vec2f(-1.0, 1.0), 1.0, 1.0);\n    let unprojected = params.inverse_projection * clip;\n    return unprojected.xyz / unprojected.w * (depth / max(-unprojected.z / unprojected.w, 0.0001));\n}\nfn nearer(origin: vec3f, a: vec3f, b: vec3f) -> vec3f {\n    return select(origin - b, a - origin, abs(a.z - origin.z) < abs(b.z - origin.z));\n}\n@compute @workgroup_size(8, 8)\nfn trace_pass(@builtin(global_invocation_id) id: vec3u) {\n    if (any(id.xy >= params.half_size)) { return; }\n    let pixel = trace_pixel(id.xy);\n    let depth = textureLoad(depth_tex, vec2i(pixel), 0).r;\n    let nr = textureLoad(normal_roughness, vec2i(pixel), 0);\n    if (depth <= 0.0 || length(nr.xyz) < 0.1) {\n        textureStore(raw_out, vec2i(id.xy), vec4f(0.0));\n        textureStore(guide_out, vec2i(id.xy), vec4f(0.0));\n        return;\n    }\n    let p = vec2i(pixel);\n    let origin = view_position(p);\n    let across = nearer(origin, view_position(p + vec2i(1, 0)), view_position(p - vec2i(1, 0)));\n    let along = nearer(origin, view_position(p + vec2i(0, 1)), view_position(p - vec2i(0, 1)));\n    let normal = normalize(nr.xyz);\n    var surface = cross(along, across);\n    surface = select(normal, normalize(surface), dot(surface, surface) > 0.0);\n    surface = select(surface, -surface, dot(surface, origin) > 0.0);\n    let roughness = clamp(nr.w, 0.0, 1.0);\n    let view_ray = normalize(origin);\n    let angle = 6.2831853 * f32(params.frame_index % 8u) * 0.125;\n    let jitter = vec3f(cos(angle), sin(angle), 0.0) * roughness * roughness * 0.15;\n    var direction = normalize(reflect(view_ray, normal) + jitter);\n    let lift = 0.05 - dot(direction, surface);\n    if (lift > 0.0) {\n        direction = normalize(direction + surface * lift);\n    }\n    var result = vec4f(0.0);\n    var hit = vec3f(0.0);\n    let fade = 1.0 - smoothstep(0.6, 0.75, roughness);\n    if (fade > 0.0) {\n        hit = march(origin + surface * depth * 0.002, direction, march_steps(roughness));\n    }\n    if (hit.z > 0.0) {\n        let at = vec2i(clamp(hit.xy, vec2f(0.0), vec2f(params.size) - vec2f(1.0)));\n        let facing = dot(textureLoad(normal_roughness, at, 0).xyz, direction);\n        let confidence = hit.z * fade * clamp(1.0 - 4.0 * (facing - 0.1), 0.0, 1.0);\n        let lod = clamp(roughness * 4.0 - 0.5, 0.0, f32(textureNumLevels(source_colour) - 1u));\n        let colour = textureSampleLevel(source_colour, colour_sampler, hit.xy / vec2f(textureDimensions(source_colour)), lod).rgb;\n        result = vec4f(colour * confidence, confidence);\n    }\n    textureStore(raw_out, vec2i(id.xy), result);\n    textureStore(guide_out, vec2i(id.xy), vec4f(depth, ssr_encode_normal(normal), roughness));\n}\n";