Extends THREE.js InstancedMesh with the ability to reference instances by a string key and add/remove instances with dynamic buffer resizing

IndexedInstancedMesh.ts

import * as THREE from "three";

type UserData = { [key: string]: unknown };

const INITIAL_SIZE = 4;

const tempMat4 = new THREE.Matrix4();
const tempColor = new THREE.Color();

/**
 * Extends InstancedMesh with the ability to reference instances by a string key
 * and add/remove instances by key with dynamic buffer resizing.
 */
export class IndexedInstancedMesh<
  TGeometry extends THREE.BufferGeometry = THREE.BufferGeometry,
  TMaterial extends THREE.Material | THREE.Material[] = THREE.Material | THREE.Material[],
  TUserData extends UserData = UserData,
> extends THREE.InstancedMesh<TGeometry, TMaterial> {
  // A map of { key -> index }
  private _map = new Map<string, number>();
  // A map of { key -> UserData }
  private _userData = new Map<string, TUserData>();
  // Effectively a map of { index -> key }
  private _keys: string[] = [];
  // Total size of the buffer attributes, which can be larger than .count (instances in use)
  private _capacity: number;

  constructor(geometry: TGeometry, material: TMaterial) {
    super(geometry, material, 0);

    this._capacity = INITIAL_SIZE;
    this.resize();
  }

  set(key: string, matrix: THREE.Matrix4, color: THREE.Color, userData: TUserData): void {
    let index: number;
    if (this._map.has(key)) {
      // Update the existing entry
      index = this._map.get(key)!;
    } else {
      // Create a new entry
      index = this.count;
      this._map.set(key, index);
      this._keys.push(key);
      this.setCount(this._map.size);
    }

    this._userData.set(key, userData);
    this.setMatrixAt(index, matrix);
    this.setColorAt(index, color);
    this.instanceMatrix.needsUpdate = true;
    this.instanceColor!.needsUpdate = true;
  }

  delete(key: string): boolean {
    if (!this._map.has(key)) return false;

    const deleteIndex = this._map.get(key)!;
    if (deleteIndex !== this.count - 1) {
      // Swap the last element with the one to be deleted
      const lastIndex = this.count - 1;
      const lastKey = this._keys[lastIndex]!;
      this._keys[deleteIndex] = lastKey;
      this._map.set(lastKey, deleteIndex);
      this.getMatrixAt(lastIndex, tempMat4);
      this.setMatrixAt(deleteIndex, tempMat4);
      this.getColorAt(lastIndex, tempColor);
      this.setColorAt(deleteIndex, tempColor);

      this.instanceMatrix.needsUpdate = true;
      this.instanceColor!.needsUpdate = true;
    }

    // Delete the last element
    this.setCount(this.count - 1);
    this._keys.pop();

    // Delete the requested key from the map
    this._map.delete(key);
    return true;
  }

  has(key: string): boolean {
    return this._map.has(key);
  }

  getUserData(key: string): TUserData | undefined {
    return this._userData.get(key);
  }

  private setCount(count: number) {
    while (count >= this._capacity) this.expand();
    this.count = count;
    this.instanceMatrix.count = count;
    this.instanceColor!.count = count;
  }

  private expand() {
    this._capacity = this._capacity + Math.trunc(this._capacity / 2) + 16;
    this.resize();
  }

  private resize() {
    const oldMatrixArray = this.instanceMatrix.array as Float32Array;
    const oldColorArray = this.instanceColor?.array as Float32Array | undefined;

    const newMatrixArray = new Float32Array(this._capacity * 16);
    const newColorArray = new Float32Array(this._capacity * 3);

    if (oldMatrixArray.length > 0) {
      newMatrixArray.set(oldMatrixArray);
    }
    if (oldColorArray && oldColorArray.length > 0) {
      newColorArray.set(oldColorArray);
    }

    this.instanceMatrix = new THREE.InstancedBufferAttribute(newMatrixArray, 16);
    this.instanceColor = new THREE.InstancedBufferAttribute(newColorArray, 3);

    this.instanceMatrix.setUsage(THREE.DynamicDrawUsage);
    this.instanceColor.setUsage(THREE.DynamicDrawUsage);
  }
}