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

Pure JavaScript game engine. Fully featured and production ready.

import { assert } from "../../assert.js";
import { lsb_32 } from "../../binary/lsb_32.js";
import { bitCount } from "../../binary/operations/bitCount.js";
import { ELEMENT_WORD_COUNT, NULL_NODE } from "./BVH.js";

// Common continuous memory region for better cache performance
const scratch_memory = new ArrayBuffer((256 * 3 + 16) * 4 + 256 * 6 * 4);

const scratch_bounds = new Float32Array(scratch_memory, 0, 256 * 6);
const scratch_areas = new Float32Array(scratch_memory, 6144, 256);
const scratch_cost = new Float32Array(scratch_memory, 7168, 256);
const scratch_partitions = new Uint32Array(scratch_memory, 8192, 256);
const scratch_treelet = new Uint32Array(scratch_memory, 9216, 16);

/**
 *
 * @param {BVH} bvh
 * @param {number[]} leaves
 * @param {number} mask
 */
function assemble_mask(bvh, leaves, mask) {

    const leaf_count = bitCount(mask);

    if (leaf_count === 1) {
        const leaf_index = lsb_32(mask);
        return leaves[leaf_index];
    }

    let partitioning_left = scratch_partitions[mask];
    let partition_right = mask & ~(partitioning_left);

    const parent = bvh.allocate_node();

    const left_node = assemble_mask(bvh, leaves, partitioning_left);
    const right_node = assemble_mask(bvh, leaves, partition_right);

    bvh.node_assign_children(parent, left_node, right_node);

    return parent;
}

/**
 *
 * @param {BVH} bvh
 * @param {number} root
 * @param {number[]} leaves
 * @param {number} mask
 */
function rebuild_treelet(bvh, root, leaves, mask) {

    let partitioning_left = scratch_partitions[mask];
    let partition_right = mask & ~(partitioning_left);

    bvh.node_assign_children(root,
        assemble_mask(bvh, leaves, partitioning_left),
        assemble_mask(bvh, leaves, partition_right),
    );
}

//cost of traversing an intermediate node
const SAH_COST_INTERMEDIATE = 1.2;
// cost of performing ray/triangle test
const SAH_COST_TRIANGLE = 1;

/**
 *
 * @param {BVH} bvh
 * @param {number} root
 * @param {number[]|Uint32Array} leaves
 * @param {number} leaf_count
 */
function optimize_treelet(bvh, root, leaves, leaf_count) {
    // console.log(`TREELET(${leaf_count}): ${leaves.slice(0, leaf_count).join(', ')}`); // DEBUG

    // see "Fast Parallel Construction of High-Quality Bounding Volume Hierarchies", Algorithm 2

    const bvh_float32 = bvh.data_float32;

    // Calculate surface area for each subset
    const mask_limit = (1 << leaf_count) - 1;

    for (let s = 1; s <= mask_limit; s++) {
        let min_x, min_y, min_z, max_x, max_y, max_z;

        // Check if 's' is a power of 2 (i.e., a single leaf)
        // (s & (s-1)) === 0 is the standard trick for this
        const is_a_leaf = (s & (s - 1)) === 0;

        if (is_a_leaf) {
            // --- CASE A: LEAF ---
            // Map the mask bit back to a linear index (0..7)
            const leaf_index = lsb_32(s);

            // Read from your compact linear array
            const node = leaves[leaf_index];

            const source_offset = node * ELEMENT_WORD_COUNT;

            min_x = bvh_float32[source_offset];
            min_y = bvh_float32[source_offset + 1];
            min_z = bvh_float32[source_offset + 2];
            max_x = bvh_float32[source_offset + 3];
            max_y = bvh_float32[source_offset + 4];
            max_z = bvh_float32[source_offset + 5];

            // Initialize base cost for leaf (SAH_COST_TRIANGLE * Area)
            // (Calculated below)
        } else {
            // --- CASE B: SUBSET ---
            // Split s into: LSB (Leaf) + Rest (Subset)
            // Both of these are < s, so they are guaranteed to be initialized.
            const mask_lsb = s & -s;       // Extract lowest set bit
            const mask_rest = s ^ mask_lsb; // The rest of the bits

            const i_lsb = mask_lsb * 6;
            const i_rest = mask_rest * 6;

            // Merge bounds
            min_x = Math.min(scratch_bounds[i_lsb],     scratch_bounds[i_rest]);
            min_y = Math.min(scratch_bounds[i_lsb + 1], scratch_bounds[i_rest + 1]);
            min_z = Math.min(scratch_bounds[i_lsb + 2], scratch_bounds[i_rest + 2]);
            max_x = Math.max(scratch_bounds[i_lsb + 3], scratch_bounds[i_rest + 3]);
            max_y = Math.max(scratch_bounds[i_lsb + 4], scratch_bounds[i_rest + 4]);
            max_z = Math.max(scratch_bounds[i_lsb + 5], scratch_bounds[i_rest + 5]);
        }

        // --- STORE ---
        const dest_offset = s * 6;

        scratch_bounds[dest_offset]     = min_x;
        scratch_bounds[dest_offset + 1] = min_y;
        scratch_bounds[dest_offset + 2] = min_z;
        scratch_bounds[dest_offset + 3] = max_x;
        scratch_bounds[dest_offset + 4] = max_y;
        scratch_bounds[dest_offset + 5] = max_z;

        // Calculate Surface Area
        const dx = max_x - min_x;
        const dy = max_y - min_y;
        const dz = max_z - min_z;

        const area = 2.0 * (dx * dy + dy * dz + dz * dx);

        scratch_areas[s] = area;

        if(is_a_leaf){
            // Initialize costs of individual leaves
            scratch_cost[s] = SAH_COST_TRIANGLE * area;
        }
    }

    // Optimize every subset of leaves
    for (let k = 2; k <= leaf_count; k++) {
        for (let s = 1; s <= mask_limit; s++) {
            if (bitCount(s) !== k) {
                continue;
            }

            // Try each way of partitioning the leaves
            let best_cost = Infinity;
            let best_partition = 0;

            const delta = (s - 1) & s;
            let partition = (-delta) & s;

            do {
                const cost = scratch_cost[partition] + scratch_cost[s ^ partition];

                if (cost < best_cost) {
                    best_cost = cost;
                    best_partition = partition;
                }

                partition = (partition - delta) & s;

            } while (partition !== 0)

            //  Calculate final SAH cost (Equation 2)
            const sah_heuristic = SAH_COST_INTERMEDIATE * scratch_areas[s] + best_cost

            scratch_cost[s] = sah_heuristic;
            scratch_partitions[s] = best_partition;
        }
    }

    // build tree
    rebuild_treelet(bvh, root, leaves, (1 << leaf_count) - 1);
}


/**
 *
 * @param {BVH} bvh
 * @param {number} root
 * @param {number} treelet_max_size
 */
function optimize_at_node(bvh, root, treelet_max_size) {
    if (root === NULL_NODE) {
        return;
    }

    if (bvh.node_is_leaf(root)) {
        return;
    }

    const left = bvh.node_get_child1(root);
    const right = bvh.node_get_child2(root);

    let treelet_size = 0;

    if (left !== NULL_NODE) {
        scratch_treelet[treelet_size++] = left;
    }

    if (right !== NULL_NODE) {
        scratch_treelet[treelet_size++] = right;
    }

    if (treelet_size === 0) {
        return;
    }

    while (treelet_size < treelet_max_size) {
        // pick node with the largest surface area
        let best_node_index = -1;
        let best_area = 0; //note that this is deliberate, expanding node with 0 area is pointless

        for (let i = 0; i < treelet_size; i++) {
            const n = scratch_treelet[i];

            if (bvh.node_is_leaf(n)) {
                // can't expand
                continue;
            }

            // TODO maintain list of areas instead of computing them every time
            const area = bvh.node_get_surface_area(n);
            if (area > best_area) {
                best_area = area;
                best_node_index = i;
            }
        }

        if (best_node_index === -1) {
            break;
        }

        // expand candidate
        const victim = scratch_treelet[best_node_index];

        const c0 = bvh.node_get_child1(victim);
        const c1 = bvh.node_get_child2(victim);

        scratch_treelet[best_node_index] = c0;
        scratch_treelet[treelet_size++] = c1;

        bvh.release_node(victim);
    }

    if (treelet_size <= 2) {
        // nothing to do
        return;
    }

    optimize_treelet(bvh, root, scratch_treelet, treelet_size);
}

const CAME_FROM_TYPE_PARENT = 0;
const CAME_FROM_TYPE_CHILD_1 = 1;
const CAME_FROM_TYPE_CHILD_2 = 2;

/**
 * Perform linear optimization across the tree, similar to rotations, but we support somewhat arbitrary depth and all possible permutations
 * Note that the root does not move, but everything below the root is subject to change
 * @see "Fast Parallel Construction of High-Quality Bounding Volume Hierarchies" by Kerras, NVIDIA 2013
 * @param {BVH} bvh
 * @param {number} [root] where to start optimization, only nodes below this one will be considered
 * @param {number} [treelet_size] the larger the size - the more optimization opportunities, but it gets exponentially slower
 */
export function ebvh_optimize_treelet(
    bvh,
    root = bvh.root,
    treelet_size = 7
) {
    assert.isNonNegativeInteger(treelet_size, 'treelet_size');
    assert.lessThanOrEqual(treelet_size, 8, 'limit is 8');

    if (root === NULL_NODE) {
        // special case
        return;
    }

    if (bvh.node_is_leaf(root)) {
        // special case
        return;
    }

    // we traverse depth-first using stack-less scheme
    let came_from_node = root;

    let came_from_type = CAME_FROM_TYPE_PARENT;

    let node = bvh.node_get_child1(came_from_node);
    if (node === NULL_NODE) {
        node = bvh.node_get_child2(came_from_node);
    }

    if (node === NULL_NODE) {
        // no children, done
        return;
    }

    const min_height = Math.ceil(Math.log2(treelet_size));

    // do a stackless traversal
    while (true) {


        if (
            came_from_type === CAME_FROM_TYPE_CHILD_2
            && bvh.node_get_height(node) >= min_height
        ) {
            // only form treelets when moving up
            optimize_at_node(bvh, node, treelet_size);
        }

        // only add to treelet when traversing up, not down

        if (bvh.node_is_leaf(node)) {

            assert.equal(came_from_type, CAME_FROM_TYPE_PARENT);

            const parent = came_from_node;
            const parent_child_1 = bvh.node_get_child1(parent);

            if (parent_child_1 === node) {
                came_from_type = CAME_FROM_TYPE_CHILD_1;
            } else {
                came_from_type = CAME_FROM_TYPE_CHILD_2;
            }

            came_from_node = node;

            node = parent;
        } else if (came_from_type === CAME_FROM_TYPE_CHILD_2) {
            // finishing traversal of this branch
            came_from_node = node;
            const parent = bvh.node_get_parent(node);

            if (bvh.node_get_child1(parent) === node) {
                came_from_type = CAME_FROM_TYPE_CHILD_1;
            } else {
                came_from_type = CAME_FROM_TYPE_CHILD_2;
            }

            if (node === root) {
                // traversed all the way back up
                break;
            }

            node = parent;

        } else if (came_from_type === CAME_FROM_TYPE_CHILD_1) {
            // traversing up from left child
            const child2 = bvh.node_get_child2(node);

            came_from_node = node;

            if (child2 === NULL_NODE) {
                // no right child, finish this branch
                came_from_type = CAME_FROM_TYPE_CHILD_2; // indicate end

                node = bvh.node_get_parent(node);
            } else {
                came_from_type = CAME_FROM_TYPE_PARENT;
                node = child2;
            }

        } else if (came_from_type === CAME_FROM_TYPE_PARENT) {
            const child1 = bvh.node_get_child1(node);
            const child2 = bvh.node_get_child2(node);

            if (child1 !== NULL_NODE) {
                came_from_type = CAME_FROM_TYPE_PARENT;
                came_from_node = node;
                node = child1;
            } else if (child2 !== NULL_NODE) {
                came_from_type = CAME_FROM_TYPE_PARENT;
                came_from_node = node;
                node = child2;
            } else {
                // this should not happen, as this would mean that node is empty and completely pointless
                // we remove the node and traverse up


                const parent_child1 = bvh.node_get_child1(came_from_node);

                bvh.release_node(node);

                if (parent_child1 === node) {
                    // we are left child
                    bvh.node_set_child1(came_from_node, NULL_NODE);
                    came_from_node = NULL_NODE;
                    came_from_type = CAME_FROM_TYPE_CHILD_1;
                } else {
                    bvh.node_set_child2(came_from_node, NULL_NODE);
                    came_from_node = NULL_NODE;
                    came_from_type = CAME_FROM_TYPE_CHILD_2;
                }

                node = came_from_node;
            }


        }

    }

    // optimize last collected treelet
    optimize_at_node(bvh, root, treelet_size);
}