three#BufferAttribute TypeScript Examples

private writeAttribute(gm: BufferGeometry, name: string, arr: Array<number>, itemSize: number): void {
        let attr: BufferAttribute = <BufferAttribute>gm.attributes[name];
        if (!attr || !attr.isBufferAttribute || attr.array.length < arr.length) {
            attr = new BufferAttribute(new Float32Array(arr.length), itemSize);
            gm.setAttribute(name, attr);
        }

        attr.copyArray(arr);
        attr.needsUpdate = true;
    }

private writeIndexAttribute(gm: BufferGeometry, arr: Array<number>): void {
        let attr: BufferAttribute = <BufferAttribute>gm.index;
        if (!attr || !attr.isBufferAttribute || attr.array.length < arr.length) {
            attr = new BufferAttribute(new Uint16Array(arr.length), 1);
            gm.index = attr;
        }

        attr.copyArray(arr);
        attr.needsUpdate = true;
    }

useTrimeshCollision = (geometry: BufferGeometry) => {
  const indices = (geometry.index as BufferAttribute).array as number[];

  const isInterleaved =
    // @ts-ignore
    geometry.attributes.position.isInterleavedBufferAttribute;

  let vertices: number[] = [];
  if (isInterleaved) {
    const attr = geometry.attributes.position as InterleavedBufferAttribute;
    const data = attr.data;
    for (let i = attr.offset; i < data.array.length; i += data.stride) {
      for (let x = 0; x < attr.itemSize; x++) {
        vertices.push(data.array[i + x]);
      }
    }
  } else {
    vertices = (geometry.attributes.position as BufferAttribute)
      .array as number[];
  }

  const [hitbox] = useTrimesh(() => ({
    type: "Static",
    args: [vertices, indices],
  }));

  return hitbox;
}

Geometry = defineComponent({
  emits: ['created'],
  props: {
    rotateX: Number,
    rotateY: Number,
    rotateZ: Number,
    attributes: { type: Array as PropType<Array<GeometryAttributeInterface>>, default: () => ([]) },
  },
  // inject for sub components
  inject: {
    mesh: MeshInjectionKey as symbol,
  },
  setup(): GeometrySetupInterface {
    return {}
  },
  created() {
    if (!this.mesh) {
      console.error('Missing parent Mesh')
      return
    }

    this.createGeometry()
    this.rotateGeometry()
    if (this.geometry) this.mesh.setGeometry(this.geometry)

    Object.keys(this.$props).forEach(prop => {
      // @ts-ignore
      watch(() => this[prop], this.refreshGeometry)
    })
  },
  unmounted() {
    this.geometry?.dispose()
  },
  methods: {
    createGeometry() {
      const bufferAttributes: Record<string, unknown> = {}
      const geometry = new BufferGeometry()
      this.attributes.forEach(attribute => {
        if (attribute.name && attribute.itemSize && attribute.array) {
          const bufferAttribute = bufferAttributes[attribute.name] = new BufferAttribute(attribute.array, attribute.itemSize, attribute.normalized)
          geometry.setAttribute(attribute.name, bufferAttribute)
        }
      })
      geometry.computeBoundingBox()
      geometry.userData.component = this
      this.geometry = geometry
      this.$emit('created', geometry)
    },
    rotateGeometry() {
      if (!this.geometry) return
      if (this.rotateX) this.geometry.rotateX(this.rotateX)
      if (this.rotateY) this.geometry.rotateY(this.rotateY)
      if (this.rotateZ) this.geometry.rotateZ(this.rotateZ)
    },
    refreshGeometry() {
      const oldGeo = this.geometry
      this.createGeometry()
      this.rotateGeometry()
      if (this.geometry && this.mesh) this.mesh.setGeometry(this.geometry)
      oldGeo?.dispose()
    },
  },
  render() { return [] },
})

export function Physics({
  drawDebug,
  drawDebugMode = DEFAULT_DEBUG_MODE,
  gravity,
  epsilon,
  fixedTimeStep,
  maxSubSteps,
  solverIterations,
  simulationSpeed = 1,
  children,
}: PropsWithChildren<AmmoPhysicsProps>) {
  const [physicsState, setPhysicsState] = useState<PhysicsState>();

  const sharedBuffersRef = useRef<SharedBuffers>({} as any);

  // Functions that are executed while the main thread holds control over the shared data
  const threadSafeQueueRef = useRef<(() => void)[]>([]);

  const physicsPerformanceInfoRef = useRef<PhysicsPerformanceInfo>({
    substepCounter: 0,
    lastTickMs: 0,
  });

  useEffect(() => {
    const uuids: string[] = [];
    const object3Ds: Record<string, Object3D> = {};
    const uuidToIndex: Record<string, number> = {};
    const IndexToUuid: Record<number, string> = {};
    const bodyOptions: Record<string, BodyConfig> = {};

    const softBodies: Record<UUID, Mesh> = {};

    const ammoWorker: Worker = createAmmoWorker();

    const workerHelpers = WorkerHelpers(ammoWorker);

    const rigidBodyBuffer = allocateCompatibleBuffer(
      4 * BUFFER_CONFIG.HEADER_LENGTH + //header
        4 * BUFFER_CONFIG.BODY_DATA_SIZE * BUFFER_CONFIG.MAX_BODIES + //matrices
        4 * BUFFER_CONFIG.MAX_BODIES //velocities
    );
    const headerIntArray = new Int32Array(
      rigidBodyBuffer,
      0,
      BUFFER_CONFIG.HEADER_LENGTH
    );
    const headerFloatArray = new Float32Array(
      rigidBodyBuffer,
      0,
      BUFFER_CONFIG.HEADER_LENGTH
    );
    const objectMatricesIntArray = new Int32Array(
      rigidBodyBuffer,
      BUFFER_CONFIG.HEADER_LENGTH * 4,
      BUFFER_CONFIG.BODY_DATA_SIZE * BUFFER_CONFIG.MAX_BODIES
    );
    const objectMatricesFloatArray = new Float32Array(
      rigidBodyBuffer,
      BUFFER_CONFIG.HEADER_LENGTH * 4,
      BUFFER_CONFIG.BODY_DATA_SIZE * BUFFER_CONFIG.MAX_BODIES
    );

    objectMatricesIntArray[0] = BufferState.UNINITIALIZED;

    const debugBuffer = allocateCompatibleBuffer(4 + 2 * DefaultBufferSize * 4);
    const debugIndex = new Uint32Array(debugBuffer, 0, 4);
    const debugVertices = new Float32Array(debugBuffer, 4, DefaultBufferSize);
    const debugColors = new Float32Array(
      debugBuffer,
      4 + DefaultBufferSize,
      DefaultBufferSize
    );
    const debugGeometry = new BufferGeometry();
    debugGeometry.setAttribute(
      "position",
      new BufferAttribute(debugVertices, 3).setUsage(DynamicDrawUsage)
    );
    debugGeometry.setAttribute(
      "color",
      new BufferAttribute(debugColors, 3).setUsage(DynamicDrawUsage)
    );

    sharedBuffersRef.current = {
      rigidBodies: {
        headerIntArray,
        headerFloatArray,
        objectMatricesFloatArray,
        objectMatricesIntArray,
      },

      softBodies: [],

      debug: {
        indexIntArray: debugIndex,
        vertexFloatArray: debugVertices,
        colorFloatArray: debugColors,
      },
    };

    const worldConfig: WorldConfig = {
      debugDrawMode: ammoDebugOptionsToNumber(drawDebugMode),
      gravity: gravity && new Vector3(gravity[0], gravity[1], gravity[2]),
      epsilon,
      fixedTimeStep,
      maxSubSteps,
      solverIterations,
    };

    workerHelpers.initWorld(worldConfig, sharedBuffersRef.current);

    const workerInitPromise = new Promise<PhysicsState>((resolve) => {
      ammoWorker.onmessage = async (event) => {
        const type: ClientMessageType = event.data.type;

        switch (type) {
          case ClientMessageType.READY: {
            if (event.data.sharedBuffers) {
              sharedBuffersRef.current = event.data.sharedBuffers;
            }

            resolve({
              workerHelpers,
              sharedBuffersRef,
              debugGeometry,
              debugBuffer,
              bodyOptions,
              uuids,
              object3Ds,
              softBodies,
              uuidToIndex,
              debugIndex,
              addRigidBody,
              removeRigidBody,
              addSoftBody,
              removeSoftBody,
              rayTest,
            });
            return;
          }
          case ClientMessageType.RIGIDBODY_READY: {
            const uuid = event.data.uuid;
            uuids.push(uuid);
            uuidToIndex[uuid] = event.data.index;
            IndexToUuid[event.data.index] = uuid;
            return;
          }
          case ClientMessageType.SOFTBODY_READY: {
            threadSafeQueueRef.current.push(() => {
              sharedBuffersRef.current.softBodies.push(
                event.data.sharedSoftBodyBuffers
              );
            });
            return;
          }
          case ClientMessageType.TRANSFER_BUFFERS: {
            sharedBuffersRef.current = event.data.sharedBuffers;
            return;
          }
          case ClientMessageType.RAYCAST_RESPONSE: {
            workerHelpers.resolveAsyncRequest(event.data);
            return;
          }
        }
        throw new Error("unknown message type" + type);
      };
    });

    workerInitPromise.then(setPhysicsState);

    function addRigidBody(
      uuid,
      mesh,
      shape: ShapeDescriptor,
      options: BodyConfig = {}
    ) {
      bodyOptions[uuid] = options;
      object3Ds[uuid] = mesh;

      if (!mesh.userData.useAmmo) {
        mesh.userData.useAmmo = {};
      }

      mesh.userData.useAmmo.rigidBody = {
        uuid,
      };

      workerHelpers.addRigidBody(uuid, mesh, shape, options);
    }

    function removeRigidBody(uuid: string) {
      uuids.splice(uuids.indexOf(uuid), 1);
      delete IndexToUuid[uuidToIndex[uuid]];
      delete uuidToIndex[uuid];
      delete bodyOptions[uuid];
      delete object3Ds[uuid].userData.useAmmo.rigidBody;
      delete object3Ds[uuid];
      workerHelpers.removeRigidBody(uuid);
    }

    function addSoftBody(uuid: UUID, mesh: Mesh, options: SoftBodyConfig = {}) {
      if (!mesh.geometry) {
        console.error("useSoftBody received: ", mesh);
        throw new Error("useSoftBody is only supported on BufferGeometries");
      }

      let indexLength: number;
      let vertexLength: number;
      let normalLength: number;

      switch (options.type) {
        case SoftBodyType.TRIMESH:
          // console.log("before merge ", mesh.geometry.attributes.position.count);
          mesh.geometry.deleteAttribute("normal");
          mesh.geometry.deleteAttribute("uv");
          mesh.geometry = mergeVertices(mesh.geometry);
          mesh.geometry.computeVertexNormals();
          // console.log("after merge ", mesh.geometry.attributes.position.count);

          indexLength =
            mesh.geometry.index!.count * mesh.geometry.index!.itemSize;
          vertexLength =
            mesh.geometry.attributes.position.count *
            mesh.geometry.attributes.position.itemSize;
          normalLength =
            mesh.geometry.attributes.normal.count *
            mesh.geometry.attributes.normal.itemSize;

          break;
        case SoftBodyType.ROPE:
          indexLength = 0;
          vertexLength =
            mesh.geometry.attributes.instanceStart.count *
            mesh.geometry.attributes.instanceStart.itemSize;
          normalLength = 0;

          break;
        default:
          throw new Error("unknown soft body type " + options.type);
      }

      const buffer = allocateCompatibleBuffer(
        indexLength * 4 + vertexLength * 4 + normalLength * 4
      );

      const sharedSoftBodyBuffers: SharedSoftBodyBuffers = {
        uuid,
        indexIntArray: new (indexLength > 65535 ? Uint32Array : Uint16Array)(
          buffer,
          0,
          indexLength
        ),
        vertexFloatArray: new Float32Array(
          buffer,
          indexLength * 4,
          vertexLength
        ),
        normalFloatArray: new Float32Array(
          buffer,
          indexLength * 4 + vertexLength * 4,
          normalLength
        ),
      };

      // Bullet softbodies operate in world-space,
      // so the transform needs to be baked into the vertex data
      mesh.updateMatrixWorld(true);
      mesh.geometry.applyMatrix4(mesh.matrixWorld);

      mesh.position.set(0, 0, 0);
      mesh.quaternion.set(0, 0, 0, 1);
      mesh.scale.set(1, 1, 1);

      mesh.frustumCulled = false;

      if (options.type === SoftBodyType.TRIMESH) {
        sharedSoftBodyBuffers.vertexFloatArray.set(
          mesh.geometry.attributes.position.array
        );

        sharedSoftBodyBuffers.indexIntArray.set(mesh.geometry.index!.array);
        sharedSoftBodyBuffers.normalFloatArray.set(
          mesh.geometry.attributes.normal.array
        );
      } else {
        for (let i = 0; i < vertexLength; i++) {
          sharedSoftBodyBuffers.vertexFloatArray[
            i * 3
          ] = mesh.geometry.attributes.instanceStart.getX(i);
          sharedSoftBodyBuffers.vertexFloatArray[
            i * 3 + 1
          ] = mesh.geometry.attributes.instanceStart.getY(i);
          sharedSoftBodyBuffers.vertexFloatArray[
            i * 3 + 2
          ] = mesh.geometry.attributes.instanceStart.getZ(i);
        }
      }

      if (isSharedArrayBufferSupported) {
        if (options.type === SoftBodyType.TRIMESH) {
          mesh.geometry.setAttribute(
            "position",
            new BufferAttribute(
              sharedSoftBodyBuffers.vertexFloatArray,
              3
            ).setUsage(DynamicDrawUsage)
          );

          mesh.geometry.setAttribute(
            "normal",
            new BufferAttribute(
              sharedSoftBodyBuffers.normalFloatArray,
              3
            ).setUsage(DynamicDrawUsage)
          );
        }
      }

      softBodies[uuid] = mesh;

      workerHelpers.addSoftBody(uuid, sharedSoftBodyBuffers, options);
    }

    function removeSoftBody(uuid: string) {
      delete softBodies[uuid];
      workerHelpers.removeSoftBody(uuid);

      sharedBuffersRef.current.softBodies = sharedBuffersRef.current.softBodies.filter(
        (ssbb) => ssbb.uuid !== uuid
      );
    }

    async function rayTest(options: RaycastOptions): Promise<RaycastHit[]> {
      const { hits } = await workerHelpers.makeAsyncRequest({
        type: MessageType.RAYCAST_REQUEST,
        ...options,
      });

      return hits.map(
        (hit: RaycastHitMessage): RaycastHit => {
          return {
            object: object3Ds[hit.uuid] || softBodies[hit.uuid],

            hitPosition: new Vector3(
              hit.hitPosition.x,
              hit.hitPosition.y,
              hit.hitPosition.z
            ),

            normal: new Vector3(hit.normal.x, hit.normal.y, hit.normal.z),
          };
        }
      );
    }

    return () => {
      ammoWorker.terminate();
      setPhysicsState(undefined);
    };
  }, []);

  useEffect(() => {
    if (!isSharedArrayBufferSupported) {
      if (drawDebug) {
        console.warn("debug visuals require SharedArrayBuffer support");
      }
      return;
    }

    if (physicsState) {
      if (drawDebug) {
        workerHelpers.enableDebug(true, physicsState.debugBuffer);
      } else {
        workerHelpers.enableDebug(false, physicsState.debugBuffer);
      }
    }
  }, [drawDebug, physicsState]);

  useEffect(() => {
    if (physicsState?.workerHelpers) {
      workerHelpers.setSimulationSpeed(simulationSpeed);
    }
  }, [physicsState?.workerHelpers, simulationSpeed]);

  if (!physicsState) {
    return null;
  }

  const { workerHelpers, debugGeometry } = physicsState;

  return (
    <AmmoPhysicsContext.Provider
      value={{
        ...workerHelpers,

        // workerHelpers Overrides
        addRigidBody: physicsState.addRigidBody,
        removeRigidBody: physicsState.removeRigidBody,

        addSoftBody: physicsState.addSoftBody,
        removeSoftBody: physicsState.removeSoftBody,

        object3Ds: physicsState.object3Ds,

        rayTest: physicsState.rayTest,

        physicsPerformanceInfoRef,
      }}
    >
      <PhysicsUpdate
        {...{
          physicsState,
          sharedBuffersRef,
          threadSafeQueueRef,
          physicsPerformanceInfoRef,
        }}
      />
      {drawDebug && <PhysicsDebug geometry={debugGeometry} />}
      {children}
    </AmmoPhysicsContext.Provider>
  );
}

export function PhysicsUpdate({
  physicsState,
  sharedBuffersRef,
  threadSafeQueueRef,
  physicsPerformanceInfoRef,
}: PhysicsUpdateProps) {
  useFrame(() => {
    if (!physicsState) {
      return;
    }

    const {
      workerHelpers,
      debugGeometry,
      bodyOptions,
      uuids,
      object3Ds,
      uuidToIndex,
      debugIndex,
      softBodies,
    } = physicsState;

    const sharedBuffers = sharedBuffersRef.current;

    if (
      // Check if the worker is finished with the buffer
      (!isSharedArrayBufferSupported &&
        sharedBuffers.rigidBodies.objectMatricesFloatArray.byteLength !== 0) ||
      (isSharedArrayBufferSupported &&
        Atomics.load(sharedBuffers.rigidBodies.headerIntArray, 0) ===
          BufferState.READY)
    ) {
      const lastSubstep = physicsPerformanceInfoRef.current.substepCounter;

      physicsPerformanceInfoRef.current.lastTickMs =
        sharedBuffers.rigidBodies.headerFloatArray[1];
      physicsPerformanceInfoRef.current.substepCounter =
        sharedBuffers.rigidBodies.headerIntArray[2];

      while (threadSafeQueueRef.current.length) {
        const fn = threadSafeQueueRef.current.shift();
        fn!();
      }

      // Skip copy if the physics worker didnt update
      if (lastSubstep !== physicsPerformanceInfoRef.current.substepCounter) {
        for (let i = 0; i < uuids.length; i++) {
          const uuid = uuids[i];
          const type = bodyOptions[uuid].type
            ? bodyOptions[uuid].type
            : BodyType.DYNAMIC;
          const object3D = object3Ds[uuid];
          if (type === BodyType.DYNAMIC) {
            matrix.fromArray(
              sharedBuffers.rigidBodies.objectMatricesFloatArray,
              uuidToIndex[uuid] * BUFFER_CONFIG.BODY_DATA_SIZE
            );

            inverse.copy(object3D.parent!.matrixWorld).invert();
            transform.multiplyMatrices(inverse, matrix);
            transform.decompose(object3D.position, object3D.quaternion, scale);
          } else {
            // sharedBuffers.rigidBodies.objectMatricesFloatArray.set(
            //   object3D.matrixWorld.elements,
            //   uuidToIndex[uuid] * BUFFER_CONFIG.BODY_DATA_SIZE
            // );
          }

          // print velocities
          // console.log(
          //   uuid,
          //   objectMatricesFloatArray[indexes[uuid] * BUFFER_CONFIG.BODY_DATA_SIZE + 16],
          //   objectMatricesFloatArray[indexes[uuid] * BUFFER_CONFIG.BODY_DATA_SIZE + 17]
          // );

          // print coliisions
          // const collisions = [];
          // for (let j = 18; j < 26; j++) {
          //   const collidingIndex = objectMatricesIntArray[uuidToIndex[uuid] * BUFFER_CONFIG.BODY_DATA_SIZE + j];
          //   if (collidingIndex !== -1) {
          //     collisions.push(IndexToUuid[collidingIndex]);
          //   }
          // }
          // console.log(uuid, collisions);
        }

        for (const softBodyBuffers of sharedBuffersRef.current.softBodies) {
          const softBodyMesh = softBodies[softBodyBuffers.uuid];

          if (softBodyMesh) {
            if ((softBodyMesh.geometry as LineGeometry).isLineGeometry) {
              (softBodyMesh.geometry as LineGeometry).setPositions(
                softBodyBuffers.vertexFloatArray
              );
              softBodyMesh.geometry.attributes.instanceStart.needsUpdate = true;
              softBodyMesh.geometry.attributes.instanceEnd.needsUpdate = true;
            } else {
              if (!isSharedArrayBufferSupported) {
                (softBodyMesh.geometry.attributes
                  .position as BufferAttribute).copyArray(
                  softBodyBuffers.vertexFloatArray
                );
                (softBodyMesh.geometry.attributes
                  .normal as BufferAttribute).copyArray(
                  softBodyBuffers.normalFloatArray
                );
              }

              softBodyMesh.geometry.attributes.position.needsUpdate = true;
              if (softBodyMesh.geometry.attributes.normal) {
                softBodyMesh.geometry.attributes.normal.needsUpdate = true;
              }
            }
          }
        }
      }

      if (isSharedArrayBufferSupported) {
        Atomics.store(
          sharedBuffers.rigidBodies.headerIntArray,
          0,
          BufferState.CONSUMED
        );
      } else {
        workerHelpers.transferSharedBuffers(sharedBuffersRef.current);
      }
    }

    if (isSharedArrayBufferSupported) {
      /* DEBUG RENDERING */
      const index = Atomics.load(debugIndex, 0);
      if (!!index) {
        debugGeometry.attributes.position.needsUpdate = true;
        debugGeometry.attributes.color.needsUpdate = true;
        debugGeometry.setDrawRange(0, index);
      }
      Atomics.store(debugIndex, 0, 0);
    }
  });

  return null;
}