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

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

const scratch_areas = new Float32Array(scratch_memory, 0, 256);
const scratch_cost = new Float32Array(scratch_memory, 1024, 256);
const scratch_partitions = new Uint32Array(scratch_memory, 2048, 256);
const scratch_treelet = new Uint32Array(scratch_memory, 3072, 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

    // Calculate surface area for each subset
    for (let i = 0; i < leaf_count; i++) {
        const leaf_0 = leaves[i];
        for (let j = i + 1; j < leaf_count; j++) {
            const leaf_1 = leaves[j];

            const s = (1 << i) | (1 << j);

            scratch_areas[s] = bvh.node_get_combined_surface_area(leaf_0, leaf_1)
        }
    }

    // Initialize costs of individual leaves
    for (let i = 0; i < leaf_count; i++) {
        const leaf = leaves[i];

        scratch_cost[1 << i] = SAH_COST_TRIANGLE * bvh.node_get_surface_area(leaf);
    }

    // Optimize every subset of leaves
    for (let k = 2; k <= leaf_count; k++) {
        for (let s = 1; s <= (1 << leaf_count) - 1; 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_child1(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);
}