Refactor shader caching to handle mesh instance on layers with different lights

src/deprecated/deprecated.js

@@ -868,12 +868,6 @@ ForwardRenderer.prototype.renderComposition = function (comp) {
     getApplication().renderComposition(comp);
 };
 
-ForwardRenderer.prototype.updateShader = function (meshInstance, objDefs, unused, pass, sortedLights) {
-    Debug.deprecated('pc.ForwardRenderer#updateShader is deprecated, use pc.MeshInstance#updatePassShader.');
-    const scene = meshInstance.material._scene || getApplication().scene;
-    return meshInstance.updatePassShader(scene, pass, sortedLights);
-};
-
 MeshInstance.prototype.syncAabb = function () {
     Debug.deprecated('pc.MeshInstance#syncAabb is deprecated.');
 };

src/scene/materials/basic-material.js

@@ -1,4 +1,3 @@
-import { Debug } from '../../core/debug.js';
 import { Color } from '../../core/math/color.js';
 
 import { ShaderProcessorOptions } from '../../platform/graphics/shader-processor-options.js';
@@ -88,14 +87,6 @@ class BasicMaterial extends Material {
 
     getShaderVariant(device, scene, objDefs, unused, pass, sortedLights, viewUniformFormat, viewBindGroupFormat, vertexFormat) {
 
-        // Note: this is deprecated function Editor and possibly other projects use: they define
-        // updateShader callback on their BasicMaterial, so we handle it here.
-        if (this.updateShader) {
-            Debug.deprecated('pc.BasicMaterial.updateShader is deprecated');
-            this.updateShader(device, scene, objDefs, null, pass, sortedLights);
-            return this.shader;
-        }
-
         const options = {
             skin: objDefs && (objDefs & SHADERDEF_SKIN) !== 0,
             screenSpace: objDefs && (objDefs & SHADERDEF_SCREENSPACE) !== 0,

src/scene/materials/material.js

@@ -79,7 +79,14 @@ class Material {
 
     id = id++;
 
-    variants = {};
+    /**
+     * The cache of shader variants generated for this material. The key represents the unique
+     * variant, the value is the shader.
+     *
+     * @type {Map<string, import('../../platform/graphics/shader.js').Shader>}
+     * @ignore
+     */
+    variants = new Map();
 
     parameters = {};
 
@@ -472,7 +479,7 @@ class Material {
     clearVariants() {
 
         // clear variants on the material
-        this.variants = {};
+        this.variants.clear();
 
         // but also clear them from all materials that reference them
         const meshInstances = this.meshInstances;
@@ -556,7 +563,7 @@ class Material {
      * are no other materials using it).
      */
     destroy() {
-        this.variants = {};
+        this.variants.clear();
         this._shader = null;
 
         for (let i = 0; i < this.meshInstances.length; i++) {

src/scene/mesh-instance.js

@@ -59,6 +59,85 @@ class Command {
     }
 }
 
+/**
+ * Internal helper class for storing the shader and related mesh bind group in the shader cache.
+ *
+ * @ignore
+ */
+class ShaderInstance {
+    /**
+     * A shader.
+     *
+     * @type {import('../platform/graphics/shader.js').Shader|undefined}
+     */
+    shader;
+
+    /**
+     * A bind group storing mesh uniforms for the shader.
+     *
+     * @type {BindGroup|null}
+     */
+    bindGroup = null;
+
+    /**
+     * Returns the mesh bind group for the shader.
+     *
+     * @param {import('../platform/graphics/graphics-device.js').GraphicsDevice} device - The
+     * graphics device.
+     * @returns {BindGroup} - The mesh bind group.
+     */
+    getBindGroup(device) {
+
+        // create bind group
+        if (!this.bindGroup) {
+            const shader = this.shader;
+            Debug.assert(shader);
+
+            // mesh uniform buffer
+            const ubFormat = shader.meshUniformBufferFormat;
+            Debug.assert(ubFormat);
+            const uniformBuffer = new UniformBuffer(device, ubFormat, false);
+
+            // mesh bind group
+            const bindGroupFormat = shader.meshBindGroupFormat;
+            Debug.assert(bindGroupFormat);
+            this.bindGroup = new BindGroup(device, bindGroupFormat, uniformBuffer);
+            DebugHelper.setName(this.bindGroup, `MeshBindGroup_${this.bindGroup.id}`);
+        }
+
+        return this.bindGroup;
+    }
+
+    destroy() {
+        const group = this.bindGroup;
+        if (group) {
+            group.defaultUniformBuffer?.destroy();
+            group.destroy();
+            this.bindGroup = null;
+        }
+    }
+}
+
+/**
+ * An entry in the shader cache, representing shaders for this mesh instance and a specific shader
+ * pass.
+ *
+ * @ignore
+ */
+class ShaderCacheEntry {
+    /**
+     * The shader instances. Looked up by lightHash, which represents an ordered set of lights.
+     *
+     * @type {Map<number, ShaderInstance>}
+     */
+    shaderInstances = new Map();
+
+    destroy() {
+        this.shaderInstances.forEach(instance => instance.destroy());
+        this.shaderInstances.clear();
+    }
+}
+
 /**
  * Callback used by {@link Layer} to calculate the "sort distance" for a {@link MeshInstance},
  * which determines its place in the render order.
@@ -105,27 +184,18 @@ class MeshInstance {
     transparent = false;
 
     /**
-     * @type {import('./materials/material.js').Material}
+     * @type {import('./materials/material.js').Material|null}
      * @private
      */
-    _material;
+    _material = null;
 
     /**
-     * An array of shaders used by the mesh instance, indexed by the shader pass constant (SHADER_FORWARD..)
+     * An array of shader cache entries, indexed by the shader pass constant (SHADER_FORWARD..). The
+     * value stores all shaders and bind groups for the shader pass for various light combinations.
      *
-     * @type {Array<import('../platform/graphics/shader.js').Shader>}
-     * @ignore
+     * @type {Array<ShaderCacheEntry|null>}
      */
-    _shader = [];
-
-    /**
-     * An array of bind groups, storing uniforms per pass. This has 1:1 relation with the _shades array,
-     * and is indexed by the shader pass constant as well.
-     *
-     * @type {Array<BindGroup>}
-     * @ignore
-     */
-    _bindGroups = [];
+    _shaderCache = [];
 
     /**
      * Create a new MeshInstance instance.
@@ -178,8 +248,6 @@ class MeshInstance {
         this._shaderDefs |= mesh.vertexBuffer.format.hasColor ? SHADERDEF_VCOLOR : 0;
         this._shaderDefs |= mesh.vertexBuffer.format.hasTangents ? SHADERDEF_TANGENTS : 0;
 
-        this._lightHash = 0;
-
         // Render options
         this.layer = LAYER_WORLD; // legacy
         /** @private */
@@ -212,20 +280,21 @@ class MeshInstance {
         this.updateKey();
 
         /**
-         * @type {import('./skin-instance.js').SkinInstance}
+         * @type {import('./skin-instance.js').SkinInstance|null}
          * @private
          */
         this._skinInstance = null;
+
         /**
-         * @type {import('./morph-instance.js').MorphInstance}
+         * @type {import('./morph-instance.js').MorphInstance|null}
          * @private
          */
         this._morphInstance = null;
 
         this.instancingData = null;
 
         /**
-         * @type {BoundingBox}
+         * @type {BoundingBox|null}
          * @private
          */
         this._customAabb = null;
@@ -400,65 +469,71 @@ class MeshInstance {
     }
 
     /**
-     * Clear the internal shader array.
+     * Clear the internal shader cache.
      *
      * @ignore
      */
     clearShaders() {
-        const shaders = this._shader;
-        for (let i = 0; i < shaders.length; i++) {
-            shaders[i] = null;
+        const shaderCache = this._shaderCache;
+        for (let i = 0; i < shaderCache.length; i++) {
+            shaderCache[i]?.destroy();
+            shaderCache[i] = null;
         }
-
-        this.destroyBindGroups();
-    }
-
-    destroyBindGroups() {
-
-        const groups = this._bindGroups;
-        for (let i = 0; i < groups.length; i++) {
-            const group = groups[i];
-            if (group) {
-                const uniformBuffer = group.defaultUniformBuffer;
-                if (uniformBuffer) {
-                    uniformBuffer.destroy();
-                }
-                group.destroy();
-            }
-        }
-        groups.length = 0;
     }
 
     /**
-     * @param {import('../platform/graphics/graphics-device.js').GraphicsDevice} device - The
-     * graphics device.
-     * @param {number} pass - Shader pass number.
-     * @returns {BindGroup} - The mesh bind group.
+     * Returns the shader instance for the specified shader pass and light hash that is compatible
+     * with this mesh instance.
+     *
+     * @param {number} shaderPass - The shader pass index.
+     * @param {number} lightHash - The hash value of the lights that are affecting this mesh instance.
+     * @param {import('./scene.js').Scene} scene - The scene.
+     * @param {import('../platform/graphics/uniform-buffer-format.js').UniformBufferFormat} [viewUniformFormat] - The
+     * format of the view uniform buffer.
+     * @param {import('../platform/graphics/bind-group-format.js').BindGroupFormat} [viewBindGroupFormat] - The
+     * format of the view bind group.
+     * @param {any} [sortedLights] - Array of arrays of lights.
+     * @returns {ShaderInstance} - the shader instance.
      * @ignore
      */
-    getBindGroup(device, pass) {
+    getShaderInstance(shaderPass, lightHash, scene, viewUniformFormat, viewBindGroupFormat, sortedLights) {
 
-        // create bind group
-        let bindGroup = this._bindGroups[pass];
-        if (!bindGroup) {
-            const shader = this._shader[pass];
-            Debug.assert(shader);
+        let shaderInstance;
+        let passEntry = this._shaderCache[shaderPass];
+        if (passEntry) {
+            shaderInstance = passEntry.shaderInstances.get(lightHash);
+        } else {
+            passEntry = new ShaderCacheEntry();
+            this._shaderCache[shaderPass] = passEntry;
+        }
 
-            // mesh uniform buffer
-            const ubFormat = shader.meshUniformBufferFormat;
-            Debug.assert(ubFormat);
-            const uniformBuffer = new UniformBuffer(device, ubFormat, false);
+        // cache miss in the shader cache of the mesh instance
+        if (!shaderInstance) {
 
-            // mesh bind group
-            const bindGroupFormat = shader.meshBindGroupFormat;
-            Debug.assert(bindGroupFormat);
-            bindGroup = new BindGroup(device, bindGroupFormat, uniformBuffer);
-            DebugHelper.setName(bindGroup, `MeshBindGroup_${bindGroup.id}`);
+            // get the shader from the material
+            const mat = this._material;
+            const shaderDefs = this._shaderDefs;
+            const variantKey = shaderPass + '_' + shaderDefs + '_' + lightHash;
+            shaderInstance = new ShaderInstance();
+            shaderInstance.shader = mat.variants.get(variantKey);
+
+            // cache miss in the material variants
+            if (!shaderInstance.shader) {
+
+                const shader = mat.getShaderVariant(this.mesh.device, scene, shaderDefs, null, shaderPass, sortedLights,
+                                                    viewUniformFormat, viewBindGroupFormat, this._mesh.vertexBuffer.format);
 
-            this._bindGroups[pass] = bindGroup;
+                // add it to the material variants cache
+                mat.variants.set(variantKey, shader);
+
+                shaderInstance.shader = shader;
+            }
+
+            // add it to the mesh instance cache
+            passEntry.shaderInstances.set(lightHash, shaderInstance);
         }
 
-        return bindGroup;
+        return shaderInstance;
     }
 
     /**
@@ -731,23 +806,6 @@ class MeshInstance {
         this._updateShaderDefs(vertexBuffer ? (this._shaderDefs | SHADERDEF_INSTANCING) : (this._shaderDefs & ~SHADERDEF_INSTANCING));
     }
 
-    /**
-     * Obtain a shader variant required to render the mesh instance within specified pass.
-     *
-     * @param {import('./scene.js').Scene} scene - The scene.
-     * @param {number} pass - The render pass.
-     * @param {any} sortedLights - Array of arrays of lights.
-     * @param {import('../platform/graphics/uniform-buffer-format.js').UniformBufferFormat} viewUniformFormat - The
-     * format of the view uniform buffer.
-     * @param {import('../platform/graphics/bind-group-format.js').BindGroupFormat} viewBindGroupFormat - The
-     * format of the view bind group.
-     * @ignore
-     */
-    updatePassShader(scene, pass, sortedLights, viewUniformFormat, viewBindGroupFormat) {
-        this._shader[pass] = this.material.getShaderVariant(this.mesh.device, scene, this._shaderDefs, null, pass, sortedLights,
-                                                            viewUniformFormat, viewBindGroupFormat, this._mesh.vertexBuffer.format);
-    }
-
     ensureMaterial(device) {
         if (!this.material) {
             Debug.warn(`Mesh attached to entity '${this.node.name}' does not have a material, using a default one.`);