@multiplekex/shallot 0.1.12 → 0.2.0

package/src/standard/render/bvh/traverse.ts

@@ -0,0 +1,467 @@
+import {
+    BVH_NODE_STRUCT_WGSL,
+    LEAF_FLAG_WGSL,
+    BLAS_NODE_STRUCT_WGSL,
+    BLAS_TRIANGLE_STRUCT_WGSL,
+} from "./structs";
+
+export const OCT_DECODE_WGSL = /* wgsl */ `
+fn octDecode(enc: u32) -> vec3<f32> {
+    let x = f32(enc & 0xFFFFu) / 65535.0 * 2.0 - 1.0;
+    let y = f32(enc >> 16u) / 65535.0 * 2.0 - 1.0;
+    let z = 1.0 - abs(x) - abs(y);
+    var n: vec3<f32>;
+    if (z < 0.0) {
+        let signX = select(-1.0, 1.0, x >= 0.0);
+        let signY = select(-1.0, 1.0, y >= 0.0);
+        n = vec3<f32>((1.0 - abs(y)) * signX, (1.0 - abs(x)) * signY, z);
+    } else {
+        n = vec3<f32>(x, y, z);
+    }
+    return normalize(n);
+}`;
+
+export const BVH_STRUCTS = /* wgsl */ `
+${BVH_NODE_STRUCT_WGSL}
+
+${LEAF_FLAG_WGSL}
+const INVALID_NODE: u32 = 0xFFFFFFFFu;
+const MAX_STACK_DEPTH: u32 = 24u;
+
+${OCT_DECODE_WGSL}
+`;
+
+export const TLAS_BLAS_STRUCTS = /* wgsl */ `
+${BLAS_NODE_STRUCT_WGSL}
+
+${BLAS_TRIANGLE_STRUCT_WGSL}
+`;
+
+export const TLAS_BLAS_BINDINGS = /* wgsl */ `
+@group(1) @binding(0) var<storage, read> tlasNodes: array<BVHNode>;
+@group(1) @binding(1) var<storage, read> tlasInstanceIds: array<u32>;
+@group(1) @binding(2) var<storage, read> blasNodes: array<BLASNode>;
+@group(1) @binding(3) var<storage, read> blasTriIds: array<u32>;
+@group(1) @binding(4) var<storage, read> blasTriangles: array<BLASTriangle>;
+@group(1) @binding(5) var<storage, read> blasMeta: array<u32>;
+@group(1) @binding(6) var<storage, read> instanceInverses: array<mat4x4<f32>>;
+`;
+
+export const TLAS_BLAS_TRAVERSAL = /* wgsl */ `
+fn isLeaf(child: u32) -> bool {
+    return (child & LEAF_FLAG) != 0u;
+}
+
+fn leafIndex(child: u32) -> u32 {
+    return child & ~LEAF_FLAG;
+}
+
+fn safeInverse(d: f32) -> f32 {
+    return select(1.0 / d, 1e30, abs(d) < 1e-10);
+}
+
+fn computeInvDir(dir: vec3<f32>) -> vec3<f32> {
+    return vec3(
+        safeInverse(dir.x),
+        safeInverse(dir.y),
+        safeInverse(dir.z)
+    );
+}
+
+fn intersectAABBDist(origin: vec3<f32>, invDir: vec3<f32>, nodeMin: vec3<f32>, nodeMax: vec3<f32>) -> f32 {
+    let t1 = (nodeMin - origin) * invDir;
+    let t2 = (nodeMax - origin) * invDir;
+
+    let tNear = min(t1, t2);
+    let tFar = max(t1, t2);
+
+    let tEnter = max(max(tNear.x, tNear.y), tNear.z);
+    let tExit = min(min(tFar.x, tFar.y), tFar.z);
+
+    if (tEnter <= tExit && tExit >= 0.0) {
+        return max(tEnter, 0.0);
+    }
+    return 1e30;
+}
+
+fn intersectBLASTriangle(ray: Ray, tri: BLASTriangle) -> HitResult {
+    var result: HitResult;
+    result.hit = false;
+    result.t = 0.0;
+    result.entityId = 0u;
+    result.u = 0.0;
+    result.v = 0.0;
+    result.normal = vec3(0.0, 1.0, 0.0);
+    result.worldPos = vec3(0.0);
+
+    let e1 = tri.e1;
+    let e2 = tri.e2;
+
+    let h = cross(ray.direction, e2);
+    let a = dot(e1, h);
+
+    if (a > -EPSILON && a < EPSILON) {
+        return result;
+    }
+
+    let f = 1.0 / a;
+    let s = ray.origin - tri.v0;
+    let u = f * dot(s, h);
+
+    if (u < 0.0 || u > 1.0) {
+        return result;
+    }
+
+    let q = cross(s, e1);
+    let v = f * dot(ray.direction, q);
+
+    if (v < 0.0 || u + v > 1.0) {
+        return result;
+    }
+
+    let t = f * dot(e2, q);
+
+    if (t > EPSILON) {
+        result.hit = true;
+        result.t = t;
+        result.u = u;
+        result.v = v;
+        let w = 1.0 - u - v;
+        let n0 = octDecode(tri.n0_enc);
+        let n1 = octDecode(tri.n1_enc);
+        let n2 = octDecode(tri.n2_enc);
+        result.normal = normalize(w * n0 + u * n1 + v * n2);
+        result.worldPos = ray.origin + t * ray.direction;
+    }
+
+    return result;
+}
+
+fn traceBLAS(
+    ray: Ray,
+    nodeOffset: u32,
+    triIdOffset: u32,
+    triOffset: u32,
+    triCount_: u32,
+    maxT: f32
+) -> HitResult {
+    var closest: HitResult;
+    closest.hit = false;
+    closest.t = maxT;
+    closest.entityId = 0u;
+    closest.u = 0.0;
+    closest.v = 0.0;
+    closest.normal = vec3(0.0, 1.0, 0.0);
+    closest.worldPos = vec3(0.0);
+
+    if (triCount_ == 0u) {
+        return closest;
+    }
+
+    if (triCount_ == 1u) {
+        let triIdx = blasTriIds[triIdOffset];
+        let tri = blasTriangles[triOffset + triIdx];
+        let hit = intersectBLASTriangle(ray, tri);
+        if (hit.hit && hit.t < maxT) {
+            return hit;
+        }
+        return closest;
+    }
+
+    let invDir = computeInvDir(ray.direction);
+
+    var stack: array<u32, MAX_STACK_DEPTH>;
+    var stackPtr = 0u;
+
+    stack[stackPtr] = 0u;
+    stackPtr++;
+
+    var iterations = 0u;
+    let maxIterations = min(triCount_ * 3u, 10000u);
+
+    while (stackPtr > 0u && iterations < maxIterations) {
+        iterations++;
+        stackPtr--;
+        let localIdx = stack[stackPtr];
+        let node = blasNodes[nodeOffset + localIdx];
+
+        let leftChild = node.leftChild;
+        let rightChild = node.rightChild;
+
+        if (leftChild != INVALID_NODE) {
+            if (isLeaf(leftChild)) {
+                let leafIdx = leafIndex(leftChild);
+                let triIdx = blasTriIds[triIdOffset + leafIdx];
+                let tri = blasTriangles[triOffset + triIdx];
+                let hit = intersectBLASTriangle(ray, tri);
+                if (hit.hit && hit.t < closest.t) {
+                    closest = hit;
+                }
+            } else {
+                let leftNode = blasNodes[nodeOffset + leftChild];
+                let leftMin = vec3(leftNode.minX, leftNode.minY, leftNode.minZ);
+                let leftMax = vec3(leftNode.maxX, leftNode.maxY, leftNode.maxZ);
+                let leftDist = intersectAABBDist(ray.origin, invDir, leftMin, leftMax);
+
+                if (leftDist < closest.t && stackPtr < MAX_STACK_DEPTH) {
+                    stack[stackPtr] = leftChild;
+                    stackPtr++;
+                }
+            }
+        }
+
+        if (rightChild != INVALID_NODE) {
+            if (isLeaf(rightChild)) {
+                let leafIdx = leafIndex(rightChild);
+                let triIdx = blasTriIds[triIdOffset + leafIdx];
+                let tri = blasTriangles[triOffset + triIdx];
+                let hit = intersectBLASTriangle(ray, tri);
+                if (hit.hit && hit.t < closest.t) {
+                    closest = hit;
+                }
+            } else {
+                let rightNode = blasNodes[nodeOffset + rightChild];
+                let rightMin = vec3(rightNode.minX, rightNode.minY, rightNode.minZ);
+                let rightMax = vec3(rightNode.maxX, rightNode.maxY, rightNode.maxZ);
+                let rightDist = intersectAABBDist(ray.origin, invDir, rightMin, rightMax);
+
+                if (rightDist < closest.t && stackPtr < MAX_STACK_DEPTH) {
+                    stack[stackPtr] = rightChild;
+                    stackPtr++;
+                }
+            }
+        }
+    }
+
+    return closest;
+}
+
+fn trace(ray: Ray) -> HitResult {
+    var closest: HitResult;
+    closest.hit = false;
+    closest.t = 1e30;
+    closest.entityId = 0u;
+    closest.u = 0.0;
+    closest.v = 0.0;
+    closest.normal = vec3(0.0, 1.0, 0.0);
+    closest.worldPos = vec3(0.0);
+
+    let count = getInstanceCount();
+    if (count == 0u) {
+        return closest;
+    }
+
+    let invDir = computeInvDir(ray.direction);
+
+    var stack: array<u32, MAX_STACK_DEPTH>;
+    var stackPtr = 0u;
+    stack[stackPtr] = 0u;
+    stackPtr++;
+
+    var iterations = 0u;
+    let maxIterations = min(count * 3u, 10000u);
+
+    while (stackPtr > 0u && iterations < maxIterations) {
+        iterations++;
+        stackPtr--;
+        let nodeIdx = stack[stackPtr];
+
+        let node = tlasNodes[nodeIdx];
+
+        var children: array<u32, 4>;
+        var dists: array<f32, 4>;
+
+        children[0] = node.child0;
+        children[1] = node.child1;
+        children[2] = node.child2;
+        children[3] = node.child3;
+
+        dists[0] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c0_minX, node.c0_minY, node.c0_minZ),
+                vec3(node.c0_maxX, node.c0_maxY, node.c0_maxZ)),
+            1e30,
+            children[0] == INVALID_NODE
+        );
+        dists[1] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c1_minX, node.c1_minY, node.c1_minZ),
+                vec3(node.c1_maxX, node.c1_maxY, node.c1_maxZ)),
+            1e30,
+            children[1] == INVALID_NODE
+        );
+        dists[2] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c2_minX, node.c2_minY, node.c2_minZ),
+                vec3(node.c2_maxX, node.c2_maxY, node.c2_maxZ)),
+            1e30,
+            children[2] == INVALID_NODE
+        );
+        dists[3] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c3_minX, node.c3_minY, node.c3_minZ),
+                vec3(node.c3_maxX, node.c3_maxY, node.c3_maxZ)),
+            1e30,
+            children[3] == INVALID_NODE
+        );
+
+        for (var i = 1u; i < 4u; i++) {
+            let keyDist = dists[i];
+            let keyChild = children[i];
+            var j = i;
+            while (j > 0u && dists[j - 1u] > keyDist) {
+                dists[j] = dists[j - 1u];
+                children[j] = children[j - 1u];
+                j--;
+            }
+            dists[j] = keyDist;
+            children[j] = keyChild;
+        }
+
+        for (var i = 3i; i >= 0i; i--) {
+            let child = children[i];
+            let dist = dists[i];
+
+            if (child == INVALID_NODE || dist >= closest.t) {
+                continue;
+            }
+
+            if (isLeaf(child)) {
+                let instanceIdx = leafIndex(child);
+                let eid = tlasInstanceIds[instanceIdx];
+                let shapeId = getShapeId(eid);
+
+                let nodeOffset = blasMeta[shapeId * 4u];
+                let triIdOffset = blasMeta[shapeId * 4u + 1u];
+                let triOffset = blasMeta[shapeId * 4u + 2u];
+                let triCount_ = blasMeta[shapeId * 4u + 3u];
+
+                if (triCount_ == 0u) {
+                    continue;
+                }
+
+                let invMatrix = instanceInverses[eid];
+                var objRay: Ray;
+                objRay.origin = (invMatrix * vec4(ray.origin, 1.0)).xyz;
+                objRay.direction = (invMatrix * vec4(ray.direction, 0.0)).xyz;
+
+                let blasHit = traceBLAS(objRay, nodeOffset, triIdOffset, triOffset, triCount_, closest.t);
+
+                if (blasHit.hit && blasHit.t < closest.t) {
+                    closest = blasHit;
+                    closest.entityId = eid;
+                    let normalMat = mat3x3(invMatrix[0].xyz, invMatrix[1].xyz, invMatrix[2].xyz);
+                    closest.normal = normalize(transpose(normalMat) * blasHit.normal);
+                    closest.worldPos = ray.origin + blasHit.t * ray.direction;
+                }
+            } else if (stackPtr < MAX_STACK_DEPTH) {
+                stack[stackPtr] = child;
+                stackPtr++;
+            }
+        }
+    }
+
+    return closest;
+}
+`;
+
+export const TLAS_BLAS_ANY_HIT = /* wgsl */ `
+fn traceAnyHit(ray: Ray, tMax: f32) -> bool {
+    let count = getInstanceCount();
+    if (count == 0u) { return false; }
+
+    let invDir = computeInvDir(ray.direction);
+
+    var stack: array<u32, MAX_STACK_DEPTH>;
+    var stackPtr = 0u;
+    stack[stackPtr] = 0u;
+    stackPtr++;
+
+    var iterations = 0u;
+    let maxIterations = min(count * 3u, 10000u);
+
+    while (stackPtr > 0u && iterations < maxIterations) {
+        iterations++;
+        stackPtr--;
+        let nodeIdx = stack[stackPtr];
+
+        let node = tlasNodes[nodeIdx];
+
+        var children: array<u32, 4>;
+        var dists: array<f32, 4>;
+
+        children[0] = node.child0;
+        children[1] = node.child1;
+        children[2] = node.child2;
+        children[3] = node.child3;
+
+        dists[0] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c0_minX, node.c0_minY, node.c0_minZ),
+                vec3(node.c0_maxX, node.c0_maxY, node.c0_maxZ)),
+            1e30,
+            children[0] == INVALID_NODE
+        );
+        dists[1] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c1_minX, node.c1_minY, node.c1_minZ),
+                vec3(node.c1_maxX, node.c1_maxY, node.c1_maxZ)),
+            1e30,
+            children[1] == INVALID_NODE
+        );
+        dists[2] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c2_minX, node.c2_minY, node.c2_minZ),
+                vec3(node.c2_maxX, node.c2_maxY, node.c2_maxZ)),
+            1e30,
+            children[2] == INVALID_NODE
+        );
+        dists[3] = select(
+            intersectAABBDist(ray.origin, invDir,
+                vec3(node.c3_minX, node.c3_minY, node.c3_minZ),
+                vec3(node.c3_maxX, node.c3_maxY, node.c3_maxZ)),
+            1e30,
+            children[3] == INVALID_NODE
+        );
+
+        for (var i = 0u; i < 4u; i++) {
+            let child = children[i];
+            let dist = dists[i];
+
+            if (child == INVALID_NODE || dist >= tMax) {
+                continue;
+            }
+
+            if (isLeaf(child)) {
+                let instanceIdx = leafIndex(child);
+                let eid = tlasInstanceIds[instanceIdx];
+                let shapeId = getShapeId(eid);
+
+                let nodeOffset = blasMeta[shapeId * 4u];
+                let triIdOffset = blasMeta[shapeId * 4u + 1u];
+                let triOffset = blasMeta[shapeId * 4u + 2u];
+                let triCount_ = blasMeta[shapeId * 4u + 3u];
+
+                if (triCount_ == 0u) {
+                    continue;
+                }
+
+                let invMatrix = instanceInverses[eid];
+                var objRay: Ray;
+                objRay.origin = (invMatrix * vec4(ray.origin, 1.0)).xyz;
+                objRay.direction = (invMatrix * vec4(ray.direction, 0.0)).xyz;
+
+                let blasHit = traceBLAS(objRay, nodeOffset, triIdOffset, triOffset, triCount_, tMax);
+                if (blasHit.hit && blasHit.t < tMax) {
+                    return true;
+                }
+            } else if (stackPtr < MAX_STACK_DEPTH) {
+                stack[stackPtr] = child;
+                stackPtr++;
+            }
+        }
+    }
+
+    return false;
+}
+`;