feat: add InstancedUniformsMesh class for setting shader uniforms per…

… instance

packages/troika-three-utils/README.md

@@ -62,6 +62,42 @@ mesh.customDepthMaterial = customMaterial.getDepthMaterial() //for shadows
 You can also declare custom `uniforms` and `defines`, inject fragment shader code to modify the output color, etc. See the JSDoc in [DerivedMaterial.js](./src/DerivedMaterial.js) for full details.
 
 
+### InstancedUniformsMesh
+
+[Source code](./src/InstancedUniformsMesh.js)
+
+This extends Three.js's [`InstancedMesh`](https://threejs.org/docs/#api/en/objects/InstancedMesh) to allow any of its material's shader uniforms to be set individually per instance. It behaves just like `InstancedMesh` but exposes a new `setUniformAt(uniformName, instanceIndex, value)` method.
+
+When you call `setUniformAt`, the geometry and the material shaders will be automatically upgraded behind the scenes to turn that uniform into an instanced buffer attribute, filling in the other indices with the uniform's default value. You can do this for any uniform of type `float`, `vec2`, `vec3`, or `vec4`. It works both for built-in Three.js materials and also for any custom ShaderMaterial.
+
+For example, here is how you could set random `emissive` and `metalness` values for each instance using a `MeshStandardMaterial`:
+
+```js
+import { InstancedUniformsMesh } from 'troika-three-utils'
+
+const count = 100
+const mesh = new InstancedUniformsMesh(
+  someGeometry,
+  new MeshStandardMaterial(),
+  count
+)
+const color = new Color()
+for (let i = 0; i < count; i++) {
+  mesh.setUniformAt('metalness', i, Math.random())
+  mesh.setUniformAt('emissive', i, color.set(Math.random() * 0xffffff))
+}
+```
+
+The type of the `value` argument should match the type of the uniform defined in the material's shader:
+
+| For uniform type: | Pass a value of this type:                |
+| ----------------- | ----------------------------------------- |
+| float             | Number                                    |
+| vec2              | `Vector2` or Array w/ length=2            |
+| vec3              | `Vector3` or `Color` or Array w/ length=3 |
+| vec4              | `Vector4` or Array w/ length=4            |
+
+
 ### BezierMesh
 
 _[Source code with JSDoc](./src/BezierMesh.js)_ | _[Online example](https://troika-examples.netlify.com/#bezier3d)_

packages/troika-three-utils/src/InstancedUniformsDerivedMaterial.js

@@ -0,0 +1,58 @@
+import { createDerivedMaterial } from './DerivedMaterial.js'
+import { getShaderUniformTypes } from './getShaderUniformTypes.js'
+import { voidMainRegExp } from './voidMainRegExp.js'
+
+const precededByUniformRE = /\buniform\s+(int|float|vec[234])\s+$/
+const attrRefReplacer = (name, index, str) => (precededByUniformRE.test(str.substr(0, index)) ? name : `troika_${name}`)
+const varyingRefReplacer = (name, index, str) => (precededByUniformRE.test(str.substr(0, index)) ? name : `troika_vary_${name}`)
+
+export function createInstancedUniformsDerivedMaterial (baseMaterial, uniformNames) {
+  const derived = createDerivedMaterial(baseMaterial, {
+    defines: {
+      TROIKA_INSTANCED_UNIFORMS: uniformNames.sort().join('|')
+    },
+
+    customRewriter ({ vertexShader, fragmentShader }) {
+      let vertexDeclarations = []
+      let vertexAssignments = []
+      let fragmentDeclarations = []
+
+      // Find what uniforms are declared in which shader and their types
+      let vertexUniforms = getShaderUniformTypes(vertexShader)
+      let fragmentUniforms = getShaderUniformTypes(fragmentShader)
+
+      // Add attributes and varyings for, and rewrite references to, instanceUniforms
+      uniformNames.forEach((name) => {
+        let vertType = vertexUniforms[name]
+        let fragType = fragmentUniforms[name]
+        if (vertType || fragType) {
+          let finder = new RegExp(`\\b${name}\\b`, 'g')
+          vertexDeclarations.push(`attribute ${vertType || fragType} troika_attr_${name};`)
+          if (vertType) {
+            vertexShader = vertexShader.replace(finder, attrRefReplacer)
+          }
+          if (fragType) {
+            fragmentShader = fragmentShader.replace(finder, varyingRefReplacer)
+            let varyingDecl = `varying ${fragType} troika_vary_${name};`
+            vertexDeclarations.push(varyingDecl)
+            fragmentDeclarations.push(varyingDecl)
+            vertexAssignments.push(`troika_vary_${name} = troika_attr_${name};`)
+          }
+        }
+      })
+
+      // Inject vertex shader declarations and assignments
+      vertexShader = `${vertexDeclarations.join('\n')}\n${vertexShader.replace(voidMainRegExp, `\n$&\n${vertexAssignments.join('\n')}`)}`
+
+      // Inject fragment shader declarations
+      if (fragmentDeclarations.length) {
+        fragmentShader = `${fragmentDeclarations.join('\n')}\n${fragmentShader}`
+      }
+
+      return { vertexShader, fragmentShader }
+    }
+  })
+
+  derived.isInstancedUniformsMaterial = true
+  return derived
+}

packages/troika-three-utils/src/InstancedUniformsMesh.js

@@ -0,0 +1,135 @@
+import { InstancedBufferAttribute, InstancedMesh, MeshBasicMaterial } from 'three'
+import { getShadersForMaterial } from './getShadersForMaterial.js'
+import { createInstancedUniformsDerivedMaterial } from './InstancedUniformsDerivedMaterial.js'
+
+const defaultMaterial = new MeshBasicMaterial()
+
+export class InstancedUniformsMesh extends InstancedMesh {
+  constructor (geometry, material, count) {
+    super(geometry, material, count)
+    this._instancedUniformNames = [] //treated as immutable
+  }
+
+  /*
+   * Getter/setter for automatically wrapping the user-supplied geometry with one that will
+   * carry our extra InstancedBufferAttribute(s)
+   */
+  get geometry () {
+    return this._derivedGeometry
+  }
+
+  set geometry (geometry) {
+    // Extend the geometry so we can add our instancing attributes but inherit everything else
+    if (geometry) {
+      geometry = Object.create(geometry)
+      geometry.attributes = Object.create(geometry.attributes)
+    }
+    this._derivedGeometry = geometry
+  }
+
+  /*
+   * Getter/setter for automatically wrapping the user-supplied material with our upgrades. We do the
+   * wrapping lazily on _read_ rather than write to avoid unnecessary wrapping on transient values.
+   */
+  get material () {
+    let derivedMaterial = this._derivedMaterial
+    const baseMaterial = this._baseMaterial || this._defaultMaterial || (this._defaultMaterial = defaultMaterial.clone())
+    const uniformNames = this._instancedUniformNames
+    if (!derivedMaterial || derivedMaterial.baseMaterial !== baseMaterial || derivedMaterial._instancedUniformNames !== uniformNames) {
+      derivedMaterial = this._derivedMaterial = createInstancedUniformsDerivedMaterial(baseMaterial, uniformNames)
+      derivedMaterial._instancedUniformNames = uniformNames
+      // dispose the derived material when its base material is disposed:
+      baseMaterial.addEventListener('dispose', function onDispose () {
+        baseMaterial.removeEventListener('dispose', onDispose)
+        derivedMaterial.dispose()
+      })
+    }
+    return derivedMaterial
+  }
+
+  set material (baseMaterial) {
+    if (Array.isArray(baseMaterial)) {
+      throw new Error('InstancedUniformsMesh does not support multiple materials')
+    }
+    // Unwrap already-derived materials
+    while (baseMaterial && baseMaterial.isInstancedUniformsMaterial) {
+      baseMaterial = baseMaterial.baseMaterial
+    }
+    this._baseMaterial = baseMaterial
+  }
+
+  get customDepthMaterial () {
+    return this.material.getDepthMaterial()
+  }
+
+  get customDistanceMaterial () {
+    return this.material.getDistanceMaterial()
+  }
+
+  /**
+   * Set the value of a shader uniform for a single instance.
+   * @param {string} name - the name of the shader uniform
+   * @param {number} index - the index of the instance to set the value for
+   * @param {number|Vector2|Vector3|Vector4|Color|Array} value - the uniform value for this instance
+   */
+  setUniformAt (name, index, value) {
+    const attrs = this.geometry.attributes
+    const attrName = `troika_attr_${name}`
+    let attr = attrs[attrName]
+    if (!attr) {
+      const defaultValue = getDefaultUniformValue(this._baseMaterial, name)
+      const itemSize = getItemSizeForValue(defaultValue)
+      attr = attrs[attrName] = new InstancedBufferAttribute(new Float32Array(itemSize * this.count), itemSize)
+      // Fill with default value:
+      if (defaultValue !== null) {
+        for (let i = 0; i < this.count; i++) {
+          setAttributeValue(attr, i, defaultValue)
+        }
+      }
+      this._instancedUniformNames = [...this._instancedUniformNames, name]
+    }
+    setAttributeValue(attr, index, value)
+    attr.needsUpdate = true
+  }
+}
+
+function setAttributeValue (attr, index, value) {
+  let size = attr.itemSize
+  if (size === 1) {
+    attr.setX(index, value)
+  } else if (size === 2) {
+    attr.setXY(index, value.x, value.y)
+  } else if (size === 3) {
+    if (value.isColor) {
+      attr.setXYZ(index, value.r, value.g, value.b)
+    } else {
+      attr.setXYZ(index, value.x, value.y, value.z)
+    }
+  } else if (size === 4) {
+    attr.setXYZW(index, value.x, value.y, value.z, value.w)
+  }
+}
+
+function getDefaultUniformValue (material, name) {
+  // Try uniforms on the material itself, then try the builtin material shaders
+  let uniforms = material.uniforms
+  if (uniforms && uniforms[name]) {
+    return uniforms[name].value
+  }
+  uniforms = getShadersForMaterial(material).uniforms
+  if (uniforms && uniforms[name]) {
+    return uniforms[name].value
+  }
+  return null
+}
+
+function getItemSizeForValue (value) {
+  return value == null ? 0
+    : typeof value === 'number' ? 1
+    : value.isVector2 ? 2
+    : value.isVector3 || value.isColor ? 3
+    : value.isVector4 ? 4
+    : Array.isArray(value) ? value.length
+    : 0
+}
+

packages/troika-three-utils/src/getShadersForMaterial.js

@@ -10,7 +10,7 @@ const MATERIAL_TYPES_TO_SHADERS = {
   MeshBasicMaterial: 'basic',
   MeshLambertMaterial: 'lambert',
   MeshPhongMaterial: 'phong',
-  MeshToonMaterial: 'phong',
+  MeshToonMaterial: 'toon',
   MeshStandardMaterial: 'physical',
   MeshPhysicalMaterial: 'physical',
   MeshMatcapMaterial: 'matcap',
@@ -31,4 +31,4 @@ const MATERIAL_TYPES_TO_SHADERS = {
 export function getShadersForMaterial(material) {
   let builtinType = MATERIAL_TYPES_TO_SHADERS[material.type]
   return builtinType ? ShaderLib[builtinType] : material //TODO fallback for unknown type?
-}
+}

packages/troika-three-utils/src/index.js

@@ -1,10 +1,9 @@
 // Troika Three.js Utilities exports
 
-
-export {createDerivedMaterial} from './DerivedMaterial.js'
-export {getShadersForMaterial} from './getShadersForMaterial.js'
-export {getShaderUniformTypes} from './getShaderUniformTypes.js'
-export {expandShaderIncludes} from './expandShaderIncludes.js'
-export {ShaderFloatArray} from './ShaderFloatArray.js'
-export {voidMainRegExp} from './voidMainRegExp.js'
-export {BezierMesh} from './BezierMesh.js'
+export { createDerivedMaterial } from './DerivedMaterial.js'
+export { getShadersForMaterial } from './getShadersForMaterial.js'
+export { getShaderUniformTypes } from './getShaderUniformTypes.js'
+export { expandShaderIncludes } from './expandShaderIncludes.js'
+export { voidMainRegExp } from './voidMainRegExp.js'
+export { InstancedUniformsMesh } from './InstancedUniformsMesh.js'
+export { BezierMesh } from './BezierMesh.js'