Updates to material chunks override functionality (#7284)

mvaligursky · Martin Valigursky · Martin Valigursky · commit 51caad7dce98 · 2025-01-17T11:54:32.000Z
* Updates to material chunks override functionality

* lint

* cleanup

---------

Co-authored-by: Martin Valigursky &lt;mvaligursky@snapchat.com&gt;
diff --git a/src/scene/materials/lit-material.js b/src/scene/materials/lit-material.js
@@ -24,8 +24,6 @@ class LitMaterial extends Material {
 
     shaderChunk = 'void evaluateFrontend() {}\n';
 
-    chunks = null;
-
     useLighting = true;
 
     useFog = true;
diff --git a/src/scene/materials/material.js b/src/scene/materials/material.js
@@ -173,13 +173,39 @@ class Material {
      */
     stencilBack = null;
 
+    /**
+     * @type {Object<string, string>}
+     * @private
+     */
+    _chunks = { };
+
     /** @protected */
     constructor() {
         if (new.target === Material) {
             Debug.error('Material class cannot be instantiated, use ShaderMaterial instead');
         }
     }
 
+    /**
+     * Sets the object containing custom shader chunks that will replace default ones.
+     *
+     * @type {Object<string, string>}
+     */
+    set chunks(value) {
+        this.clearVariants();
+        this._chunks = value;
+    }
+
+    /**
+     * Gets the object containing custom shader chunks.
+     *
+     * @type {Object<string, string>}
+     */
+    get chunks() {
+        this.clearVariants();
+        return this._chunks;
+    }
+
     /**
      * Sets the offset for the output depth buffer value. Useful for decals to prevent z-fighting.
      * Typically a small negative value (-0.1) is used to render the mesh slightly closer to the
@@ -534,6 +560,14 @@ class Material {
         this.defines.clear();
         source.defines.forEach((value, key) => this.defines.set(key, value));
 
+        // chunks
+        const srcChunks = source._chunks;
+        for (const p in srcChunks) {
+            if (srcChunks.hasOwnProperty(p)) {
+                this._chunks[p] = srcChunks[p];
+            }
+        }
+
         return this;
     }
 
diff --git a/src/scene/materials/shader-material.js b/src/scene/materials/shader-material.js
@@ -110,7 +110,8 @@ class ShaderMaterial extends Material {
             gamma: params.cameraShaderParams.shaderOutputGamma,
             toneMapping: params.cameraShaderParams.toneMapping,
             fog: params.cameraShaderParams.fog,
-            shaderDesc: this.shaderDesc
+            shaderDesc: this.shaderDesc,
+            chunks: this.chunks ?? {} // override chunks from the material
         };
 
         const processingOptions = new ShaderProcessorOptions(params.viewUniformFormat, params.viewBindGroupFormat, params.vertexFormat);
diff --git a/src/scene/materials/standard-material.js b/src/scene/materials/standard-material.js
@@ -588,29 +588,9 @@ class StandardMaterial extends Material {
             this[`_${name}`] = _props[name].value();
         });
 
-        /**
-         * @type {Object<string, string>}
-         * @private
-         */
-        this._chunks = { };
         this._uniformCache = { };
     }
 
-    /**
-     * Object containing custom shader chunks that will replace default ones.
-     *
-     * @type {Object<string, string>}
-     */
-    set chunks(value) {
-        this._dirtyShader = true;
-        this._chunks = value;
-    }
-
-    get chunks() {
-        this._dirtyShader = true;
-        return this._chunks;
-    }
-
     /**
      * Copy a `StandardMaterial`.
      *
@@ -625,13 +605,6 @@ class StandardMaterial extends Material {
             this[k] = source[k];
         });
 
-        // clone chunks
-        for (const p in source._chunks) {
-            if (source._chunks.hasOwnProperty(p)) {
-                this._chunks[p] = source._chunks[p];
-            }
-        }
-
         // clone user attributes
         this.userAttributes = new Map(source.userAttributes);
 
diff --git a/src/scene/shader-lib/programs/lit-shader.js b/src/scene/shader-lib/programs/lit-shader.js
@@ -1,8 +1,7 @@
 import {
     SEMANTIC_ATTR8, SEMANTIC_ATTR9, SEMANTIC_ATTR12, SEMANTIC_ATTR13, SEMANTIC_ATTR14, SEMANTIC_ATTR15,
     SEMANTIC_BLENDINDICES, SEMANTIC_BLENDWEIGHT, SEMANTIC_COLOR, SEMANTIC_NORMAL, SEMANTIC_POSITION, SEMANTIC_TANGENT,
-    SEMANTIC_TEXCOORD0, SEMANTIC_TEXCOORD1,
-    SHADERTAG_MATERIAL
+    SEMANTIC_TEXCOORD0, SEMANTIC_TEXCOORD1
 } from '../../../platform/graphics/constants.js';
 import {
     BLEND_ADDITIVEALPHA, BLEND_NORMAL, BLEND_PREMULTIPLIED,
@@ -21,7 +20,6 @@ import { ChunkUtils } from '../chunk-utils.js';
 import { LightsBuffer } from '../../lighting/lights-buffer.js';
 import { ShaderPass } from '../../shader-pass.js';
 import { validateUserChunks } from '../chunks/chunk-validation.js';
-import { ShaderUtils } from '../../../platform/graphics/shader-utils.js';
 import { ChunkBuilder } from '../chunk-builder.js';
 import { ShaderGenerator } from './shader-generator.js';
 import { Debug } from '../../../core/debug.js';
@@ -1531,33 +1529,6 @@ class LitShader {
 
         Debug.assert(!this.fshader.includes('litShaderArgs'), 'Automatic compatibility with shaders using litShaderArgs has been removed. Please update the shader to use the new system.');
     }
-
-    getDefinition(options) {
-
-        const vIncludes = new Map();
-        vIncludes.set('transformCoreVS', this.chunks.transformCoreVS);
-        vIncludes.set('transformInstancingVS', this.chunks.transformInstancingVS);
-        vIncludes.set('skinVS', this.chunks.skinVS);
-        vIncludes.set('skinBatchVS', this.chunks.skinBatchVS);
-
-        const defines = new Map(options.defines);
-
-        const definition = ShaderUtils.createDefinition(this.device, {
-            name: 'LitShader',
-            attributes: this.attributes,
-            vertexCode: this.vshader,
-            fragmentCode: this.fshader,
-            vertexIncludes: vIncludes,
-            fragmentDefines: defines,
-            vertexDefines: defines
-        });
-
-        if (this.shaderPassInfo.isForward) {
-            definition.tag = SHADERTAG_MATERIAL;
-        }
-
-        return definition;
-    }
 }
 
 export { LitShader };
diff --git a/src/scene/shader-lib/programs/lit.js b/src/scene/shader-lib/programs/lit.js
@@ -2,9 +2,12 @@ import { ChunkBuilder } from '../chunk-builder.js';
 import { LitShader } from './lit-shader.js';
 import { LitOptionsUtils } from './lit-options-utils.js';
 import { ShaderGenerator } from './shader-generator.js';
+import { SHADERTAG_MATERIAL } from '../../../platform/graphics/constants.js';
+import { ShaderUtils } from '../../../platform/graphics/shader-utils.js';
 
 /**
  * @import { GraphicsDevice } from '../../../platform/graphics/graphics-device.js'
+ * @import { LitMaterialOptions } from '../../materials/lit-material-options.js'
  */
 
 const dummyUvs = [0, 1, 2, 3, 4, 5, 6, 7];
@@ -24,7 +27,7 @@ class ShaderGeneratorLit extends ShaderGenerator {
 
     /**
      * @param {GraphicsDevice} device - The graphics device.
-     * @param {object} options - The options to be passed to the backend.
+     * @param {LitMaterialOptions} options - The options to be passed to the backend.
      * @returns {object} Returns the created shader definition.
      */
     createShaderDefinition(device, options) {
@@ -47,7 +50,29 @@ class ShaderGeneratorLit extends ShaderGenerator {
         litShader.generateVertexShader(usedUvSets, usedUvSets, mapTransforms);
         litShader.generateFragmentShader(decl.code, code.code, func.code, 'vUv0');
 
-        return litShader.getDefinition(options);
+        const vIncludes = new Map(Object.entries({
+            ...Object.getPrototypeOf(litShader.chunks), // the prototype stores the default chunks
+            ...litShader.chunks,  // user overrides are supplied as instance properties
+            ...options.litOptions.chunks
+        }));
+
+        const defines = new Map(options.defines);
+
+        const definition = ShaderUtils.createDefinition(device, {
+            name: 'LitShader',
+            attributes: litShader.attributes,
+            vertexCode: litShader.vshader,
+            fragmentCode: litShader.fshader,
+            vertexIncludes: vIncludes,
+            fragmentDefines: defines,
+            vertexDefines: defines
+        });
+
+        if (litShader.shaderPassInfo.isForward) {
+            definition.tag = SHADERTAG_MATERIAL;
+        }
+
+        return definition;
     }
 }
 
diff --git a/src/scene/shader-lib/programs/shader-generator-shader.js b/src/scene/shader-lib/programs/shader-generator-shader.js
@@ -22,6 +22,10 @@ const fShader = `
 
 class ShaderGeneratorShader extends ShaderGenerator {
     generateKey(options) {
+
+        // Note: options.chunks are not included in the key as currently shader variants are removed
+        // from the material when its chunks are modified.
+
         const desc = options.shaderDesc;
         const vsHash = desc.vertexCode ? hashCode(desc.vertexCode) : 0;
         const fsHash = desc.fragmentCode ? hashCode(desc.fragmentCode) : 0;
@@ -66,7 +70,7 @@ class ShaderGeneratorShader extends ShaderGenerator {
         defines.set('GAMMA', gammaNames[options.gamma]);
     }
 
-    createVertexDefinition(definitionOptions, options, shaderPassInfo) {
+    createVertexDefinition(definitionOptions, options, shaderPassInfo, sharedIncludes) {
 
         const desc = options.shaderDesc;
 
@@ -77,10 +81,7 @@ class ShaderGeneratorShader extends ShaderGenerator {
             definitionOptions.vertexCode = desc.vertexCode;
 
         } else {
-            const includes = new Map(Object.entries({
-                ...shaderChunks,
-                ...options.chunks
-            }));
+            const includes = new Map(sharedIncludes);
             const defines = new Map(options.defines);
             this.addSharedDefines(defines, options);
 
@@ -103,7 +104,7 @@ class ShaderGeneratorShader extends ShaderGenerator {
         }
     }
 
-    createFragmentDefinition(definitionOptions, options, shaderPassInfo) {
+    createFragmentDefinition(definitionOptions, options, shaderPassInfo, sharedIncludes) {
 
         const desc = options.shaderDesc;
 
@@ -114,11 +115,7 @@ class ShaderGeneratorShader extends ShaderGenerator {
             definitionOptions.fragmentCode = desc.fragmentCode;
 
         } else {
-            const includes = new Map(Object.entries({
-                ...shaderChunks,
-                ...options.chunks
-            }));
-
+            const includes = new Map(sharedIncludes);
             includes.set('shaderPassDefines', shaderPassInfo.shaderDefines);
             includes.set('gamma', ShaderGenerator.gammaCode(options.gamma));
             includes.set('fog', ShaderGenerator.fogCode(options.fog));
@@ -146,9 +143,14 @@ class ShaderGeneratorShader extends ShaderGenerator {
             meshBindGroupFormat: desc.meshBindGroupFormat
         };
 
+        const sharedIncludes = new Map(Object.entries({
+            ...shaderChunks,  // default chunks
+            ...options.chunks // material override chunks
+        }));
+
         this.createAttributesDefinition(definitionOptions, options);
-        this.createVertexDefinition(definitionOptions, options, shaderPassInfo);
-        this.createFragmentDefinition(definitionOptions, options, shaderPassInfo);
+        this.createVertexDefinition(definitionOptions, options, shaderPassInfo, sharedIncludes);
+        this.createFragmentDefinition(definitionOptions, options, shaderPassInfo, sharedIncludes);
 
         return ShaderUtils.createDefinition(device, definitionOptions);
     }
diff --git a/src/scene/shader-lib/programs/standard.js b/src/scene/shader-lib/programs/standard.js
@@ -11,6 +11,8 @@ import { ChunkUtils } from '../chunk-utils.js';
 import { StandardMaterialOptions } from '../../materials/standard-material-options.js';
 import { LitOptionsUtils } from './lit-options-utils.js';
 import { ShaderGenerator } from './shader-generator.js';
+import { ShaderUtils } from '../../../platform/graphics/shader-utils.js';
+import { SHADERTAG_MATERIAL } from '../../../platform/graphics/constants.js';
 
 /**
  * @import { GraphicsDevice } from '../../../platform/graphics/graphics-device.js'
@@ -528,7 +530,29 @@ class ShaderGeneratorStandard extends ShaderGenerator {
 
         litShader.generateFragmentShader(decl.code, code.code, func.code, lightingUv);
 
-        return litShader.getDefinition(options);
+        const vIncludes = new Map(Object.entries({
+            ...Object.getPrototypeOf(litShader.chunks), // the prototype stores the default chunks
+            ...litShader.chunks,  // user overrides are supplied as instance properties
+            ...options.litOptions.chunks
+        }));
+
+        const defines = new Map(options.defines);
+
+        const definition = ShaderUtils.createDefinition(device, {
+            name: 'StandardShader',
+            attributes: litShader.attributes,
+            vertexCode: litShader.vshader,
+            fragmentCode: litShader.fshader,
+            vertexIncludes: vIncludes,
+            fragmentDefines: defines,
+            vertexDefines: defines
+        });
+
+        if (litShader.shaderPassInfo.isForward) {
+            definition.tag = SHADERTAG_MATERIAL;
+        }
+
+        return definition;
     }
 }
 
