@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 { NULL_NODE } from "./BVH.js";

/**
 * Karras 2012 — "Maximizing Parallelism in the Construction of BVHs, Octrees, and k-d Trees".
 *
 * For a sorted array of N Morton codes, build a binary radix tree of N-1 internal
 * nodes and N leaves. Each internal node corresponds to the longest common prefix
 * of a contiguous range of Morton codes; its split point is where the first bit
 * after that prefix flips from 0 to 1.
 *
 * The caller must pre-allocate:
 *   - `leaf_count` leaf nodes in `bvh`, each with its AABB and user-data set;
 *     their IDs are passed in `leaf_nodes`, sorted by `sorted_morton_codes`.
 *   - `leaf_count - 1` internal nodes (when `leaf_count >= 2`); their IDs are
 *     passed in `internal_nodes`, in arbitrary order.
 *
 * The returned root index is wired into the tree; internal-node bounds and
 * heights are filled in by a bottom-up refit after structural construction.
 *
 * @see "Thinking Parallel, Part III: Tree Construction on the GPU", 2012 by Tero Karras
 */

/**
 * Compute the position of the first bit where the codes at `first` and `last`
 * differ in their longest-common-prefix sense — i.e. the split point for the
 * range [first, last]. Returns an index `s` such that all codes in [first, s]
 * share more leading bits with the firstCode than codes in [s+1, last] do.
 *
 * @param {number[]|Uint32Array} sortedMortonCodes
 * @param {number} first
 * @param {number} last
 * @return {number}
 */
function find_split(sortedMortonCodes, first, last) {
    const firstCode = sortedMortonCodes[first];
    const lastCode = sortedMortonCodes[last];

    // Identical Morton codes (multiple leaves at the same point): split the
    // range in the middle so we still build a balanced subtree.
    if (firstCode === lastCode) {
        return (first + last) >>> 1;
    }

    // Number of highest bits that match between first and last.
    const commonPrefix = Math.clz32(firstCode ^ lastCode);

    // Binary search for the highest index that still shares more than
    // `commonPrefix` bits with the firstCode.
    let split = first;
    let step = last - first;

    do {
        step = (step + 1) >>> 1;
        const newSplit = split + step;

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

    return split;
}

/**
 * Determine the [first, last] code range covered by internal node `idx`.
 *
 * @see https://github.com/mbartling/cuda-bvh/blob/7f2f98d9d29956c3559632e59104ba66f31f80b8/kernels/bvh.cu#L276
 * @param {number[]|Uint32Array} output 2-element output [first, last]
 * @param {number[]|Uint32Array} sortedMortonCodes
 * @param {number} leaf_count
 * @param {number} idx
 */
function determineRange(output, sortedMortonCodes, leaf_count, idx) {
    // For the root (idx === 0) the range is unambiguously [0, leaf_count - 1].
    // We short-circuit so callers don't have to special-case it.
    if (idx === 0) {
        output[0] = 0;
        output[1] = leaf_count - 1;
        return;
    }

    const codeIdx = sortedMortonCodes[idx];
    const prefixRight = Math.clz32(codeIdx ^ sortedMortonCodes[idx + 1]);
    const prefixLeft  = Math.clz32(codeIdx ^ sortedMortonCodes[idx - 1]);

    const direction = (prefixRight - prefixLeft) > 0 ? 1 : -1;

    // Minimum prefix length we require to stay in this range — anything that
    // shares fewer bits than this is "outside" the node's subtree.
    const min_prefix_range = Math.clz32(codeIdx ^ sortedMortonCodes[idx - direction]);

    // Expand maximum range with exponential growth, bounded by array.
    let lmax = 2;
    let next_key = idx + lmax * direction;

    while (
        next_key >= 0 &&
        next_key < leaf_count &&
        Math.clz32(codeIdx ^ sortedMortonCodes[next_key]) > min_prefix_range
    ) {
        lmax *= 2;
        next_key = idx + lmax * direction;
    }

    // Binary-search for the actual length within [0, lmax].
    let length = 0;
    do {
        lmax = (lmax + 1) >> 1;
        const new_val = idx + (length + lmax) * direction;

        if (new_val >= 0 && new_val < leaf_count) {
            const prefix = Math.clz32(codeIdx ^ sortedMortonCodes[new_val]);
            if (prefix > min_prefix_range) {
                length += lmax;
            }
        }
    } while (lmax > 1);

    const j = idx + length * direction;

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

/**
 * Bottom-up bounds and height refit for a radix-built tree.
 *
 * Walks the tree once in pre-order to collect every reachable node, then
 * processes them in reverse — which is post-order for trees — combining
 * each internal node's AABB from its children and updating its height.
 *
 * @param {BVH} bvh
 * @param {number} root
 */
function refit_radix_tree(bvh, root) {
    if (root === NULL_NODE) {
        return;
    }

    const order = [];
    const stack = [root];

    while (stack.length > 0) {
        const node = stack.pop();
        order.push(node);

        if (!bvh.node_is_leaf(node)) {
            stack.push(bvh.node_get_child1(node));
            stack.push(bvh.node_get_child2(node));
        }
    }

    for (let i = order.length - 1; i >= 0; i--) {
        const node = order[i];
        if (bvh.node_is_leaf(node)) {
            continue;
        }
        const c1 = bvh.node_get_child1(node);
        const c2 = bvh.node_get_child2(node);
        bvh.node_set_combined_aabb(node, c1, c2);
        bvh.node_set_height(
            node,
            1 + Math.max(bvh.node_get_height(c1), bvh.node_get_height(c2))
        );
    }
}

/**
 * Build a binary radix tree over the given leaves and Morton codes.
 *
 * @param {BVH} bvh
 * @param {number[]|Uint32Array} leaf_nodes
 *        Pre-allocated leaf node IDs, ordered to match `sorted_morton_codes`
 *        (i.e. leaf_nodes[i]'s code is sorted_morton_codes[i]).
 * @param {number[]|Uint32Array} sorted_morton_codes
 *        Morton codes in non-decreasing order. May contain duplicates.
 * @param {number} leaf_count
 *        Number of leaves (must equal the length of `leaf_nodes` and
 *        `sorted_morton_codes` that the function will actually look at).
 * @param {number[]|Uint32Array} internal_nodes
 *        Pre-allocated internal node IDs, of length at least `leaf_count - 1`.
 *        Ignored when leaf_count < 2.
 * @returns {number} The root node id, or NULL_NODE if leaf_count is 0.
 */
export function ebvh_build_hierarchy_radix(
    bvh,
    leaf_nodes,
    sorted_morton_codes,
    leaf_count,
    internal_nodes
) {
    assert.isNonNegativeInteger(leaf_count, 'leaf_count');

    // Empty input: nothing to build.
    if (leaf_count === 0) {
        return NULL_NODE;
    }

    // Single leaf is its own root; no internal node needed.
    if (leaf_count === 1) {
        const root = leaf_nodes[0];
        bvh.node_set_parent(root, NULL_NODE);
        return root;
    }

    // General case: build leaf_count - 1 internal nodes per Karras 2012.
    const range = new Uint32Array(2);
    const internal_node_count = leaf_count - 1;

    for (let idx = 0; idx < internal_node_count; idx++) {
        determineRange(range, sorted_morton_codes, leaf_count, idx);

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

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

        // child A: the [first, split] subrange's owner. If that subrange is a
        // single element it's a leaf; otherwise it's the internal node at
        // index `split` (Karras's convention: internal node k owns the range
        // whose left boundary is at code-index k).
        const childA = (split === first)
            ? leaf_nodes[split]
            : internal_nodes[split];

        // child B: the [split+1, last] subrange's owner.
        const childB = (split + 1 === last)
            ? leaf_nodes[split + 1]
            : internal_nodes[split + 1];

        const parent = internal_nodes[idx];
        bvh.node_assign_children_only(parent, childA, childB);
    }

    // Internal node 0 covers [0, leaf_count - 1] — the whole tree — and is
    // therefore the root.
    const root = internal_nodes[0];
    bvh.node_set_parent(root, NULL_NODE);

    // Walk back up to fill in correct AABBs and heights for every internal
    // node. Heights set during construction by node_assign_children_only are
    // stale because children may not have been built yet at parent-creation
    // time.
    refit_radix_tree(bvh, root);

    return root;
}