@woosh/meep-engine/src/core/bvh2/binary/2/BinaryUint32BVH.js

Pure JavaScript game engine. Fully featured and production ready.

import { assert } from "../../../assert.js";
import { ceilPowerOfTwo } from "../../../binary/operations/ceilPowerOfTwo.js";
import { array_copy } from "../../../collection/array/array_copy.js";
import { array_swap } from "../../../collection/array/array_swap.js";
import { SCRATCH_UINT32_TRAVERSAL_STACK } from "../../../collection/SCRATCH_UINT32_TRAVERSAL_STACK.js";
import { aabb3_array_combine } from "../../../geom/3d/aabb/aabb3_array_combine.js";
import { aabb3_array_set } from "../../../geom/3d/aabb/aabb3_array_set.js";
import { aabb3_compute_half_surface_area } from "../../../geom/3d/aabb/aabb3_compute_half_surface_area.js";
import { v3_morton_encode_bounded } from "../../../geom/3d/morton/v3_morton_encode_bounded.js";

/**
 * @readonly
 * @type {number}
 */
export const BVH_BOX_BYTE_SIZE = 6;

/**
 * In words (4 byte)
 * @readonly
 * @type {number}
 */
export const BVH_BINARY_NODE_SIZE = 6;
/**
 * @readonly
 * @type {number}
 */
export const BVH_LEAF_NODE_SIZE = 7;

/**
 *
 * @param {Float32Array} data
 * @param {number} destination
 * @param {number} source
 */
function copy_box_zero_size(data, destination, source) {
    assert.isNonNegativeInteger(destination, 'destination');
    assert.isNonNegativeInteger(source, 'source');

    const x = data[source];
    const y = data[source + 1];
    const z = data[source + 2];

    assert.notNaN(x, 'x');
    assert.notNaN(y, 'y');
    assert.notNaN(z, 'z');

    aabb3_array_set(data, destination, x, y, z, x, y, z);
}

/**
 * Assumes data will be normalized to 0...1 value range
 * @param {Float32Array} data
 * @param {number} address
 * @param {number[]} bounds
 * @returns {number}
 */
function build_morton(data, address, bounds) {

    const x0 = data[address];
    const y0 = data[address + 1];
    const z0 = data[address + 2];
    const x1 = data[address + 3];
    const y1 = data[address + 4];
    const z1 = data[address + 5];

    const cx = (x0 + x1) * 0.5;
    const cy = (y0 + y1) * 0.5;
    const cz = (z0 + z1) * 0.5;

    return v3_morton_encode_bounded(cx, cy, cz, bounds);

}

const stack = SCRATCH_UINT32_TRAVERSAL_STACK;

/**
 * Memory-efficient LBVH implementation.
 * LBVH is fast to build, is quite fast to query and has good memory usage due to implicit addressing.
 * LBVH is static, so it's not suitable for dynamic usecases, if your usecase requires updates to the BVH - use {@link BVH} instead, which is a fully dynamic BVH implementation.
 *
 * @see https://en.wikipedia.org/wiki/Bounding_volume_hierarchy
 * @see BVH
 */
export class BinaryUint32BVH {
    /**
     *
     * @private
     * @type {ArrayBuffer}
     */
    __data_buffer;

    /**
     * @readonly
     * @type {Float32Array}
     * @private
     */
    __data_float32;

    /**
     * @readonly
     * @private
     * @type {Uint32Array}
     */
    __data_uint32;

    /**
     *
     * @type {number}
     * @private
     */
    __node_count_binary = 0;

    /**
     *
     * @type {number}
     * @private
     */
    __node_count_leaf = 0;

    constructor() {
        this.data = new ArrayBuffer(320);
    }

    /**
     * In bytes
     * @returns {number}
     */
    estimateByteSize() {
        return this.data.byteLength + 248;
    }

    getTotalBoxCount() {
        return this.__node_count_binary + this.__node_count_leaf;
    }

    get binary_node_count() {
        return this.__node_count_binary;
    }

    get leaf_node_count() {
        return this.__node_count_leaf;
    }

    /**
     *
     * @returns {number}
     */
    getLeafBlockAddress() {
        return this.__node_count_binary * BVH_BINARY_NODE_SIZE;
    }

    get float32() {
        return this.__data_float32;
    }

    get uint32() {
        return this.__data_uint32;
    }

    /**
     *
     * @param {ArrayBuffer} buffer
     */
    set data(buffer) {
        assert.defined(buffer, 'buffer');

        this.__data_buffer = buffer;

        this.__data_float32 = new Float32Array(this.__data_buffer);
        this.__data_uint32 = new Uint32Array(this.__data_buffer);
    }

    get data() {
        return this.__data_buffer;
    }

    /**
     * Resolve index of the node to address where the node data starts, this is required to know where AABB is stored in memory
     * @param {number} node_index
     * @returns {number}
     */
    getNodeAddress(node_index) {
        const binary_node_count = this.__node_count_binary;
        const leaf_node_index = node_index - binary_node_count;

        if (leaf_node_index < 0) {
            // binary node
            return node_index * BVH_BINARY_NODE_SIZE;
        } else {
            // leaf node
            return binary_node_count * BVH_BINARY_NODE_SIZE + leaf_node_index * BVH_LEAF_NODE_SIZE;
        }
    }

    initialize_structure() {
        // compute memory requirements
        const word_count = this.__node_count_binary * BVH_BINARY_NODE_SIZE + this.__node_count_leaf * BVH_LEAF_NODE_SIZE;
        const storage_size = word_count * 4;

        // possibly resize the storage
        if (this.__data_buffer.byteLength < storage_size) {
            this.data = new ArrayBuffer(storage_size);
        }
    }

    /**
     *
     * @param {number} count
     */
    setLeafCount(count) {
        assert.isNonNegativeInteger(count, 'count');

        this.__node_count_leaf = count;

        const twoLeafLimit = ceilPowerOfTwo(count);

        if (count <= 1) {
            // special case
            this.__node_count_binary = twoLeafLimit;
        } else {
            this.__node_count_binary = twoLeafLimit - 1;
        }

    }

    /**
     *
     * @param {number} index
     * @return {number}
     */
    getLeafAddress(index) {
        assert.isNonNegativeInteger(index, 'index');
        assert.lessThan(index, this.__node_count_leaf, 'leaf index overflow');

        const leaf_block_address = this.__node_count_binary * BVH_BINARY_NODE_SIZE;

        return index * BVH_LEAF_NODE_SIZE + leaf_block_address;
    }

    /**
     *
     * @param {number} index
     * @param {number} payload
     * @param {number} x0
     * @param {number} y0
     * @param {number} z0
     * @param {number} x1
     * @param {number} y1
     * @param {number} z1
     */
    setLeafData(
        index, payload,
        x0, y0, z0,
        x1, y1, z1
    ) {
        assert.isNonNegativeInteger(index, 'index');
        assert.lessThan(index, this.__node_count_leaf, 'leaf index overflow');

        assert.isNumber(x0, 'x0');
        assert.isNumber(y0, 'y0');
        assert.isNumber(z0, 'z0');
        assert.isNumber(x1, 'x1');
        assert.isNumber(y1, 'y1');
        assert.isNumber(z1, 'z1');

        assert.notNaN(x0, 'x0');
        assert.notNaN(y0, 'y0');
        assert.notNaN(z0, 'z0');
        assert.notNaN(x1, 'x1');
        assert.notNaN(y1, 'y1');
        assert.notNaN(z1, 'z1');

        assert.isNonNegativeInteger(payload, 'payload');

        const address = this.getLeafAddress(index);

        aabb3_array_set(
            this.__data_float32,
            address,
            x0, y0, z0,
            x1, y1, z1
        );

        this.__data_uint32[address + 6] = payload;
    }

    /**
     * Read bounds of a box at the given address
     * @param {number} address where the box is found
     * @param {number[]|Float32Array} destination where to write the box coordinates (x0,y0,z0,x1,y1,z1)
     * @param {number} destination_offset offset within the destination array where to start writing results
     */
    readBounds(address, destination, destination_offset) {
        array_copy(this.__data_float32, address, destination, destination_offset, 6);
    }

    /**
     *
     * @param {number} leaf_index
     * @returns {number}
     */
    readLeafPayload(leaf_index) {
        const block_address = this.getLeafBlockAddress();

        const address = block_address + leaf_index * BVH_LEAF_NODE_SIZE + BVH_BOX_BYTE_SIZE;

        return this.__data_uint32[address];
    }

    compute_total_surface_area() {
        let result = 0;

        const box = new Float32Array(6);

        for (let i = 0; i < this.getTotalBoxCount(); i++) {
            let address;

            if (i < this.__node_count_binary) {
                address = i * BVH_BINARY_NODE_SIZE;
            } else {
                const li = i - this.__node_count_binary;
                address = li * BVH_LEAF_NODE_SIZE + this.getLeafBlockAddress();
            }

            this.readBounds(address, box, 0);
            result += aabb3_compute_half_surface_area(box[0], box[1], box[2], box[3], box[4], box[5]);
        }

        return result;
    }

    /**
     * Sort leaf nodes according to their morton codes
     * @param {number[]} bounds
     */
    sort_morton(bounds) {

        const leaf_block_address = this.__node_count_binary * BVH_BINARY_NODE_SIZE;

        if (this.__node_count_leaf < 2) {
            // no swaps available
            return;
        }

        const stack_top = stack.pointer;

        let stackPointer = stack_top;
        let i, j;

        stack[stackPointer++] = 0; // first node
        stack[stackPointer++] = this.__node_count_leaf - 1; // last node

        const data = this.__data_float32;

        while (stackPointer > stack_top) {
            stackPointer -= 2;

            const right = stack[stackPointer + 1];
            const left = stack[stackPointer];

            i = left;
            j = right;

            const pivot_index = (left + right) >> 1;

            const pivot_address = pivot_index * BVH_LEAF_NODE_SIZE + leaf_block_address;

            const pivot = build_morton(data, pivot_address, bounds);

            /* partition */
            while (i <= j) {

                while (build_morton(data, i * BVH_LEAF_NODE_SIZE + leaf_block_address, bounds) < pivot) {
                    i++;
                }

                while (build_morton(data, j * BVH_LEAF_NODE_SIZE + leaf_block_address, bounds) > pivot) {
                    j--;
                }

                if (i <= j) {

                    if (i !== j) {
                        //do swap
                        this.__swap_leaves(i, j);
                    }

                    i++;
                    j--;
                }
            }

            /* recursion */
            if (left < j) {
                stack[stackPointer++] = left;
                stack[stackPointer++] = j;
            }
            if (i < right) {
                stack[stackPointer++] = i;
                stack[stackPointer++] = right;
            }
        }
    }

    /**
     * Does not update intermediate node bounds
     * @param {number} i
     * @param {number} j
     * @private
     */
    __swap_leaves(i, j) {
        const leaf_block_address = this.getLeafBlockAddress();
        const a = i * BVH_LEAF_NODE_SIZE + leaf_block_address;
        const b = j * BVH_LEAF_NODE_SIZE + leaf_block_address;

        array_swap(
            this.__data_float32, a,
            this.__data_float32, b,
            BVH_LEAF_NODE_SIZE
        );
    }

    /**
     * Assemble leaf nodes into hierarchy, set binary node bounds iteratively bottom up
     */
    build() {
        const binary_node_count = this.__node_count_binary;

        const leaf_node_block_address = binary_node_count * BVH_BINARY_NODE_SIZE;

        let level = Math.floor(Math.log(binary_node_count) / Math.log(2));

        let i, offset, level_node_count;
        //NOTE: building the first level separately allows us to avoid some switching logic needed to determine what is the type of lower level node
        //build one level above leaf nodes
        level_node_count = Math.pow(2, level);
        offset = (level_node_count - 1) * BVH_BINARY_NODE_SIZE;

        let parentIndex;

        const node_count_leaf = this.__node_count_leaf;

        const float32 = this.__data_float32;

        // build bottom-most level, just above the leaves
        for (i = 0; i < level_node_count; i++) {
            const leaf_index_0 = i << 1;
            const leaf_index_1 = leaf_index_0 + 1;

            const leaf_offset_0 = leaf_node_block_address + leaf_index_0 * BVH_LEAF_NODE_SIZE;
            const leaf_offset_1 = leaf_node_block_address + leaf_index_1 * BVH_LEAF_NODE_SIZE;

            if (leaf_index_1 < node_count_leaf) {
                // both children nodes are valid
                aabb3_array_combine(
                    float32, offset,
                    float32, leaf_offset_0,
                    float32, leaf_offset_1
                );
            } else if (leaf_index_0 < node_count_leaf) {
                // only left child node is valid
                array_copy(float32, leaf_offset_0, float32, offset, 6);
            } else {
                //initialize to 0-size box same position as previous node
                copy_box_zero_size(this.__data_float32, offset, (offset - BVH_BINARY_NODE_SIZE));
            }

            offset += BVH_BINARY_NODE_SIZE;
        }

        level--;

        //build intermediate nodes
        for (; level >= 0; level--) {
            level_node_count = 1 << level;
            parentIndex = (level_node_count - 1);

            for (i = 0; i < level_node_count; i++) {

                const childIndex0 = (parentIndex << 1) + 1;

                const address_parent = parentIndex * BVH_BINARY_NODE_SIZE;

                const address_child_0 = childIndex0 * BVH_BINARY_NODE_SIZE;
                const address_child_1 = address_child_0 + BVH_BINARY_NODE_SIZE;

                aabb3_array_combine(
                    float32, address_parent,
                    float32, address_child_0,
                    float32, address_child_1
                );

                parentIndex++;
            }
        }
    }
}