yawn/addons/handles/src/materials/PBRMaterial.ts

import type { Scene } from "../Scene";
import type { Texture } from "./Texture";

export type PBRTextureInfo = {
  texCoord?: number;
  offset?: ArrayLike<number>;
  scale?: ArrayLike<number>;
  rotation?: number;
};

export type PBRMaterialOptions = {
  baseColor?: ArrayLike<number>;
  metallic?: number;
  roughness?: number;
  emissive?: ArrayLike<number>;
  normalScale?: number;
  alphaCutoff?: number;
  baseColorTexture?: Texture;
  metallicRoughnessTexture?: Texture;
  normalTexture?: Texture;
  emissiveTexture?: Texture;
  occlusionTexture?: Texture;
  transmissionTexture?: Texture;
  thicknessTexture?: Texture;
  clearcoatTexture?: Texture;
  iridescenceTexture?: Texture;
  iridescenceThicknessTexture?: Texture;
  anisotropyTexture?: Texture;
  specularTexture?: Texture;
  specularColorTexture?: Texture;
  textureInfos?: ArrayLike<PBRTextureInfo | undefined>;
  flags?: number;
  occlusionStrength?: number;
  transmission?: number;
  ior?: number;
  thickness?: number;
  attenuationDistance?: number;
  attenuationColor?: ArrayLike<number>;
  emissiveStrength?: number;
  clearcoat?: number;
  clearcoatRoughness?: number;
  specular?: number;
  sheenColor?: ArrayLike<number>;
  sheenRoughness?: number;
  specularColor?: ArrayLike<number>;
  iridescence?: number;
  iridescenceIor?: number;
  iridescenceThicknessMinimum?: number;
  iridescenceThicknessMaximum?: number;
  anisotropy?: number;
  anisotropyRotation?: number;
  specularExtensionColor?: ArrayLike<number>;
};

/** Conventional PBR values, extensions, texture IDs, and UV transforms. */
export class PBRMaterial {
  readonly scene: Scene;
  id = -1;
  readonly ready: Promise<this>;
  #disposed = false;

  constructor(scene: Scene, options: PBRMaterialOptions = {}) {
    this.scene = scene;
    const textures = [
      options.baseColorTexture,
      options.metallicRoughnessTexture,
      options.normalTexture,
      options.emissiveTexture,
      options.occlusionTexture,
      options.transmissionTexture,
      options.thicknessTexture,
      options.clearcoatTexture,
      options.iridescenceTexture,
      options.iridescenceThicknessTexture,
      options.anisotropyTexture,
      options.specularTexture,
      options.specularColorTexture,
    ].filter((texture): texture is Texture => texture !== undefined);
    this.ready = Promise.all(textures.map((texture) => texture.ready)).then(
      async () => {
        const id = await scene.allocateMaterial();
        this.id = id;
        const color = Array.from(options.baseColor ?? [1, 1, 1, 1]);
        const emissive = Array.from(options.emissive ?? [0, 0, 0]);
        const attenuation = Array.from(options.attenuationColor ?? [1, 1, 1]);
        const secondary = Array.from(
          options.flags && options.flags & 16
            ? (options.specularColor ?? [1, 1, 1])
            : (options.sheenColor ?? [0, 0, 0]),
        );
        const specularExtensionColor = Array.from(
          options.specularExtensionColor ?? [1, 1, 1],
        );
        if (
          color.length !== 4 ||
          emissive.length !== 3 ||
          attenuation.length !== 3 ||
          secondary.length !== 3 ||
          specularExtensionColor.length !== 3
        )
          throw new TypeError("PBR material vectors");
        scene
          .array("materials")
          .write(id, [
            ...color,
            options.metallic ?? 0,
            options.roughness ?? 0.7,
            ...emissive,
            options.normalScale ?? 1,
            options.alphaCutoff ?? 0.5,
            options.flags ?? 0,
          ]);
        scene
          .array("materialTextures")
          .write(id, [
            (options.baseColorTexture?.id ?? -1) + 1,
            (options.metallicRoughnessTexture?.id ?? -1) + 1,
            (options.normalTexture?.id ?? -1) + 1,
            (options.emissiveTexture?.id ?? -1) + 1,
            (options.occlusionTexture?.id ?? -1) + 1,
            (options.transmissionTexture?.id ?? -1) + 1,
            (options.thicknessTexture?.id ?? -1) + 1,
            (options.clearcoatTexture?.id ?? -1) + 1,
            (options.iridescenceTexture?.id ?? -1) + 1,
            (options.iridescenceThicknessTexture?.id ?? -1) + 1,
            (options.anisotropyTexture?.id ?? -1) + 1,
            (options.specularTexture?.id ?? -1) + 1,
            (options.specularColorTexture?.id ?? -1) + 1,
            0,
            0,
            0,
          ]);
        scene
          .array("materialExtensions")
          .write(id, [
            options.transmission ?? 0,
            options.ior ?? 1.5,
            options.thickness ?? 0,
            options.attenuationDistance ?? 0,
            ...attenuation,
            options.emissiveStrength ?? 1,
            options.clearcoat ?? 0,
            options.clearcoatRoughness ?? 0,
            options.specular ?? 1,
            options.occlusionStrength ?? 1,
            ...secondary,
            options.sheenRoughness ?? 0,
            options.iridescence ?? 0,
            options.iridescenceIor ?? 1.3,
            options.iridescenceThicknessMinimum ?? 100,
            options.iridescenceThicknessMaximum ?? 400,
            options.anisotropy ?? 0,
            options.anisotropyRotation ?? 0,
            0,
            0,
            ...specularExtensionColor,
            0,
          ]);
        const transforms = new Float32Array(104);
        for (let slot = 0; slot < 13; slot++) {
          const info = options.textureInfos?.[slot];
          const offset = Array.from(info?.offset ?? [0, 0]);
          const scale = Array.from(info?.scale ?? [1, 1]);
          if (offset.length !== 2 || scale.length !== 2)
            throw new TypeError("PBR texture transform");
          const rotation = info?.rotation ?? 0;
          transforms.set(
            [
              offset[0],
              offset[1],
              scale[0],
              scale[1],
              Math.sin(rotation),
              Math.cos(rotation),
              info?.texCoord ?? 0,
              0,
            ],
            slot * 8,
          );
        }
        scene.array("materialTextureInfo").write(id, transforms);
        return this;
      },
    );
  }

  #values() {
    if (this.id < 0)
      throw new Error("Await material.ready before reading or writing it");
    return this.scene.array("materials").row(this.id);
  }

  get baseColor() {
    return this.#values().subarray(0, 4);
  }
  set baseColor(value: ArrayLike<number>) {
    if (value.length !== 4) throw new RangeError("baseColor");
    this.#values().set(value, 0);
  }
  get metallic() {
    return this.#values()[4];
  }
  set metallic(value: number) {
    this.#values()[4] = value;
  }
  get roughness() {
    return this.#values()[5];
  }
  set roughness(value: number) {
    this.#values()[5] = value;
  }
  get emissive() {
    return this.#values().subarray(6, 9);
  }
  set emissive(value: ArrayLike<number>) {
    if (value.length !== 3) throw new RangeError("emissive");
    this.#values().set(value, 6);
  }

  async dispose() {
    await this.ready;
    if (this.#disposed) return;
    this.#disposed = true;
    await this.scene.core.deleteObject("materials", this.id);
  }
}