@woosh/meep-engine/src/core/bvh2/bvh3/ebvh_build_hierarchy_radix.js

Pure JavaScript game engine. Fully featured and production ready.

import { assert } from "../../assert.js";
import { clamp } from "../../math/clamp.js";
import { NULL_NODE } from "./BVH.js";
import { ebvh_update_hierarchy_bounds } from "./ebvh_update_hierarchy_bounds.js";

/**
 *
 * @see "Thinking Parallel, Part III: Tree Construction on the GPU", 2012 by Tero Karras
 *
 * @param {number[]|Uint32Array} sortedMortonCodes
 * @param {number} first
 * @param {number} last
 * @return {number}
 */
function find_split(
    sortedMortonCodes,
    first,
    last
) {
    // Identical Morton codes => split the range in the middle.

    const firstCode = sortedMortonCodes[first];
    const lastCode = sortedMortonCodes[last];

    if (firstCode === lastCode) {
        return (first + last) >>> 1;
    }

    // Calculate the number of highest bits that are the same
    // for all objects, using the count-leading-zeros intrinsic.

    const commonPrefix = Math.clz32(firstCode ^ lastCode);

    // Use binary search to find where the next bit differs.
    // Specifically, we are looking for the highest object that
    // shares more than commonPrefix bits with the first one.

    let split = first; // initial guess
    let step = last - first;

    do {
        step = (step + 1) >>> 1; // exponential decrease
        const newSplit = split + step; // proposed new position

        if (newSplit < last) {
            const splitCode = sortedMortonCodes[newSplit];
            const splitPrefix = Math.clz32(firstCode ^ splitCode);
            if (splitPrefix > commonPrefix) {
                split = newSplit; // accept proposal
            }
        }
    }
    while (step > 1);

    return split;
}

/**
 *
 * @param {number} indexA
 * @param {number} indexB
 * @param {number} elementCount
 * @param {number[]} mortonCodes
 * @return {number}
 * @see https://github.com/turanszkij/WickedEngine/blob/506749de321c2ab66fd33fbe41efb95afbbb7ff8/WickedEngine/shaders/bvh_hierarchyCS.hlsl#L28C1-L48C2
 */
function GetLongestCommonPrefix(indexA, indexB, elementCount, mortonCodes) {
    assert.isNonNegativeInteger(indexA, 'indexA');
    assert.isNonNegativeInteger(indexB, 'indexB');

    if (indexA >= elementCount || indexB >= elementCount) {
        return -1;
    } else {
        const mortonCodeA = mortonCodes[indexA];
        const mortonCodeB = mortonCodes[indexB];
        if (mortonCodeA !== mortonCodeB) {
            return Math.clz32(mortonCodeA ^ mortonCodeB);
        } else {
            // TODO: Technically this should be primitive ID
            return Math.clz32(indexA ^ indexB) + 31;
        }
    }
}

/**
 * @see https://github.com/turanszkij/WickedEngine/blob/506749de321c2ab66fd33fbe41efb95afbbb7ff8/WickedEngine/shaders/bvh_hierarchyCS.hlsl#L28C1-L48C2
 * @param {number[]|Uint32Array} output
 * @param {number[]} sortedMortonCodes
 * @param {number} numTriangles
 * @param {number} idx
 */
function determineRangeW(output, sortedMortonCodes, numTriangles, idx) {
    if (idx === 0) {
        // root
        output[0] = 0;
        output[1] = numTriangles - 1;
        return;
    }

    let d = GetLongestCommonPrefix(idx, idx + 1, numTriangles, sortedMortonCodes) - GetLongestCommonPrefix(idx, idx - 1, numTriangles, sortedMortonCodes);
    d = clamp(d, -1, 1);

    let minPrefix = GetLongestCommonPrefix(idx, idx - d, numTriangles, sortedMortonCodes);

    // TODO: Consider starting this at a higher number
    let maxLength = 2;
    while (GetLongestCommonPrefix(idx, idx + maxLength * d, numTriangles, sortedMortonCodes) > minPrefix) {
        maxLength <<= 2;
    }

    let length = 0;
    for (let t = maxLength >>> 1; t > 0; t >>>= 1) {
        if (GetLongestCommonPrefix(idx, idx + (length + t) * d, numTriangles, sortedMortonCodes) > minPrefix) {
            length = length + t;
        }
    }

    let j = idx + length * d;

    output[0] = Math.min(idx, j);
    output[1] = Math.max(idx, j);

}


/**
 * @see https://github.com/mbartling/cuda-bvh/blob/7f2f98d9d29956c3559632e59104ba66f31f80b8/kernels/bvh.cu#L276C1-L352C2
 * @param {number[]|Uint32Array} output
 * @param {number[]} sortedMortonCodes
 * @param {number} numTriangles
 * @param {number} idx
 */
function determineRange(output, sortedMortonCodes, numTriangles, idx) {

    //determine the range of keys covered by each internal node (as well as its children)
    //direction is found by looking at the neighboring keys ki-1 , ki , ki+1
    //the index is either the beginning of the range or the end of the range
    let direction = 0;
    let common_prefix_with_left = 0;
    let common_prefix_with_right = 0;

    common_prefix_with_right = Math.clz32(sortedMortonCodes[idx] ^ sortedMortonCodes[idx + 1]);

    if (idx === 0) {
        common_prefix_with_left = -1;
    } else {
        common_prefix_with_left = Math.clz32(sortedMortonCodes[idx] ^ sortedMortonCodes[idx - 1]);

    }

    direction = ((common_prefix_with_right - common_prefix_with_left) > 0) ? 1 : -1;

    let min_prefix_range = 0;

    if (idx === 0) {
        min_prefix_range = -1;

    } else {
        min_prefix_range = Math.clz32(sortedMortonCodes[idx] ^ sortedMortonCodes[idx - direction]);
    }

    let lmax = 2;
    let next_key = idx + lmax * direction;

    while ((next_key >= 0) && (next_key < numTriangles) && (Math.clz32(sortedMortonCodes[idx] ^ sortedMortonCodes[next_key]) > min_prefix_range)) {
        lmax *= 2;
        next_key = idx + lmax * direction;
    }
    //find the other end using binary search
    let l = 0;

    do {
        lmax = (lmax + 1) >> 1; // exponential decrease
        const new_val = idx + (l + lmax) * direction;

        if (new_val >= 0 && new_val < numTriangles) {
            const Code = sortedMortonCodes[new_val];
            const Prefix = Math.clz32(sortedMortonCodes[idx] ^ Code);
            if (Prefix > min_prefix_range) {
                l = l + lmax;
            }
        }
    } while (lmax > 1);

    const j = idx + l * direction;

    let left = 0;
    let right = 0;

    if (idx < j) {
        left = idx;
        right = j;
    } else {
        left = j;
        right = idx;
    }

    // printf("idx : (%d) returning range (%d, %d) \n" , idx , left, right);

    output[0] = left
    output[1] = right;
}

/**
 * DO NOT USE, currently broken
 * @param {BVH} bvh
 * @param {number[]|Uint32Array} leaf_nodes
 * @param {number[]} sorted_morton_codes
 * @param {number} leaf_count
 * @param {number[]} internal_nodes
 * @returns {number} new root
 */
export function ebvh_build_hierarchy_radix(
    bvh,
    leaf_nodes,
    sorted_morton_codes,
    leaf_count,
    internal_nodes,
) {
    console.warn("IMPLEMENTATION IS INCOMPLETE, DO NOT USE")

    assert.isNonNegativeInteger(leaf_count, 'leaf_count');

    // Construct leaf nodes.
    // Note: This step can be avoided by storing
    // the tree in a slightly different way.

    const range = new Uint32Array(2);

    // Construct internal nodes.
    const intermediate_node_count = leaf_count - 1;
    for (let idx = 0; idx < intermediate_node_count; idx++) // in parallel
    {
        // Find out which range of objects the node corresponds to.
        // (This is where the magic happens!)

        determineRange(range, sorted_morton_codes, leaf_count, idx);

        const first = range[0];
        const last = range[1];

        // Determine where to split the range.

        const split = find_split(sorted_morton_codes, first, last);
        assert.isNonNegativeInteger(split, 'split');

        // Select childA.

        let childA;
        if (split === first) {
            childA = leaf_nodes[split];
        } else {
            childA = internal_nodes[split];
        }

        // Select childB.

        let childB;
        if (split + 1 === last) {
            childB = leaf_nodes[split + 1];
        } else {
            childB = internal_nodes[split + 1];
        }

        // Record parent-child relationships.
        const parent = internal_nodes[idx];

        bvh.node_assign_children_only(parent, childA, childB);
    }

    // Node 0 is the root.
    const root = internal_nodes[0];

    bvh.node_set_parent(root, NULL_NODE);

    // update bounds
    ebvh_update_hierarchy_bounds(bvh, root);

    return root;
}