@multiplekex/shallot 0.2.3 → 0.2.5

package/src/standard/raytracing/shaders.ts

@@ -0,0 +1,815 @@
+import type { SurfaceData } from "../render/surface";
+import {
+    compileVertexBody,
+    WGSL_LIGHTING_CALC,
+    SKY_DIR_WGSL,
+    HAZE_WGSL,
+    SKY_WGSL,
+    NOISE_WGSL,
+} from "../render/surface/shaders";
+import {
+    BVH_STRUCTS,
+    TLAS_BLAS_STRUCTS,
+    TLAS_BLAS_BINDINGS,
+    TLAS_BLAS_TRAVERSAL,
+    TLAS_BLAS_ANY_HIT,
+} from "./bvh/traverse";
+import { RAY_STRUCT_WGSL, HIT_RESULT_STRUCT_WGSL } from "./bvh/structs";
+import {
+    SURFACE_DATA_STRUCT_WGSL,
+    SCENE_STRUCT_WGSL,
+    SKY_STRUCT_WGSL,
+    DATA_STRUCT_WGSL,
+} from "../render/surface/structs";
+
+const SPECULAR_WGSL = /* wgsl */ `
+const DIELECTRIC_F0: f32 = 0.04;
+
+fn blinnPhongSpecular(N: vec3<f32>, L: vec3<f32>, V: vec3<f32>, roughness: f32) -> f32 {
+    let H = normalize(L + V);
+    let NdotH = max(dot(N, H), 0.0);
+    let shininess = pow(2.0, (1.0 - roughness) * 10.0);
+    let intensity = (1.0 - roughness) * (1.0 - roughness);
+    return pow(NdotH, shininess) * intensity;
+}
+
+fn schlickFresnel(cosTheta: f32, F0: vec3<f32>) -> vec3<f32> {
+    return F0 + (vec3(1.0) - F0) * pow(1.0 - cosTheta, 5.0);
+}
+
+fn computeF0Vec(baseColor: vec3<f32>, metallic: f32) -> vec3<f32> {
+    return mix(vec3(DIELECTRIC_F0), baseColor, metallic);
+}
+`;
+
+const REFRACTION_WGSL = /* wgsl */ `
+fn refractRay(I: vec3<f32>, N: vec3<f32>, eta: f32) -> vec4<f32> {
+    let cosI = -dot(I, N);
+    let sinT2 = eta * eta * (1.0 - cosI * cosI);
+    if (sinT2 > 1.0) {
+        return vec4(reflect(I, N), 1.0);
+    }
+    let cosT = sqrt(1.0 - sinT2);
+    return vec4(normalize(eta * I + (eta * cosI - cosT) * N), 0.0);
+}
+
+fn fresnelSchlickIOR(cosTheta: f32, n1: f32, n2: f32, metallic: f32) -> f32 {
+    let r0 = pow((n1 - n2) / (n1 + n2), 2.0);
+    let baseFresnel = r0 + (1.0 - r0) * pow(1.0 - cosTheta, 5.0);
+    let minReflection = metallic * 0.5;
+    return max(baseFresnel, minReflection);
+}
+`;
+
+const SHADOW_HELPER_WGSL = /* wgsl */ `
+fn getShadow(surface: SurfaceData) -> f32 {
+    let L = -scene.sunDirection.xyz;
+    let NdotL = max(dot(surface.normal, L), 0.0);
+    if (NdotL <= 0.0) { return 1.0; }
+    let origin = surface.worldPos + surface.normal * 0.001;
+    return sampleSoftShadow(origin, L, scene.shadowSoftness, scene.shadowSamples);
+}
+`;
+
+const SHADOW_WGSL = /* wgsl */ `
+const GOLDEN_ANGLE: f32 = 2.39996323;
+const SHADOW_TRANSPARENCY_EPSILON: f32 = 0.001;
+const MAX_SHADOW_TRANSPARENT_DEPTH: u32 = 2u;
+
+fn buildTangentBasis(dir: vec3<f32>) -> mat3x3<f32> {
+    let up = select(vec3(0.0, 1.0, 0.0), vec3(1.0, 0.0, 0.0), abs(dir.y) > 0.99);
+    let tangent = normalize(cross(dir, up));
+    let bitangent = cross(dir, tangent);
+    return mat3x3(tangent, bitangent, dir);
+}
+
+fn traceShadowWithTransparency(origin: vec3<f32>, sunDir: vec3<f32>, tMax: f32) -> f32 {
+    var shadowFactor = 1.0;
+    var currentPos = origin;
+
+    for (var i = 0u; i < MAX_SHADOW_TRANSPARENT_DEPTH; i++) {
+        var shadowRay: Ray;
+        shadowRay.origin = currentPos;
+        shadowRay.direction = sunDir;
+
+        let hit = trace(shadowRay);
+        if (!hit.hit || hit.t > tMax) {
+            break;
+        }
+
+        let shadow = getData(hit.entityId);
+        let opacity = shadow.baseColor.a;
+        shadowFactor *= (1.0 - opacity);
+
+        if (opacity >= 1.0 || shadowFactor < 0.02) {
+            break;
+        }
+
+        currentPos = hit.worldPos + sunDir * SHADOW_TRANSPARENCY_EPSILON;
+    }
+
+    return shadowFactor;
+}
+
+fn sampleSoftShadow(origin: vec3<f32>, sunDir: vec3<f32>, softness: f32, samples: u32) -> f32 {
+    if (samples == 0u) {
+        return 1.0;
+    }
+    if (samples == 1u || softness <= 0.0) {
+        return traceShadowWithTransparency(origin, sunDir, 1000.0);
+    }
+
+    let basis = buildTangentBasis(sunDir);
+    var totalShadow = 0.0;
+
+    let noise = fract(sin(dot(origin.xy, vec2(12.9898, 78.233))) * 43758.5453);
+    let rotationOffset = noise * 6.28318530718;
+
+    for (var i = 0u; i < samples; i++) {
+        let angle = f32(i) * GOLDEN_ANGLE + rotationOffset;
+        let z = (f32(i) + 0.5) / f32(samples);
+        let r = sqrt(1.0 - z * z) * softness * 0.1;
+
+        let localOffset = vec3(cos(angle) * r, sin(angle) * r, 0.0);
+        let sampleDir = normalize(sunDir + basis * localOffset);
+
+        totalShadow += traceShadowWithTransparency(origin, sampleDir, 1000.0);
+    }
+
+    return totalShadow / f32(samples);
+}
+`;
+
+function compileApplyLighting(lit?: boolean, shadows?: boolean, reflections?: boolean): string {
+    if (lit === false) {
+        return "return surface.baseColor;";
+    }
+
+    if (shadows && reflections) {
+        // shadowFactor is now passed as parameter, computed once in main kernel
+        return `
+    let V = -rayDir;
+    let L = -scene.sunDirection.xyz;
+    let NdotL = max(dot(surface.normal, L), 0.0);
+    let NdotV = max(dot(surface.normal, V), 0.0);
+
+    let F0 = computeF0Vec(surface.baseColor, surface.metallic);
+    let F = schlickFresnel(NdotV, F0);
+
+    let ambient = scene.ambientColor.rgb * scene.ambientColor.a;
+    let sunDiffuse = scene.sunColor.rgb * NdotL * shadowFactor;
+    let diffuseWeight = 1.0 - surface.metallic;
+    let diffuseColor = surface.baseColor * (ambient + sunDiffuse) * diffuseWeight + surface.emission;
+
+    let specTerm = blinnPhongSpecular(surface.normal, L, V, surface.roughness);
+    let specular = scene.sunColor.rgb * specTerm * F * NdotL * shadowFactor;
+
+    var finalColor = diffuseColor * (vec3(1.0) - F) + specular;
+
+    if (scene.reflectionDepth > 0u) {
+        finalColor = traceReflections(
+            surface.worldPos, surface.normal, rayDir,
+            diffuseColor + specular, surface.baseColor,
+            surface.roughness, surface.metallic,
+            scene.reflectionDepth
+        );
+    }
+    return finalColor;`;
+    }
+
+    if (shadows) {
+        return `
+    let V = -rayDir;
+    let L = -scene.sunDirection.xyz;
+    let NdotL = max(dot(surface.normal, L), 0.0);
+    let NdotV = max(dot(surface.normal, V), 0.0);
+
+    let F0 = computeF0Vec(surface.baseColor, surface.metallic);
+    let F = schlickFresnel(NdotV, F0);
+    let diffuseWeight = (1.0 - surface.metallic) * (vec3(1.0) - F);
+    let ambient = scene.ambientColor.rgb * scene.ambientColor.a;
+    let sunDiffuse = scene.sunColor.rgb * NdotL * shadowFactor;
+    let diffuseColor = surface.baseColor * (ambient + sunDiffuse) * diffuseWeight + surface.emission;
+    let specTerm = blinnPhongSpecular(surface.normal, L, V, surface.roughness);
+    let specular = scene.sunColor.rgb * specTerm * F * NdotL * shadowFactor;
+    return diffuseColor + specular;`;
+    }
+
+    if (reflections) {
+        return `
+    let V = -rayDir;
+    let L = -scene.sunDirection.xyz;
+    let NdotL = max(dot(surface.normal, L), 0.0);
+    let NdotV = max(dot(surface.normal, V), 0.0);
+
+    let F0 = computeF0Vec(surface.baseColor, surface.metallic);
+    let F = schlickFresnel(NdotV, F0);
+
+    let ambient = scene.ambientColor.rgb * scene.ambientColor.a;
+    let sunDiffuse = scene.sunColor.rgb * NdotL;
+    let diffuseWeight = 1.0 - surface.metallic;
+    let diffuseColor = surface.baseColor * (ambient + sunDiffuse) * diffuseWeight + surface.emission;
+
+    let specTerm = blinnPhongSpecular(surface.normal, L, V, surface.roughness);
+    let specular = scene.sunColor.rgb * specTerm * F * NdotL;
+
+    var finalColor = diffuseColor * (vec3(1.0) - F) + specular;
+
+    if (scene.reflectionDepth > 0u) {
+        finalColor = traceReflections(
+            surface.worldPos, surface.normal, rayDir,
+            diffuseColor + specular, surface.baseColor,
+            surface.roughness, surface.metallic,
+            scene.reflectionDepth
+        );
+    }
+    return finalColor;`;
+    }
+
+    return `${WGSL_LIGHTING_CALC}
+    return surface.baseColor * lighting + surface.emission;`;
+}
+
+const EPSILON = 1e-7;
+
+export const RT_STRUCTS = /* wgsl */ `
+${RAY_STRUCT_WGSL}
+${HIT_RESULT_STRUCT_WGSL}
+${SURFACE_DATA_STRUCT_WGSL}
+${SCENE_STRUCT_WGSL}
+${SKY_STRUCT_WGSL}
+${DATA_STRUCT_WGSL}
+`;
+
+export const RT_BINDINGS = /* wgsl */ `
+@group(0) @binding(0) var<uniform> scene: Scene;
+@group(0) @binding(1) var<storage, read> data: array<Data>;
+@group(0) @binding(2) var output_scene: texture_storage_2d<rgba8unorm, write>;
+@group(0) @binding(3) var output_depth: texture_storage_2d<r32float, write>;
+@group(0) @binding(4) var output_entityId: texture_storage_2d<r32uint, write>;
+@group(0) @binding(5) var<uniform> sky: Sky;
+`;
+
+export const RT_UTILS = /* wgsl */ `
+fn getData(eid: u32) -> Data {
+    return data[eid];
+}
+
+fn getSurfaceType(eid: u32) -> u32 {
+    return data[eid].flags & 0xFFu;
+}
+
+fn getVolume(eid: u32) -> u32 {
+    return (data[eid].flags >> 8u) & 0xFu;
+}
+
+fn getShapeId(eid: u32) -> u32 {
+    return data[eid].flags >> 16u;
+}
+
+fn getInstanceCount() -> u32 {
+    return scene.instanceCount;
+}
+`;
+
+export const RT_RAY_GEN = /* wgsl */ `
+${SKY_DIR_WGSL}
+
+struct PrimaryRay {
+    origin: vec3<f32>,
+    direction: vec3<f32>,
+    skyDir: vec3<f32>,
+}
+
+fn generateRay(screenX: f32, screenY: f32) -> PrimaryRay {
+    var result: PrimaryRay;
+    result.skyDir = computeSkyDir(screenX, screenY);
+
+    let width = scene.viewport.x;
+    let height = scene.viewport.y;
+    let ndcX = screenX * 2.0 - 1.0;
+    let ndcY = 1.0 - screenY * 2.0;
+    let aspect = width / height;
+
+    let cameraWorld = scene.cameraWorld;
+    let camPosX = cameraWorld[3][0];
+    let camPosY = cameraWorld[3][1];
+    let camPosZ = cameraWorld[3][2];
+
+    if (scene.cameraMode > 0.5) {
+        let r00 = cameraWorld[0][0]; let r10 = cameraWorld[0][1]; let r20 = cameraWorld[0][2];
+        let r01 = cameraWorld[1][0]; let r11 = cameraWorld[1][1]; let r21 = cameraWorld[1][2];
+        let r02 = cameraWorld[2][0]; let r12 = cameraWorld[2][1]; let r22 = cameraWorld[2][2];
+
+        let halfHeight = scene.cameraSize;
+        let halfWidth = halfHeight * aspect;
+        let offsetX = ndcX * halfWidth;
+        let offsetY = ndcY * halfHeight;
+        let fwdX = -r02; let fwdY = -r12; let fwdZ = -r22;
+
+        result.origin = vec3(
+            camPosX + r00 * offsetX + r01 * offsetY + fwdX * scene.near,
+            camPosY + r10 * offsetX + r11 * offsetY + fwdY * scene.near,
+            camPosZ + r20 * offsetX + r21 * offsetY + fwdZ * scene.near
+        );
+        result.direction = vec3(fwdX, fwdY, fwdZ);
+    } else {
+        let dir = result.skyDir;
+        result.origin = vec3(camPosX + dir.x * scene.near, camPosY + dir.y * scene.near, camPosZ + dir.z * scene.near);
+        result.direction = dir;
+    }
+
+    return result;
+}
+`;
+
+export const RT_INTERSECTION = /* wgsl */ `
+const EPSILON: f32 = ${EPSILON};
+`;
+
+export function compileRTSurface(data: SurfaceData): string {
+    return compileUberShader([data]);
+}
+
+function compileSurfaceVariant(id: number, data: SurfaceData): string {
+    const vertexBody = compileVertexBody(data.vertex);
+    const fragmentBody = data.fragment ?? "";
+    const lightingCode = compileApplyLighting(data.lit, true, false);
+
+    return `
+fn userVertexTransform_${id}(worldPos: vec3<f32>, normal: vec3<f32>, eid: u32) -> vec3<f32> {
+    ${vertexBody}
+}
+
+fn userFragment_${id}(surface: ptr<function, SurfaceData>, position: vec4<f32>) {
+    ${fragmentBody}
+}
+
+fn applyLighting_${id}(surface: SurfaceData, rayDir: vec3<f32>, shadowFactor: f32) -> vec3<f32> {
+    ${lightingCode}
+}
+`;
+}
+
+function compileDispatchFunctions(surfaceCount: number): string {
+    const vertexCases = Array.from(
+        { length: surfaceCount },
+        (_, i) => `        case ${i}u: { return userVertexTransform_${i}(worldPos, normal, eid); }`
+    ).join("\n");
+
+    const fragmentCases = Array.from(
+        { length: surfaceCount },
+        (_, i) => `        case ${i}u: { userFragment_${i}(surface, position); }`
+    ).join("\n");
+
+    const lightingCases = Array.from(
+        { length: surfaceCount },
+        (_, i) =>
+            `        case ${i}u: { return applyLighting_${i}(surface, rayDir, shadowFactor); }`
+    ).join("\n");
+
+    return `
+fn dispatchVertexTransform(surfaceId: u32, worldPos: vec3<f32>, normal: vec3<f32>, eid: u32) -> vec3<f32> {
+    switch surfaceId {
+${vertexCases}
+        default: { return userVertexTransform_0(worldPos, normal, eid); }
+    }
+}
+
+fn dispatchFragment(surfaceId: u32, surface: ptr<function, SurfaceData>, position: vec4<f32>) {
+    switch surfaceId {
+${fragmentCases}
+        default: { userFragment_0(surface, position); }
+    }
+}
+
+fn dispatchLighting(surfaceId: u32, surface: SurfaceData, rayDir: vec3<f32>, shadowFactor: f32) -> vec3<f32> {
+    switch surfaceId {
+${lightingCases}
+        default: { return applyLighting_0(surface, rayDir, shadowFactor); }
+    }
+}
+`;
+}
+
+export function compileUberShader(surfaces: SurfaceData[]): string {
+    const surfaceVariants = surfaces.map((s, i) => compileSurfaceVariant(i, s)).join("\n");
+    const dispatchFunctions = compileDispatchFunctions(surfaces.length);
+
+    return /* wgsl */ `
+${RT_STRUCTS}
+${BVH_STRUCTS}
+${TLAS_BLAS_STRUCTS}
+${RT_BINDINGS}
+${RT_UTILS}
+${TLAS_BLAS_BINDINGS}
+${RT_RAY_GEN}
+${RT_INTERSECTION}
+${TLAS_BLAS_TRAVERSAL}
+${TLAS_BLAS_ANY_HIT}
+${SHADOW_WGSL}
+${SHADOW_HELPER_WGSL}
+${SPECULAR_WGSL}
+${REFRACTION_WGSL}
+const PI: f32 = 3.14159265359;
+${NOISE_WGSL}
+${SKY_WGSL}
+${HAZE_WGSL}
+
+${surfaceVariants}
+${dispatchFunctions}
+
+const KIND_PRIMARY: u32 = 0u;
+const KIND_REFRACT: u32 = 1u;
+const KIND_REFLECT: u32 = 2u;
+const KIND_BEHIND:  u32 = 3u;
+const MAX_TASKS: u32 = 16u;
+const MIN_CONTRIBUTION: f32 = 0.02;
+const REFLECTION_EPSILON: f32 = 0.001;
+const TRANSPARENCY_EPSILON: f32 = 0.001;
+const REFRACTION_EPSILON: f32 = 0.001;
+
+const VOLUME_SOLID: u32 = 0u;
+const VOLUME_HALF_SPACE: u32 = 1u;
+
+var<private> tOrigin:      array<vec3f, MAX_TASKS>;
+var<private> tDirection:   array<vec3f, MAX_TASKS>;
+var<private> tKind:        array<u32,   MAX_TASKS>;
+var<private> tWeight:      array<vec3f, MAX_TASKS>;
+var<private> tIOR:         array<f32,   MAX_TASKS>;
+var<private> tMediumColor: array<vec3f, MAX_TASKS>;
+var<private> tTotalDist:   array<f32,   MAX_TASKS>;
+var<private> tInMedium:    array<u32,   MAX_TASKS>;
+var<private> tBounces:     array<u32,   MAX_TASKS>;
+var<private> tVolumeType:  array<u32,   MAX_TASKS>;
+var<private> tRemaining:   array<f32,   MAX_TASKS>;
+
+@compute @workgroup_size(8, 8)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let width = u32(scene.viewport.x);
+    let height = u32(scene.viewport.y);
+
+    if (gid.x >= width || gid.y >= height) {
+        return;
+    }
+
+    let screenX = (f32(gid.x) + 0.5) / f32(width);
+    let screenY = (f32(gid.y) + 0.5) / f32(height);
+
+    let primary = generateRay(screenX, screenY);
+
+    tOrigin[0] = primary.origin;
+    tDirection[0] = primary.direction;
+    tKind[0] = KIND_PRIMARY;
+    tWeight[0] = vec3(1.0);
+    tRemaining[0] = 1.0;
+    var taskCount = 1u;
+
+    var finalColor = vec3(0.0);
+    var depth = 1e30;
+    var entityId = 0u;
+    var firstHit = true;
+
+    for (var i = 0u; i < MAX_TASKS; i++) {
+        if (i >= taskCount) {
+            break;
+        }
+
+        let avgWeight = (tWeight[i].x + tWeight[i].y + tWeight[i].z) / 3.0;
+        if (avgWeight < MIN_CONTRIBUTION) {
+            continue;
+        }
+
+        var ray: Ray;
+        ray.origin = tOrigin[i];
+        ray.direction = tDirection[i];
+        let hit = trace(ray);
+        let kind = tKind[i];
+        let weight = tWeight[i];
+
+        if (kind == KIND_PRIMARY) {
+            if (!hit.hit) {
+                finalColor += sampleSky(primary.skyDir) * weight;
+                continue;
+            }
+
+            let eid = hit.entityId;
+            let d = getData(eid);
+            let opacity = d.baseColor.a;
+            let surfaceId = getSurfaceType(eid);
+
+            if (firstHit) {
+                depth = hit.t;
+                entityId = eid;
+                firstHit = false;
+            }
+
+            let finalWorldPos = dispatchVertexTransform(surfaceId, hit.worldPos, hit.normal, eid);
+
+            var surface: SurfaceData;
+            surface.baseColor = d.baseColor.rgb;
+            surface.roughness = d.pbr.x;
+            surface.metallic = d.pbr.y;
+            surface.opacity = opacity;
+            surface.emission = d.emission.rgb * d.emission.a;
+            surface.normal = hit.normal;
+            surface.worldPos = finalWorldPos;
+
+            dispatchFragment(surfaceId, &surface, vec4(f32(gid.x), f32(gid.y), hit.t, 1.0));
+
+            let ior = d.pbr.z;
+            if (ior > 1.0 && scene.refractionDepth > 0u) {
+                let entering = dot(ray.direction, surface.normal) < 0.0;
+                let volumeType = getVolume(eid);
+
+                if (entering || volumeType == VOLUME_SOLID) {
+                    let n = select(-surface.normal, surface.normal, entering);
+                    let n1 = select(ior, 1.0, entering);
+                    let n2 = select(1.0, ior, entering);
+                    let cosI = abs(dot(ray.direction, n));
+                    let fresnel = fresnelSchlickIOR(cosI, n1, n2, surface.metallic);
+
+                    if (taskCount < MAX_TASKS) {
+                        let refractResult = refractRay(ray.direction, n, n1 / n2);
+                        let isTIR = refractResult.w > 0.5;
+                        tOrigin[taskCount] = select(
+                            surface.worldPos - n * REFRACTION_EPSILON,
+                            surface.worldPos + n * REFRACTION_EPSILON,
+                            isTIR
+                        );
+                        tDirection[taskCount] = refractResult.xyz;
+                        tKind[taskCount] = KIND_REFRACT;
+                        tWeight[taskCount] = weight * (1.0 - fresnel);
+                        tIOR[taskCount] = ior;
+                        tMediumColor[taskCount] = surface.baseColor;
+                        tTotalDist[taskCount] = 0.0;
+                        tInMedium[taskCount] = select(0u, 1u, entering);
+                        tBounces[taskCount] = scene.refractionDepth;
+                        tVolumeType[taskCount] = volumeType;
+                        taskCount++;
+                    }
+
+                    let litColor = dispatchLighting(surfaceId, surface, primary.direction, getShadow(surface));
+                    let fresnelWeight = weight * fresnel;
+                    let refractionWeight = weight * (1.0 - fresnel);
+
+                    let rSmoothness = 1.0 - surface.roughness;
+                    let rRoughnessAtten = rSmoothness * rSmoothness;
+
+                    if (fresnel >= MIN_CONTRIBUTION && scene.reflectionDepth > 0u && taskCount < MAX_TASKS) {
+                        finalColor += applyHaze(litColor, hit.t) * refractionWeight;
+
+                        tOrigin[taskCount] = surface.worldPos + surface.normal * REFLECTION_EPSILON;
+                        tDirection[taskCount] = reflect(ray.direction, surface.normal);
+                        tKind[taskCount] = KIND_REFLECT;
+                        tWeight[taskCount] = fresnelWeight * rRoughnessAtten;
+                        tBounces[taskCount] = scene.reflectionDepth;
+                        tRemaining[taskCount] = 1.0;
+                        taskCount++;
+                    } else {
+                        finalColor += applyHaze(litColor, hit.t) * fresnelWeight;
+                    }
+                    continue;
+                }
+            }
+
+            let litColor = dispatchLighting(surfaceId, surface, primary.direction, getShadow(surface));
+
+            let V = -ray.direction;
+            let NdotV = max(dot(surface.normal, V), 0.0);
+            let F0 = computeF0Vec(surface.baseColor, surface.metallic);
+            let F = schlickFresnel(NdotV, F0);
+            let smoothness = 1.0 - surface.roughness;
+            let roughnessAtten = smoothness * smoothness;
+            let reflWeight = F * roughnessAtten;
+            let avgRefl = (reflWeight.x + reflWeight.y + reflWeight.z) / 3.0;
+
+            if (avgRefl >= MIN_CONTRIBUTION && scene.reflectionDepth > 0u && taskCount < MAX_TASKS) {
+                finalColor += applyHaze(litColor, hit.t) * weight * opacity * (vec3(1.0) - reflWeight);
+
+                tOrigin[taskCount] = surface.worldPos + surface.normal * REFLECTION_EPSILON;
+                tDirection[taskCount] = reflect(ray.direction, surface.normal);
+                tKind[taskCount] = KIND_REFLECT;
+                tWeight[taskCount] = weight * opacity * reflWeight;
+                tBounces[taskCount] = scene.reflectionDepth;
+                tRemaining[taskCount] = 1.0;
+                taskCount++;
+            } else {
+                finalColor += applyHaze(litColor, hit.t) * weight * opacity;
+            }
+
+            if (opacity < 1.0) {
+                let nextWeight = weight * (1.0 - opacity);
+                let avgNext = (nextWeight.x + nextWeight.y + nextWeight.z) / 3.0;
+                if (avgNext >= MIN_CONTRIBUTION && taskCount < MAX_TASKS) {
+                    tOrigin[taskCount] = hit.worldPos + ray.direction * TRANSPARENCY_EPSILON;
+                    tDirection[taskCount] = ray.direction;
+                    tKind[taskCount] = KIND_PRIMARY;
+                    tWeight[taskCount] = nextWeight;
+                    tRemaining[taskCount] = 1.0;
+                    taskCount++;
+                }
+            }
+        } else if (kind == KIND_REFRACT) {
+            let isHalfSpace = tVolumeType[i] == VOLUME_HALF_SPACE;
+            let inMedium = tInMedium[i] == 1u;
+            let totalDist = tTotalDist[i];
+
+            if (!hit.hit) {
+                if (isHalfSpace && inMedium) {
+                    let depthFade = exp(-totalDist * 0.1);
+                    finalColor += weight * tMediumColor[i] * depthFade * scene.ambientColor.rgb * scene.ambientColor.a;
+                } else {
+                    var skyColor = sampleSky(ray.direction);
+                    if (inMedium && totalDist > 0.0) {
+                        let absorption = exp(-totalDist * 0.5);
+                        skyColor = mix(skyColor * tMediumColor[i], skyColor, absorption);
+                    }
+                    finalColor += weight * skyColor;
+                }
+                continue;
+            }
+
+            var newTotalDist = totalDist;
+            if (inMedium) {
+                newTotalDist += hit.t;
+            }
+
+            let refEid = hit.entityId;
+            let refracted = getData(refEid);
+            let refOpacity = refracted.baseColor.a;
+            let refIOR = refracted.pbr.z;
+
+            if (refIOR <= 1.0) {
+                let refSurfaceId = getSurfaceType(refEid);
+                let refWorldPos = dispatchVertexTransform(refSurfaceId, hit.worldPos, hit.normal, refEid);
+
+                var refSurface: SurfaceData;
+                refSurface.baseColor = refracted.baseColor.rgb;
+                refSurface.roughness = refracted.pbr.x;
+                refSurface.metallic = refracted.pbr.y;
+                refSurface.opacity = refOpacity;
+                refSurface.emission = refracted.emission.rgb * refracted.emission.a;
+                refSurface.normal = hit.normal;
+                refSurface.worldPos = refWorldPos;
+
+                dispatchFragment(refSurfaceId, &refSurface, vec4(0.0, 0.0, hit.t, 1.0));
+
+                var surfaceColor = dispatchLighting(refSurfaceId, refSurface, ray.direction, getShadow(refSurface));
+
+                if (newTotalDist > 0.0) {
+                    let absorption = exp(-newTotalDist * 0.5);
+                    surfaceColor = mix(surfaceColor * tMediumColor[i], surfaceColor, absorption);
+                }
+
+                surfaceColor = applyHaze(surfaceColor, hit.t);
+
+                if (refOpacity >= 1.0) {
+                    finalColor += weight * surfaceColor;
+                } else {
+                    finalColor += weight * surfaceColor * refOpacity;
+
+                    if (taskCount < MAX_TASKS) {
+                        tOrigin[taskCount] = hit.worldPos + ray.direction * REFRACTION_EPSILON;
+                        tDirection[taskCount] = ray.direction;
+                        tKind[taskCount] = KIND_BEHIND;
+                        tWeight[taskCount] = weight * (1.0 - refOpacity);
+                        taskCount++;
+                    }
+                }
+                continue;
+            }
+
+            let entering = dot(ray.direction, hit.normal) < 0.0;
+
+            if (tBounces[i] > 0u && taskCount < MAX_TASKS) {
+                let n = select(-hit.normal, hit.normal, entering);
+                let n1 = select(tIOR[i], 1.0, entering);
+                let n2 = select(1.0, tIOR[i], entering);
+                let eta = n1 / n2;
+
+                let refractResult = refractRay(ray.direction, n, eta);
+                let isTIR = refractResult.w > 0.5;
+
+                var nextInMedium = inMedium;
+                if (!isTIR) {
+                    nextInMedium = entering;
+                }
+
+                tOrigin[taskCount] = select(
+                    hit.worldPos - n * REFRACTION_EPSILON,
+                    hit.worldPos + n * REFRACTION_EPSILON,
+                    isTIR
+                );
+                tDirection[taskCount] = refractResult.xyz;
+                tKind[taskCount] = KIND_REFRACT;
+                tWeight[taskCount] = weight;
+                tIOR[taskCount] = refIOR;
+                tMediumColor[taskCount] = select(tMediumColor[i], refracted.baseColor.rgb, entering);
+                tTotalDist[taskCount] = select(newTotalDist, 0.0, entering);
+                tInMedium[taskCount] = select(0u, 1u, nextInMedium);
+                tBounces[taskCount] = tBounces[i] - 1u;
+                tVolumeType[taskCount] = tVolumeType[i];
+                taskCount++;
+            } else {
+                if (isHalfSpace && inMedium) {
+                    let depthFade = exp(-newTotalDist * 0.1);
+                    finalColor += weight * tMediumColor[i] * depthFade * scene.ambientColor.rgb * scene.ambientColor.a;
+                } else {
+                    var skyColor = sampleSky(ray.direction);
+                    if (inMedium && newTotalDist > 0.0) {
+                        let absorption = exp(-newTotalDist * 0.5);
+                        skyColor = mix(skyColor * tMediumColor[i], skyColor, absorption);
+                    }
+                    finalColor += weight * skyColor;
+                }
+            }
+        } else if (kind == KIND_REFLECT) {
+            if (!hit.hit) {
+                finalColor += weight * sampleSky(ray.direction);
+                continue;
+            }
+
+            let eid = hit.entityId;
+            let hitData = getData(eid);
+            let hitSurfaceId = getSurfaceType(eid);
+            let hitWorldPos = dispatchVertexTransform(hitSurfaceId, hit.worldPos, hit.normal, eid);
+
+            var hitSurface: SurfaceData;
+            hitSurface.baseColor = hitData.baseColor.rgb;
+            hitSurface.roughness = hitData.pbr.x;
+            hitSurface.metallic = hitData.pbr.y;
+            hitSurface.opacity = hitData.baseColor.a;
+            hitSurface.emission = hitData.emission.rgb * hitData.emission.a;
+            hitSurface.normal = hit.normal;
+            hitSurface.worldPos = hitWorldPos;
+
+            dispatchFragment(hitSurfaceId, &hitSurface, vec4(0.0, 0.0, hit.t, 1.0));
+
+            let hitOpacity = hitSurface.opacity;
+            let hitLit = dispatchLighting(hitSurfaceId, hitSurface, ray.direction, getShadow(hitSurface));
+
+            finalColor += applyHaze(hitLit, hit.t) * weight * hitOpacity * tRemaining[i];
+
+            if (hitOpacity < 1.0) {
+                tRemaining[i] *= (1.0 - hitOpacity);
+                if (tRemaining[i] >= MIN_CONTRIBUTION && taskCount < MAX_TASKS) {
+                    tOrigin[taskCount] = hit.worldPos + ray.direction * REFLECTION_EPSILON;
+                    tDirection[taskCount] = ray.direction;
+                    tKind[taskCount] = KIND_REFLECT;
+                    tWeight[taskCount] = weight;
+                    tBounces[taskCount] = tBounces[i];
+                    tRemaining[taskCount] = tRemaining[i];
+                    taskCount++;
+                } else {
+                    finalColor += sampleSky(ray.direction) * weight * tRemaining[i];
+                }
+                continue;
+            }
+
+            if (tBounces[i] > 1u && taskCount < MAX_TASKS) {
+                let hitSmoothness = 1.0 - hitSurface.roughness;
+                let hitRoughnessAtten = hitSmoothness * hitSmoothness;
+                let opaqueF0 = computeF0Vec(hitSurface.baseColor, hitSurface.metallic);
+                let opaqueNdotV = max(dot(hit.normal, -ray.direction), 0.0);
+                let opaqueF = schlickFresnel(opaqueNdotV, opaqueF0);
+                let nextThroughput = weight * opaqueF * hitRoughnessAtten;
+                let avgNext = (nextThroughput.x + nextThroughput.y + nextThroughput.z) / 3.0;
+
+                if (avgNext >= MIN_CONTRIBUTION) {
+                    tOrigin[taskCount] = hit.worldPos + hit.normal * REFLECTION_EPSILON;
+                    tDirection[taskCount] = reflect(ray.direction, hit.normal);
+                    tKind[taskCount] = KIND_REFLECT;
+                    tWeight[taskCount] = nextThroughput;
+                    tBounces[taskCount] = tBounces[i] - 1u;
+                    tRemaining[taskCount] = 1.0;
+                    taskCount++;
+                }
+            }
+        } else if (kind == KIND_BEHIND) {
+            if (!hit.hit) {
+                finalColor += weight * sampleSky(ray.direction);
+                continue;
+            }
+
+            let behind = getData(hit.entityId);
+            let bL = -scene.sunDirection.xyz;
+            let bNdotL = max(dot(hit.normal, bL), 0.0);
+            let bAmbient = scene.ambientColor.rgb * scene.ambientColor.a;
+            let bSun = scene.sunColor.rgb * bNdotL;
+            let bDiffuse = behind.baseColor.rgb * (bAmbient + bSun) * (1.0 - behind.pbr.y);
+            let bEmission = behind.emission.rgb * behind.emission.a;
+            finalColor += weight * applyHaze(bDiffuse + bEmission, hit.t);
+        }
+    }
+
+    textureStore(output_scene, vec2<i32>(gid.xy), vec4(finalColor, 1.0));
+    textureStore(output_depth, vec2<i32>(gid.xy), vec4(depth, 0.0, 0.0, 0.0));
+    textureStore(output_entityId, vec2<i32>(gid.xy), vec4(entityId, 0u, 0u, 0u));
+}
+`;
+}
+
+export function compileForwardShader(surface: SurfaceData): string {
+    return compileUberShader([surface]);
+}