Updated the way shaders include tonemapping chunks (#7174)

* Updated the way shaders include tonemapping chunks

* lint

---------

Co-authored-by: Martin Valigursky <[email protected]>

src/extras/render-passes/render-pass-compose.js

@@ -2,9 +2,7 @@ import { math } from '../../core/math/math.js';
 import { Color } from '../../core/math/color.js';
 import { RenderPassShaderQuad } from '../../scene/graphics/render-pass-shader-quad.js';
 import { shaderChunks } from '../../scene/shader-lib/chunks/chunks.js';
-import { TONEMAP_LINEAR } from '../../scene/constants.js';
-import { ShaderGenerator } from '../../scene/shader-lib/programs/shader-generator.js';
-
+import { TONEMAP_LINEAR, tonemapNames } from '../../scene/constants.js';
 
 // Contrast Adaptive Sharpening (CAS) is used to apply the sharpening. It's based on AMD's
 // FidelityFX CAS, WebGL implementation: https://www.shadertoy.com/view/wtlSWB. It's best to run it
@@ -13,6 +11,11 @@ import { ShaderGenerator } from '../../scene/shader-lib/programs/shader-generato
 // before the tone-mapping.
 
 const fragmentShader = /* glsl */ `
+
+    #include "tonemappingPS"
+    #include "decodePS"
+    #include "gamma2_2PS"
+
     varying vec2 uv0;
     uniform sampler2D sceneTexture;
     uniform vec2 sceneTextureInvRes;
@@ -422,23 +425,24 @@ class RenderPassCompose extends RenderPassShaderQuad {
             if (this._key !== key) {
                 this._key = key;
 
-                const defines =
-                    (this.bloomTexture ? '#define BLOOM\n' : '') +
-                    (this.ssaoTexture ? '#define SSAO\n' : '') +
-                    (this.gradingEnabled ? '#define GRADING\n' : '') +
-                    (this.vignetteEnabled ? '#define VIGNETTE\n' : '') +
-                    (this.fringingEnabled ? '#define FRINGING\n' : '') +
-                    (this.taaEnabled ? '#define TAA\n' : '') +
-                    (this.isSharpnessEnabled ? '#define CAS\n' : '') +
-                    (this._srgb ? '' : '#define GAMMA_CORRECT_OUTPUT\n') +
-                    (this._debug ? `#define DEBUG_COMPOSE ${this._debug}\n` : '');
-
-                const fsChunks =
-                shaderChunks.decodePS +
-                shaderChunks.gamma2_2PS +
-                ShaderGenerator.tonemapCode(this.toneMapping);
-
-                this.shader = this.createQuadShader(`ComposeShader-${key}`, defines + fsChunks + fragmentShader);
+                const defines = new Map();
+                defines.set('TONEMAP', tonemapNames[this.toneMapping]);
+                if (this.bloomTexture) defines.set('BLOOM', true);
+                if (this.ssaoTexture) defines.set('SSAO', true);
+                if (this.gradingEnabled) defines.set('GRADING', true);
+                if (this.vignetteEnabled) defines.set('VIGNETTE', true);
+                if (this.fringingEnabled) defines.set('FRINGING', true);
+                if (this.taaEnabled) defines.set('TAA', true);
+                if (this.isSharpnessEnabled) defines.set('CAS', true);
+                if (!this._srgb) defines.set('GAMMA_CORRECT_OUTPUT', true);
+                if (this._debug) defines.set('DEBUG_COMPOSE', this._debug);
+
+                const includes = new Map(Object.entries(shaderChunks));
+
+                this.shader = this.createQuadShader(`ComposeShader-${key}`, fragmentShader, {
+                    fragmentIncludes: includes,
+                    fragmentDefines: defines
+                });
             }
         }
     }

src/scene/constants.js

@@ -660,6 +660,17 @@ export const TONEMAP_NEUTRAL = 5;
  */
 export const TONEMAP_NONE = 6;
 
+// names of the tonemaps
+export const tonemapNames = [
+    'LINEAR',
+    'FILMIC',
+    'HEJL',
+    'ACES',
+    'ACES2',
+    'NEUTRAL',
+    'NONE'
+];
+
 /**
  * No specular occlusion.
  *

src/scene/gsplat/gsplat-material.js

@@ -1,6 +1,6 @@
 import { CULLFACE_NONE, SEMANTIC_ATTR13, SEMANTIC_POSITION } from '../../platform/graphics/constants.js';
 import { ShaderProcessorOptions } from '../../platform/graphics/shader-processor-options.js';
-import { BLEND_NONE, BLEND_NORMAL, DITHER_NONE, GAMMA_NONE, GAMMA_SRGB, TONEMAP_ACES, TONEMAP_ACES2, TONEMAP_FILMIC, TONEMAP_HEJL, TONEMAP_LINEAR, TONEMAP_NEUTRAL, TONEMAP_NONE } from '../constants.js';
+import { BLEND_NONE, BLEND_NORMAL, DITHER_NONE, GAMMA_NONE, GAMMA_SRGB, tonemapNames } from '../constants.js';
 import { ShaderMaterial } from '../materials/shader-material.js';
 import { getProgramLibrary } from '../shader-lib/get-program-library.js';
 import { getMaterialShaderDefines } from '../shader-lib/utils.js';
@@ -10,16 +10,6 @@ import { ShaderGenerator } from '../shader-lib/programs/shader-generator.js';
 import { ShaderPass } from '../shader-pass.js';
 import { hashCode } from '../../core/hash.js';
 
-const tonemapNames = {
-    [TONEMAP_FILMIC]: 'FILMIC',
-    [TONEMAP_LINEAR]: 'LINEAR',
-    [TONEMAP_HEJL]: 'HEJL',
-    [TONEMAP_ACES]: 'ACES',
-    [TONEMAP_ACES2]: 'ACES2',
-    [TONEMAP_NEUTRAL]: 'NEUTRAL',
-    [TONEMAP_NONE]: 'NONE'
-};
-
 const gammaNames = {
     [GAMMA_NONE]: 'NONE',
     [GAMMA_SRGB]: 'SRGB'

src/scene/shader-lib/chunks/chunks.js

@@ -181,13 +181,14 @@ import TBNPS from './lit/frag/TBN.js';
 import TBNderivativePS from './lit/frag/TBNderivative.js';
 import TBNObjectSpacePS from './lit/frag/TBNObjectSpace.js';
 import thicknessPS from './standard/frag/thickness.js';
-import tonemappingAcesPS from './common/frag/tonemappingAces.js';
-import tonemappingAces2PS from './common/frag/tonemappingAces2.js';
-import tonemappingFilmicPS from './common/frag/tonemappingFilmic.js';
-import tonemappingHejlPS from './common/frag/tonemappingHejl.js';
-import tonemappingLinearPS from './common/frag/tonemappingLinear.js';
-import tonemappingNeutralPS from './common/frag/tonemappingNeutral.js';
-import tonemappingNonePS from './common/frag/tonemappingNone.js';
+import tonemappingPS from './common/frag/tonemapping/tonemapping.js';
+import tonemappingAcesPS from './common/frag/tonemapping/tonemappingAces.js';
+import tonemappingAces2PS from './common/frag/tonemapping/tonemappingAces2.js';
+import tonemappingFilmicPS from './common/frag/tonemapping/tonemappingFilmic.js';
+import tonemappingHejlPS from './common/frag/tonemapping/tonemappingHejl.js';
+import tonemappingLinearPS from './common/frag/tonemapping/tonemappingLinear.js';
+import tonemappingNeutralPS from './common/frag/tonemapping/tonemappingNeutral.js';
+import tonemappingNonePS from './common/frag/tonemapping/tonemappingNone.js';
 import transformVS from './common/vert/transform.js';
 import transformCoreVS from './common/vert/transformCore.js';
 import transformInstancingVS from './common/vert/transformInstancing.js';
@@ -390,6 +391,7 @@ const shaderChunks = {
     TBNderivativePS,
     TBNObjectSpacePS,
     thicknessPS,
+    tonemappingPS,
     tonemappingAcesPS,
     tonemappingAces2PS,
     tonemappingFilmicPS,

src/scene/shader-lib/chunks/common/frag/tonemapping/tonemapping.js

@@ -0,0 +1,17 @@
+export default /* glsl */`
+
+#if TONEMAP == FILMIC
+    #include "tonemappingFilmicPS"
+#elif TONEMAP == LINEAR
+    #include "tonemappingLinearPS"
+#elif TONEMAP == HEJL
+    #include "tonemappingHejlPS"
+#elif TONEMAP == ACES
+    #include "tonemappingAcesPS"
+#elif TONEMAP == ACES2
+    #include "tonemappingAces2PS"
+#elif TONEMAP == NEUTRAL
+    #include "tonemappingNeutralPS"
+#endif
+
+`;

src/scene/shader-lib/chunks/gsplat/gsplatOutput.js

@@ -1,18 +1,6 @@
 export default /* glsl */`
 
-#if TONEMAP == FILMIC
-    #include "tonemappingFilmicPS"
-#elif TONEMAP == LINEAR
-    #include "tonemappingLinearPS"
-#elif TONEMAP == HEJL
-    #include "tonemappingHejlPS"
-#elif TONEMAP == ACES
-    #include "tonemappingAcesPS"
-#elif TONEMAP == ACES2
-    #include "tonemappingAces2PS"
-#elif TONEMAP == NEUTRAL
-    #include "tonemappingNeutralPS"
-#endif
+#include "tonemappingPS"
 
 #if TONEMAP != NONE
     #if GAMMA == SRGB

src/scene/shader-lib/chunks/skybox/frag/skybox.js

@@ -1,7 +1,7 @@
 export default /* glsl */`
     #include "decodePS"
     #include "gamma"
-    #include "tonemapping"
+    #include "tonemappingPS"
     #include "envMultiplyPS"
 
     varying vec3 vViewDir;

src/scene/shader-lib/programs/particle.js

@@ -1,5 +1,5 @@
 import { ShaderUtils } from '../../../platform/graphics/shader-utils.js';
-import { BLEND_ADDITIVE, BLEND_MULTIPLICATIVE, BLEND_NORMAL } from '../../constants.js';
+import { BLEND_ADDITIVE, BLEND_MULTIPLICATIVE, BLEND_NORMAL, tonemapNames } from '../../constants.js';
 import { shaderChunks } from '../chunks/chunks.js';
 import { ShaderGenerator } from './shader-generator.js';
 
@@ -85,7 +85,7 @@ class ShaderGeneratorParticle extends ShaderGenerator {
 
         fshader += shaderChunks.decodePS;
         fshader += ShaderGenerator.gammaCode(options.gamma);
-        fshader += ShaderGenerator.tonemapCode(options.toneMap);
+        fshader += '#include "tonemappingPS"\n';
         fshader += ShaderGenerator.fogCode(options.fog);
 
         if (options.normal === 2) fshader += '\nuniform sampler2D normalMap;\n';
@@ -105,11 +105,20 @@ class ShaderGeneratorParticle extends ShaderGenerator {
         }
         fshader += shaderChunks.particle_endPS;
 
+        const includes = new Map(Object.entries({
+            ...shaderChunks,
+            ...options.chunks
+        }));
+
+        const fragmentDefines = new Map(options.defines);
+        fragmentDefines.set('TONEMAP', tonemapNames[options.toneMap]);
+
         return ShaderUtils.createDefinition(device, {
             name: 'ParticleShader',
             vertexCode: vshader,
             fragmentCode: fshader,
-            fragmentDefines: options.defines,
+            fragmentDefines: fragmentDefines,
+            fragmentIncludes: includes,
             vertexDefines: options.defines
         });
     }

src/scene/shader-lib/programs/shader-generator-shader.js

@@ -1,6 +1,7 @@
 import { hashCode } from '../../../core/hash.js';
 import { SEMANTIC_ATTR15, SEMANTIC_BLENDINDICES, SEMANTIC_BLENDWEIGHT, SHADERLANGUAGE_WGSL } from '../../../platform/graphics/constants.js';
 import { ShaderUtils } from '../../../platform/graphics/shader-utils.js';
+import { tonemapNames } from '../../constants.js';
 import { ShaderPass } from '../../shader-pass.js';
 import { shaderChunks } from '../chunks/chunks.js';
 import { ShaderGenerator } from './shader-generator.js';
@@ -14,7 +15,7 @@ const fShader = `
     #include "shaderPassDefines"
     #include "decodePS"
     #include "gamma"
-    #include "tonemapping"
+    #include "tonemappingPS"
     #include "fog"
     #include "userCode"
 `;
@@ -109,16 +110,18 @@ class ShaderGeneratorShader extends ShaderGenerator {
             definitionOptions.fragmentCode = desc.fragmentCode;
 
         } else {
-            const includes = new Map();
-            const defines = new Map(options.defines);
+            const includes = new Map(Object.entries({
+                ...shaderChunks,
+                ...options.chunks
+            }));
 
             includes.set('shaderPassDefines', shaderPassInfo.shaderDefines);
-            includes.set('decodePS', shaderChunks.decodePS);
             includes.set('gamma', ShaderGenerator.gammaCode(options.gamma));
-            includes.set('tonemapping', ShaderGenerator.tonemapCode(options.toneMapping));
             includes.set('fog', ShaderGenerator.fogCode(options.fog));
             includes.set('userCode', desc.fragmentCode);
-            includes.set('pick', shaderChunks.pickPS);
+
+            const defines = new Map(options.defines);
+            defines.set('TONEMAP', tonemapNames[options.toneMapping]);
 
             definitionOptions.fragmentCode = fShader;
             definitionOptions.fragmentIncludes = includes;

src/scene/shader-lib/programs/skybox.js

@@ -4,7 +4,7 @@ import { ChunkUtils } from '../chunk-utils.js';
 
 import { ShaderUtils } from '../../../platform/graphics/shader-utils.js';
 import { ShaderGenerator } from './shader-generator.js';
-import { SKYTYPE_INFINITE } from '../../constants.js';
+import { SKYTYPE_INFINITE, tonemapNames } from '../../constants.js';
 
 class ShaderGeneratorSkybox extends ShaderGenerator {
     generateKey(options) {
@@ -17,17 +17,20 @@ class ShaderGeneratorSkybox extends ShaderGenerator {
 
         // defines
         const defines = new Map();
+        defines.set('TONEMAP', tonemapNames[options.toneMapping]);
         defines.set('SKYBOX_DECODE_FNC', ChunkUtils.decodeFunc(options.encoding));
         if (options.skymesh !== SKYTYPE_INFINITE) defines.set('SKYMESH', '');
         if (options.type === 'cubemap') {
             defines.set('SKY_CUBEMAP', '');
         }
 
         // includes
-        const includes = new Map();
+        const includes = new Map(Object.entries({
+            ...shaderChunks,
+            ...options.chunks
+        }));
         includes.set('decodePS', shaderChunks.decodePS);
         includes.set('gamma', ShaderGenerator.gammaCode(options.gamma));
-        includes.set('tonemapping', ShaderGenerator.tonemapCode(options.toneMapping));
         includes.set('envMultiplyPS', shaderChunks.envMultiplyPS);
 
         if (options.type !== 'cubemap') {