@woosh/meep-engine/src/engine/graphics/geometry/instancing/InstancedMeshGroup.js

Pure JavaScript game engine. Fully featured and production ready.

import {
    BufferGeometry,
    DynamicDrawUsage,
    InstancedBufferAttribute,
    InstancedBufferGeometry,
    Line,
    LineSegments,
    Mesh,
    MeshDepthMaterial,
    RGBADepthPacking
} from 'three';
import { Cache } from "../../../../core/cache/Cache.js";
import { array_copy } from "../../../../core/collection/array/array_copy.js";
import { typed_array_copy } from "../../../../core/collection/array/typed/typed_array_copy.js";
import { max3 } from "../../../../core/math/max3.js";
import { min2 } from "../../../../core/math/min2.js";
import { computeMaterialEquality } from "../../../asset/loaders/material/computeMaterialEquality.js";
import { computeMaterialHash } from "../../../asset/loaders/material/computeMaterialHash.js";
import { StaticMaterialCache } from "../../../asset/loaders/material/StaticMaterialCache.js";
import { DrawMode } from "../../ecs/mesh-v2/DrawMode.js";
import { composeCompile } from "../../material/composeCompile.js";
import ThreeFactory from "../../three/ThreeFactory.js";
import { geometry_copy } from "./geometry_copy.js";
import { rewriteMaterial } from "./rewriteMaterial.js";

/**
 * @typedef {Object} CacheKey
 * @property {Material} color
 * @property {Material} depth
 */

/**
 * @readonly
 * @type {Cache<Material,CacheKey>}
 */
const material_cache = new Cache({
    maxWeight: 1024,
    keyHashFunction: computeMaterialHash,
    keyEqualityFunction: computeMaterialEquality
});

export class InstancedMeshGroup {
    /**
     *
     * @constructor
     */
    constructor() {
        /**
         * Instanced geometry
         * @type {InstancedBufferGeometry|null}
         * @private
         */
        this.__threeGeometry = null;

        /**
         * Geometry of a single instance
         * @type {BufferGeometry|null}
         * @private
         */
        this.__threeInstanceGeometry = null;

        /**
         * @private
         * @type {number}
         */
        this.count = 0;

        /**
         *
         * @type {number}
         */
        this.capacity = 0;

        /**
         *
         * @type {number}
         */
        this.growFactor = 1.2;
        /**
         * Minimum spare capacity increase during growing
         * @type {number}
         */
        this.growConstant = 16;

        /**
         *
         * @type {number}
         */
        this.shrinkFactor = 0.5;

        /**
         * Minimum capacity reduction for shrinkage to occur
         * @type {number}
         */
        this.shrinkConstant = 64;

        /**
         *
         * @type {InstancedBufferAttribute}
         * @private
         */
        this.__attributeTransform = null;

        /**
         * Shortcut to buffer attribute data
         * @type {Float32Array}
         * @private
         */
        this.__attributeTransformArray = null;


        /**
         *
         * TODO implement usage
         * @type {InstancedBufferAttribute}
         * @private
         */
        this.__attributeColor = null;

        /**
         * TODO implement usage
         * @type {Uint8Array}
         * @private
         */
        this.__attributeColorArray = null;

        /**
         *
         * @type {Usage}
         * @private
         */
        this.__instanceUsage = DynamicDrawUsage;

        /**
         *
         * @type {Material}
         * @private
         */
        this.__material = null;

        this.indices = [];
        this.references = [];

        /**
         *
         * @type {DrawMode}
         * @private
         */
        this.__draw_mode = DrawMode.Triangles;

        this.mesh = ThreeFactory.createMesh();

        this.use_color = true;
    }

    __build_mesh() {

        let new_object;
        switch (this.__draw_mode) {
            case DrawMode.Triangles:
                new_object = new Mesh(this.mesh.geometry, this.mesh.material);
                break;
            case DrawMode.Lines:
                new_object = new Line(this.mesh.geometry, this.mesh.material);
                break;
            case DrawMode.LineSegments:
                new_object = new LineSegments(this.mesh.geometry, this.mesh.material);
                break;

            default:
                throw new Error(`Unsupported DrawMode '${this.__draw_mode}'`);
        }

        new_object.matrixAutoUpdate = false;
        new_object.frustumCulled = false;
        new_object.matrixWorldNeedsUpdate = false;

        if (this.mesh !== null) {
            new_object.castShadow = this.mesh.castShadow;
            new_object.receiveShadow = this.mesh.receiveShadow;

            new_object.customDepthMaterial = this.mesh.customDepthMaterial;
        }

        this.mesh = new_object;
    }

    /**
     *
     * @param {DrawMode} v
     */
    set draw_mode(v) {
        if (v === this.__draw_mode) {
            // no change
            return;
        }

        this.__draw_mode = v;

        this.__build_mesh();

    }

    /**
     *
     * @param {Usage} v
     */
    set instance_usage(v) {
        if (this.__instanceUsage === v) {
            return;
        }

        this.__instanceUsage = v;

        if (this.__attributeTransform !== null) {
            this.__attributeTransform.setUsage(v);
        }

        if (this.__attributeColor !== null) {
            this.__attributeColor.setUsage(v);
        }
    }

    /**
     *
     * @return {Usage}
     */
    get instance_usage() {
        return this.__instanceUsage;
    }

    dispose() {
        if (this.__threeGeometry !== null) {
            this.__threeGeometry.dispose();
        }
    }

    /**
     *
     * @param {THREE.BufferGeometry} geometry
     */
    setGeometry(geometry) {
        if (geometry.isBufferGeometry !== true) {
            throw new Error(`Expected THREE.BufferedGeometry, got something else instead.`);
        }

        this.__threeInstanceGeometry = geometry;

        this.build();
    }

    /**
     *
     * @param {THREE.Material|THREE.ShaderMaterial} sourceMaterial
     * @returns {CacheKey}
     */
    #buildMaterial(sourceMaterial) {
        //console.warn(`building material : {id:${sourceMaterial.id}, name: ${sourceMaterial.name}, type: ${sourceMaterial.type}`)

        const material = sourceMaterial.clone();

        if (material.isShaderMaterial) {
            // transfer source material uniforms
            for (const uniformsKey in material.uniforms) {
                material.uniforms[uniformsKey] = sourceMaterial.uniforms[uniformsKey];
            }
        }


        material.onBeforeCompile = composeCompile(sourceMaterial.onBeforeCompile, rewriteMaterial);

        //we need a custom depth material to ensure shadows will be drawn correctly
        const depthMaterial = new MeshDepthMaterial({
            depthPacking: RGBADepthPacking,
        });

        //if the source material uses alpha testing - enable it on the depth material
        if (material.alphaTest !== 0) {
            depthMaterial.alphaTest = material.alphaTest;
            depthMaterial.map = material.map;
        }

        depthMaterial.onBeforeCompile = rewriteMaterial;

        const cachedMaterial = StaticMaterialCache.Global.acquire(material);

        return {
            color: cachedMaterial,
            depth: StaticMaterialCache.Global.acquire(depthMaterial)
        };

    }

    /**
     *
     * @param {THREE.Material|THREE.ShaderMaterial} sourceMaterial
     */
    setMaterial(sourceMaterial) {
        const m = material_cache.getOrCompute(sourceMaterial, this.#buildMaterial, this);

        this.__material = m.color;

        this.mesh.material = m.color;
        this.mesh.customDepthMaterial = m.depth;
    }

    /**
     *
     * @param {number} size
     */
    setCapacity(size) {
        if (this.capacity === size) {
            // no change
            return;
        }

        this.capacity = size;

        this.build();
    }

    /**
     *
     * @param {number} size
     */
    ensureCapacity(size) {

        const currentCapacity = this.capacity;

        if (currentCapacity < size) {
            // grow

            const newCapacityRaw = max3(
                size,
                currentCapacity * this.growFactor,
                currentCapacity + this.growConstant
            );

            const newCapacityInteger = Math.ceil(newCapacityRaw);

            this.setCapacity(newCapacityInteger);

        } else if (
            min2(
                currentCapacity * this.shrinkFactor,
                currentCapacity - this.shrinkConstant
            ) > size
        ) {
            // shrink
            this.setCapacity(size);
        }

    }

    /**
     *
     * @param {number} index
     * @param {number[]|ArrayLike<number>|Float32Array} transform 4x4 transform matrix of the instance
     */
    setTransformAt(index, transform) {
        array_copy(transform, 0, this.__attributeTransformArray, index * 16, 16);
    }

    /**
     *
     * @param {number} index
     * @param {number[]|ArrayLike<number>|Float32Array} color RGBA color in uint8 format (0...255), LDR
     */
    setColorAt(index, color) {
        array_copy(color, 0, this.__attributeColorArray, index * 4, 4);
    }

    /**
     *
     * @param {number} index
     * @param {number} r
     * @param {number} g
     * @param {number} b
     * @param {number} a
     */
    setColorByComponentAt(index, r, g, b, a) {
        const color_array = this.__attributeColorArray;

        const i4 = index * 4;

        color_array[i4] = r;
        color_array[i4 + 1] = g;
        color_array[i4 + 2] = b;
        color_array[i4 + 3] = a;
    }

    requestAttributeUpdate() {
        this.__attributeTransform.needsUpdate = true;

        if (this.__attributeColor !== null) {
            this.__attributeColor.needsUpdate = true;
        }
    }

    /**
     * Swap position in attribute arrays of two elements
     * @param {int} indexA
     * @param {int} indexB
     */
    swap(indexA, indexB) {
        throw new Error('Not Implemented Yet');
    }

    /**
     *
     * @param {int} size
     */
    setCount(size) {
        this.ensureCapacity(size);

        this.count = size;
        this.__threeGeometry.instanceCount = size;
    }

    /**
     *
     * @return {number|int}
     */
    getCount() {
        return this.count;
    }

    /**
     *
     * @param {int} reference
     * @returns {int}
     */
    add(reference) {
        //get index
        const index = this.count;
        //grow
        this.setCount(index + 1);

        this.indices[reference] = index;
        this.references[index] = reference;

        return index;
    }

    /**
     *
     * @param {function(index:int,reference:int)} visitor
     * @param {*} [thisArg]
     */
    traverseReferences(visitor, thisArg) {
        const indices = this.indices;
        const index_count = indices.length;

        for (let i = 0; i < index_count; i++) {
            const index = indices[i];

            if (index === undefined) {
                // unoccupied
                continue;
            }

            visitor.call(thisArg, index, i);
        }
    }

    /**
     *
     * @param {int} reference
     */
    remove(reference) {
        //dereference
        const index = this.indices[reference];

        if (index === undefined) {
            //reference is not known
            throw new Error(`Reference '${reference}' was not found`);
        }

        delete this.indices[reference];

        const lastIndex = this.count - 1;

        const references = this.references;

        if (index === lastIndex) {
            //easy case, reference is placed at the end, no swap needed, we can just forget about it
            delete references[index];
        } else {
            //not at the end, move the very last reference to this place, effectively shrinking the set


            const a_transform = this.__attributeTransform;

            a_transform.array.copyWithin(index * 16, lastIndex * 16, lastIndex * 16 + 16);
            a_transform.needsUpdate = true;

            const a_color = this.__attributeColor;

            if (a_color !== null) {
                a_color.array.copyWithin(index * 4, lastIndex * 4, lastIndex * 4 + 4);
                a_color.needsUpdate = true;
            }


            //update moved reference index
            const movedReference = references[lastIndex];

            if (movedReference === undefined) {
                //moved reference not found
                throw new Error(`Moved reference was not found`);
            }

            delete references[lastIndex];

            //assume the place of removed reference
            references[index] = movedReference;
            this.indices[movedReference] = index;
        }

        //update size
        this.setCount(this.count - 1);
    }

    build() {
        if (this.__threeGeometry !== null) {
            // dispose old geometry's GPU memory
            this.__threeGeometry.dispose();
        }

        this.__threeGeometry = new InstancedBufferGeometry();

        const instance = this.__threeInstanceGeometry;
        const geometry = this.__threeGeometry;

        geometry.instanceCount = this.count;
        geometry.dynamic = true;

        //copy instance attributes
        geometry_copy(geometry, instance);

        //build instanced attributes
        const newTransformArray = new Float32Array(this.capacity * 16);

        if (this.__attributeTransform !== null) {
            const oldTransformArray = this.__attributeTransform.array;
            typed_array_copy(oldTransformArray, newTransformArray);

            if (newTransformArray.length > oldTransformArray.length) {
                //fill the rest of the array with 0s
                // TODO consider writing identity matrix instead
                newTransformArray.fill(0, oldTransformArray.length);
            }
        }

        //rewrite old attributes
        this.__attributeTransformArray = newTransformArray;
        this.__attributeTransform = new InstancedBufferAttribute(newTransformArray, 16);
        this.__attributeTransform.setUsage(this.__instanceUsage);

        //add attributes to newly created geometry
        geometry.setAttribute("instanceMatrix", this.__attributeTransform);

        if (this.use_color) {

            // color attribute
            const newColorArray = new Uint8Array(this.capacity * 4);

            if (this.__attributeColor !== null) {
                const oldColorArray = this.__attributeColor.array;
                typed_array_copy(oldColorArray, newColorArray);

                if (newColorArray.length > oldColorArray.length) {
                    newColorArray.fill(255, oldColorArray.length);
                }
            }

            this.__attributeColorArray = newColorArray;
            this.__attributeColor = new InstancedBufferAttribute(newColorArray, 4);
            this.__attributeColor.normalized = true;
            this.__attributeColor.setUsage(this.__instanceUsage);

            geometry.setAttribute("instanceColor", this.__attributeColor);
        } else {
            this.__attributeColorArray = null;
            this.__attributeColor = null;
        }

        this.mesh.geometry = geometry;
    }

    /**
     *
     * @param {THREE.BufferGeometry} geometry
     * @param {Material} material
     * @returns {InstancedMeshGroup}
     */
    static from(geometry, material) {
        const r = new InstancedMeshGroup();

        r.setGeometry(geometry);
        r.setMaterial(material);

        return r;
    }
}