refactor: move Label and LabelMaterial into separate files

src/Label.ts

@@ -0,0 +1,225 @@
+import * as THREE from "three";
+
+import { LabelMaterial } from "./LabelMaterial.ts";
+import { LabelPool } from "./LabelPool.ts";
+
+const tempVec2 = new THREE.Vector2();
+
+/**
+ * Since THREE.js r151, InstancedMesh supports bounding sphere calculations using instanceMatrix.
+ * However, Label does not use instanceMatrix and the resulting bounding spheres are have NaN
+ * values. Instead, fall back to using the (non-instanced) bounding sphere of the geometry, which at
+ * least provides a semi-correct value based on the label's `position`.
+ */
+class InstancedMeshWithBasicBoundingSphere extends THREE.InstancedMesh {
+  override computeBoundingSphere(): void {
+    this.geometry.computeBoundingSphere();
+    const boundingSphere = this.geometry.boundingSphere;
+    if (boundingSphere) {
+      (this.boundingSphere ??= new THREE.Sphere()).copy(boundingSphere);
+    }
+  }
+}
+
+export class Label extends THREE.Object3D {
+  text = "";
+  mesh: THREE.InstancedMesh;
+  geometry: THREE.InstancedBufferGeometry;
+  material: LabelMaterial;
+  pickingMaterial: LabelMaterial;
+
+  instanceAttrData: Float32Array;
+  instanceAttrBuffer: THREE.InstancedInterleavedBuffer;
+
+  instanceBoxPosition: THREE.InterleavedBufferAttribute;
+  instanceCharPosition: THREE.InterleavedBufferAttribute;
+  instanceUv: THREE.InterleavedBufferAttribute;
+  instanceBoxSize: THREE.InterleavedBufferAttribute;
+  instanceCharSize: THREE.InterleavedBufferAttribute;
+
+  lineHeight = 1;
+
+  public labelPool: LabelPool;
+
+  constructor(labelPool: LabelPool) {
+    super();
+    this.labelPool = labelPool;
+
+    this.geometry = new THREE.InstancedBufferGeometry();
+
+    this.geometry.setAttribute("position", LabelPool.QUAD_POSITIONS);
+    this.geometry.setAttribute("uv", LabelPool.QUAD_UVS);
+
+    this.instanceAttrData = new Float32Array();
+    this.instanceAttrBuffer = new THREE.InstancedInterleavedBuffer(this.instanceAttrData, 10, 1);
+    this.instanceBoxPosition = new THREE.InterleavedBufferAttribute(this.instanceAttrBuffer, 2, 0);
+    this.instanceCharPosition = new THREE.InterleavedBufferAttribute(this.instanceAttrBuffer, 2, 2);
+    this.instanceUv = new THREE.InterleavedBufferAttribute(this.instanceAttrBuffer, 2, 4);
+    this.instanceBoxSize = new THREE.InterleavedBufferAttribute(this.instanceAttrBuffer, 2, 6);
+    this.instanceCharSize = new THREE.InterleavedBufferAttribute(this.instanceAttrBuffer, 2, 8);
+    this.geometry.setAttribute("instanceBoxPosition", this.instanceBoxPosition);
+    this.geometry.setAttribute("instanceCharPosition", this.instanceCharPosition);
+    this.geometry.setAttribute("instanceUv", this.instanceUv);
+    this.geometry.setAttribute("instanceBoxSize", this.instanceBoxSize);
+    this.geometry.setAttribute("instanceCharSize", this.instanceCharSize);
+
+    this.material = new LabelMaterial({ atlasTexture: labelPool.atlasTexture });
+    this.pickingMaterial = new LabelMaterial({ picking: true });
+
+    this.mesh = new InstancedMeshWithBasicBoundingSphere(this.geometry, this.material, 0);
+    this.mesh.userData.pickingMaterial = this.pickingMaterial;
+
+    this.mesh.onBeforeRender = (renderer, _scene, _camera, _geometry, _material, _group) => {
+      renderer.getSize(tempVec2);
+      this.material.uniforms.uCanvasSize!.value[0] = tempVec2.x;
+      this.material.uniforms.uCanvasSize!.value[1] = tempVec2.y;
+      this.pickingMaterial.uniforms.uCanvasSize!.value[0] = tempVec2.x;
+      this.pickingMaterial.uniforms.uCanvasSize!.value[1] = tempVec2.y;
+    };
+
+    this.add(this.mesh);
+
+    labelPool.addEventListener("scaleFactorChange", () => {
+      // Trigger recalculation of scale uniform
+      this.setLineHeight(this.lineHeight);
+    });
+
+    labelPool.addEventListener("atlasChange", () => {
+      this._handleAtlasChange();
+    });
+    this._handleAtlasChange();
+  }
+
+  private _handleAtlasChange() {
+    this.material.uniforms.uTextureSize!.value[0] = this.labelPool.atlasTexture.image.width;
+    this.material.uniforms.uTextureSize!.value[1] = this.labelPool.atlasTexture.image.height;
+    this.setLineHeight(this.lineHeight);
+    this._needsUpdateLayout = true;
+    this._updateLayoutIfNeeded();
+  }
+
+  dispose(): void {
+    this.geometry.dispose();
+    this.material.dispose();
+    this.pickingMaterial.dispose();
+    this.mesh.dispose();
+  }
+
+  private reallocateAttributeBufferIfNeeded(numChars: number) {
+    const requiredLength = numChars * 10 * Float32Array.BYTES_PER_ELEMENT;
+    if (this.instanceAttrData.byteLength < requiredLength) {
+      this.instanceAttrData = new Float32Array(requiredLength);
+      this.instanceAttrBuffer = new THREE.InstancedInterleavedBuffer(this.instanceAttrData, 10, 1);
+      this.instanceBoxPosition.data = this.instanceAttrBuffer;
+      this.instanceCharPosition.data = this.instanceAttrBuffer;
+      this.instanceUv.data = this.instanceAttrBuffer;
+      this.instanceBoxSize.data = this.instanceAttrBuffer;
+      this.instanceCharSize.data = this.instanceAttrBuffer;
+    }
+  }
+
+  setText(text: string): void {
+    if (text !== this.text) {
+      this.text = text;
+      this._needsUpdateLayout = true;
+      this.labelPool.updateAtlas(text);
+      this._updateLayoutIfNeeded();
+    }
+  }
+
+  private _needsUpdateLayout = false;
+  private _updateLayoutIfNeeded() {
+    if (!this._needsUpdateLayout) {
+      return;
+    }
+    const layoutInfo = this.labelPool.fontManager.layout(this.text);
+    this.material.uniforms.uLabelSize!.value[0] = layoutInfo.width;
+    this.material.uniforms.uLabelSize!.value[1] = layoutInfo.height;
+    this.pickingMaterial.uniforms.uLabelSize!.value[0] = layoutInfo.width;
+    this.pickingMaterial.uniforms.uLabelSize!.value[1] = layoutInfo.height;
+
+    this.geometry.instanceCount = this.mesh.count = layoutInfo.chars.length;
+
+    this.reallocateAttributeBufferIfNeeded(layoutInfo.chars.length);
+
+    let i = 0;
+    for (const char of layoutInfo.chars) {
+      // instanceBoxPosition
+      this.instanceAttrData[i++] = char.left;
+      this.instanceAttrData[i++] = layoutInfo.height - char.boxTop - char.boxHeight;
+      // instanceCharPosition
+      this.instanceAttrData[i++] = char.left;
+      this.instanceAttrData[i++] =
+        layoutInfo.height - char.boxTop - char.boxHeight + char.top - char.boxTop;
+      // instanceUv
+      this.instanceAttrData[i++] = char.atlasX;
+      this.instanceAttrData[i++] = char.atlasY;
+      // instanceBoxSize
+      this.instanceAttrData[i++] = char.xAdvance;
+      this.instanceAttrData[i++] = char.boxHeight;
+      // instanceCharSize
+      this.instanceAttrData[i++] = char.width;
+      this.instanceAttrData[i++] = char.height;
+    }
+    this.instanceAttrBuffer.needsUpdate = true;
+    this._needsUpdateLayout = false;
+  }
+
+  /** Values should be in working (linear-srgb) color space */
+  setColor(r: number, g: number, b: number, a = 1): void {
+    this.material.uniforms.uColor!.value[0] = r;
+    this.material.uniforms.uColor!.value[1] = g;
+    this.material.uniforms.uColor!.value[2] = b;
+    this.material.uniforms.uColor!.value[3] = a;
+    this.#updateTransparency();
+  }
+
+  /** Values should be in working (linear-srgb) color space */
+  setBackgroundColor(r: number, g: number, b: number, a = 1): void {
+    this.material.uniforms.uBackgroundColor!.value[0] = r;
+    this.material.uniforms.uBackgroundColor!.value[1] = g;
+    this.material.uniforms.uBackgroundColor!.value[2] = b;
+    this.material.uniforms.uBackgroundColor!.value[3] = a;
+    this.#updateTransparency();
+  }
+
+  #updateTransparency(): void {
+    const bgOpacity = this.material.uniforms.uBackgroundColor!.value[3];
+    const fgOpacity = this.material.uniforms.uColor!.value[3];
+    const transparent = bgOpacity < 1 || fgOpacity < 1;
+    this.material.transparent = transparent;
+    this.material.depthWrite = !transparent;
+  }
+
+  // eslint-disable-next-line @foxglove/no-boolean-parameters
+  setBillboard(billboard: boolean): void {
+    this.material.uniforms.uBillboard!.value = billboard;
+    this.pickingMaterial.uniforms.uBillboard!.value = billboard;
+  }
+
+  /**
+   * Enable or disable size attenuation. Setting this to `false` also requires that billboarding is
+   * enabled.
+   */
+  // eslint-disable-next-line @foxglove/no-boolean-parameters
+  setSizeAttenuation(sizeAttenuation: boolean): void {
+    this.material.uniforms.uSizeAttenuation!.value = sizeAttenuation;
+    this.pickingMaterial.uniforms.uSizeAttenuation!.value = sizeAttenuation;
+  }
+
+  setAnchorPoint(x: number, y: number): void {
+    this.material.uniforms.uAnchorPoint!.value[0] = x;
+    this.material.uniforms.uAnchorPoint!.value[1] = y;
+    this.pickingMaterial.uniforms.uAnchorPoint!.value[0] = x;
+    this.pickingMaterial.uniforms.uAnchorPoint!.value[1] = y;
+  }
+
+  setLineHeight(lineHeight: number): void {
+    this.lineHeight = lineHeight;
+    const scale =
+      (this.lineHeight * this.labelPool.scaleFactor) /
+      this.labelPool.fontManager.atlasData.lineHeight;
+    this.material.uniforms.uScale!.value = scale;
+    this.pickingMaterial.uniforms.uScale!.value = scale;
+  }
+}

src/LabelMaterial.ts

@@ -0,0 +1,145 @@
+import * as THREE from "three";
+
+export class LabelMaterial extends THREE.ShaderMaterial {
+  picking: boolean;
+
+  constructor(params: {
+    atlasTexture?: THREE.Texture<THREE.DataTextureImageData>;
+    picking?: boolean;
+  }) {
+    super({
+      glslVersion: THREE.GLSL3,
+      vertexShader: /* glsl */ `\
+#include <common>
+#include <logdepthbuf_pars_vertex>
+
+precision highp float;
+precision highp int;
+
+uniform bool uBillboard;
+uniform bool uSizeAttenuation;
+uniform float uScale;
+uniform vec2 uLabelSize;
+uniform vec2 uTextureSize;
+uniform vec2 uAnchorPoint;
+uniform vec2 uCanvasSize;
+
+in vec2 instanceBoxPosition, instanceCharPosition;
+in vec2 instanceUv;
+in vec2 instanceBoxSize, instanceCharSize;
+out mediump vec2 vUv;
+out mediump vec2 vInsideChar;
+out mediump vec2 vPosInLabel;
+void main() {
+  // Adjust uv coordinates so they are in the 0-1 range in the character region
+  vec2 boxUv = (uv * instanceBoxSize - (instanceCharPosition - instanceBoxPosition)) / instanceCharSize;
+  vInsideChar = boxUv;
+  vUv = (instanceUv + boxUv * instanceCharSize) / uTextureSize;
+  vec2 vertexPos = (instanceBoxPosition + position.xy * instanceBoxSize - uAnchorPoint * uLabelSize) * uScale;
+  vPosInLabel = (instanceBoxPosition + position.xy * instanceBoxSize);
+
+  // Adapted from THREE.ShaderLib.sprite
+  if (uBillboard) {
+    if (uSizeAttenuation) {
+      vec4 mvPosition = modelViewMatrix * vec4( 0.0, 0.0, 0.0, 1.0 );
+      mvPosition.xy += vertexPos;
+      gl_Position = projectionMatrix * mvPosition;
+    } else {
+      vec4 mvPosition = modelViewMatrix * vec4(0., 0., 0., 1.);
+
+      // Adapted from THREE.ShaderLib.sprite
+      vec2 scale;
+      scale.x = length(vec3(modelMatrix[0].xyz));
+      scale.y = length(vec3(modelMatrix[1].xyz));
+
+      gl_Position = projectionMatrix * mvPosition;
+
+      // Add position after projection to maintain constant pixel size
+      gl_Position.xy += vertexPos * 2. / uCanvasSize * scale * gl_Position.w;
+    }
+  } else {
+    gl_Position = projectionMatrix * modelViewMatrix * vec4(vertexPos, 0.0, 1.0);
+  }
+
+  #include <logdepthbuf_vertex>
+}
+`,
+      fragmentShader:
+        params.picking === true
+          ? /* glsl */ `\
+#ifdef GL_FRAGMENT_PRECISION_HIGH
+  precision highp float;
+#else
+  precision mediump float;
+#endif
+
+#include <logdepthbuf_pars_fragment>
+
+uniform vec4 objectId;
+out vec4 outColor;
+void main() {
+  outColor = objectId;
+
+  #include <logdepthbuf_fragment>
+}
+`
+          : /* glsl */ `\
+#ifdef GL_FRAGMENT_PRECISION_HIGH
+  precision highp float;
+#else
+  precision mediump float;
+#endif
+
+#include <logdepthbuf_pars_fragment>
+
+uniform sampler2D uMap;
+uniform mediump vec4 uColor, uBackgroundColor;
+uniform float uScale;
+uniform vec2 uLabelSize;
+in mediump vec2 vUv;
+in mediump vec2 vPosInLabel;
+in mediump vec2 vInsideChar;
+out vec4 outColor;
+
+// From https://github.com/Jam3/three-bmfont-text/blob/e17efbe4e9392a83d4c5ee35c67eca5a11a13395/shaders/sdf.js
+float aastep(float threshold, float value) {
+  float afwidth = length(vec2(dFdx(value), dFdy(value))) * 0.70710678118654757;
+  return smoothstep(threshold - afwidth, threshold + afwidth, value);
+}
+
+void main() {
+  float dist = texture(uMap, vUv).a;
+  vec4 color = uBackgroundColor * (1.0 - dist) + uColor * dist;
+  outColor = mix(uBackgroundColor, uColor, aastep(0.75, dist));
+
+  bool insideChar = vInsideChar.x >= 0.0 && vInsideChar.x <= 1.0 && vInsideChar.y >= 0.0 && vInsideChar.y <= 1.0;
+  outColor = insideChar ? outColor : uBackgroundColor;
+  outColor = sRGBTransferOETF(outColor); // assumes output encoding is srgb
+
+  #include <logdepthbuf_fragment>
+}
+`,
+      uniforms: {
+        objectId: { value: [NaN, NaN, NaN, NaN] },
+        uAnchorPoint: { value: [0.5, 0.5] },
+        uBillboard: { value: false },
+        uSizeAttenuation: { value: true },
+        uLabelSize: { value: [0, 0] },
+        uCanvasSize: { value: [0, 0] },
+        uScale: { value: 0 },
+        uTextureSize: {
+          value: [params.atlasTexture?.image.width ?? 0, params.atlasTexture?.image.height ?? 0],
+        },
+        uMap: { value: params.atlasTexture },
+        uColor: { value: [0, 0, 0, 1] },
+        uBackgroundColor: { value: [1, 1, 1, 1] },
+      },
+
+      side: THREE.DoubleSide,
+      transparent: false,
+      depthWrite: true,
+    });
+
+    this.picking = params.picking ?? false;
+  }
+}

src/LabelPool.stories.tsx

@@ -3,7 +3,7 @@ import { useEffect, useRef, useState, type ReactElement } from "react";
 import * as THREE from "three";
 import { OrbitControls } from "three/examples/jsm/controls/OrbitControls.js";
 
-import type { Label } from "./LabelPool.ts";
+import type { Label } from "./Label.ts";
 import { LabelPool } from "./LabelPool.ts";
 
 const meta: Meta<typeof BasicTemplate> = {