Merge pull request #13256 from Popov72/postprocess-customize

PostProcess: Add hooks to alter shader code used by post processes

packages/dev/core/src/PostProcesses/postProcess.ts

@@ -27,6 +27,28 @@ declare type Animation = import("../Animations/animation").Animation;
 declare type PrePassRenderer = import("../Rendering/prePassRenderer").PrePassRenderer;
 declare type PrePassEffectConfiguration = import("../Rendering/prePassEffectConfiguration").PrePassEffectConfiguration;
 
+/**
+ * Allows for custom processing of the shader code used by a post process
+ */
+export type PostProcessCustomShaderCodeProcessing = {
+    /**
+     * If provided, will be called two times with the vertex and fragment code so that this code can be updated after the #include have been processed
+     */
+    processCodeAfterIncludes?: (postProcessName: string, shaderType: string, code: string) => string;
+    /**
+     * If provided, will be called two times with the vertex and fragment code so that this code can be updated before it is compiled by the GPU
+     */
+    processFinalCode?: (postProcessName: string, shaderType: string, code: string) => string;
+    /**
+     * If provided, will be called before creating the effect to collect additional custom bindings (defines, uniforms, samplers)
+     */
+    defineCustomBindings?: (postProcessName: string, defines: Nullable<string>, uniforms: string[], samplers: string[]) => Nullable<string>;
+    /**
+     * If provided, will be called when binding inputs to the shader code to allow the user to add custom bindings
+     */
+    bindCustomBindings?: (postProcessName: string, effect: Effect) => void;
+};
+
 /**
  * Size options for a post process
  */
@@ -42,6 +64,27 @@ export class PostProcess {
     /** @internal */
     public _parentContainer: Nullable<AbstractScene> = null;
 
+    private static _CustomShaderCodeProcessing: { [postProcessName: string]: PostProcessCustomShaderCodeProcessing } = {};
+
+    /**
+     * Registers a shader code processing with a post process name.
+     * @param postProcessName name of the post process. Use null for the fallback shader code processing. This is the shader code processing that will be used in case no specific shader code processing has been associated to a post process name
+     * @param customShaderCodeProcessing shader code processing to associate to the post process name
+     * @returns
+     */
+    public static RegisterShaderCodeProcessing(postProcessName: Nullable<string>, customShaderCodeProcessing?: PostProcessCustomShaderCodeProcessing) {
+        if (!customShaderCodeProcessing) {
+            delete PostProcess._CustomShaderCodeProcessing[postProcessName ?? ""];
+            return;
+        }
+
+        PostProcess._CustomShaderCodeProcessing[postProcessName ?? ""] = customShaderCodeProcessing;
+    }
+
+    private static _GetShaderCodeProcessing(postProcessName: string) {
+        return PostProcess._CustomShaderCodeProcessing[postProcessName] ?? PostProcess._CustomShaderCodeProcessing[""];
+    }
+
     /**
      * Gets or sets the unique id of the post process
      */
@@ -209,6 +252,7 @@ export class PostProcess {
     protected _indexParameters: any;
     private _shareOutputWithPostProcess: Nullable<PostProcess>;
     private _texelSize = Vector2.Zero();
+
     /** @internal */
     public _forcedOutputTexture: Nullable<RenderTargetWrapper>;
 
@@ -498,17 +542,39 @@ export class PostProcess {
         vertexUrl?: string,
         fragmentUrl?: string
     ) {
+        const customShaderCodeProcessing = PostProcess._GetShaderCodeProcessing(this.name);
+        if (customShaderCodeProcessing?.defineCustomBindings) {
+            const newUniforms = uniforms?.slice() ?? [];
+            newUniforms.push(...this._parameters);
+
+            const newSamplers = samplers?.slice() ?? [];
+            newSamplers.push(...this._samplers);
+
+            defines = customShaderCodeProcessing.defineCustomBindings(this.name, defines, newUniforms, newSamplers);
+            uniforms = newUniforms;
+            samplers = newSamplers;
+        }
         this._postProcessDefines = defines;
         this._drawWrapper.effect = this._engine.createEffect(
             { vertex: vertexUrl ?? this._vertexUrl, fragment: fragmentUrl ?? this._fragmentUrl },
-            ["position"],
-            uniforms || this._parameters,
-            samplers || this._samplers,
-            defines !== null ? defines : "",
-            undefined,
-            onCompiled,
-            onError,
-            indexParameters || this._indexParameters
+            {
+                attributes: ["position"],
+                uniformsNames: uniforms || this._parameters,
+                uniformBuffersNames: [],
+                samplers: samplers || this._samplers,
+                defines: defines !== null ? defines : "",
+                fallbacks: null,
+                onCompiled: onCompiled ?? null,
+                onError: onError ?? null,
+                indexParameters: indexParameters || this._indexParameters,
+                processCodeAfterIncludes: customShaderCodeProcessing?.processCodeAfterIncludes
+                    ? (shaderType: string, code: string) => customShaderCodeProcessing!.processCodeAfterIncludes!(this.name, shaderType, code)
+                    : null,
+                processFinalCode: customShaderCodeProcessing?.processFinalCode
+                    ? (shaderType: string, code: string) => customShaderCodeProcessing!.processFinalCode!(this.name, shaderType, code)
+                    : null,
+            },
+            this._engine
         );
     }
 
@@ -787,6 +853,8 @@ export class PostProcess {
         this._drawWrapper.effect.setVector2("scale", this._scaleRatio);
         this.onApplyObservable.notifyObservers(this._drawWrapper.effect);
 
+        PostProcess._GetShaderCodeProcessing(this.name)?.bindCustomBindings?.(this.name, this._drawWrapper.effect);
+
         return this._drawWrapper.effect;
     }
 

packages/dev/core/src/Shaders/ShadersInclude/imageProcessingFunctions.fx

@@ -93,8 +93,12 @@
 	}
 #endif
 
+#define CUSTOM_IMAGEPROCESSINGFUNCTIONS_DEFINITIONS
+
 vec4 applyImageProcessing(vec4 result) {
 
+	#define CUSTOM_IMAGEPROCESSINGFUNCTIONS_UPDATERESULT_ATSTART
+
 #ifdef EXPOSURE
 	result.rgb *= exposureLinear;
 #endif
@@ -174,5 +178,7 @@ vec4 applyImageProcessing(vec4 result) {
 	result.rgb = saturate(result.rgb + vec3(dither));
 #endif
 
+	#define CUSTOM_IMAGEPROCESSINGFUNCTIONS_UPDATERESULT_ATEND
+
 	return result;
 }

packages/dev/core/src/Shaders/circleOfConfusion.fragment.fx

@@ -18,7 +18,13 @@ uniform float cocPrecalculation;
 void main(void)
 {
     float depth = texture2D(depthSampler, vUV).r;
+
+    #define CUSTOM_COC_DEPTH
+
     float pixelDistance = (cameraMinMaxZ.x + cameraMinMaxZ.y * depth) * 1000.0; // actual distance from the lens in scene units/1000 (eg. millimeter)
+
+    #define CUSTOM_COC_PIXELDISTANCE
+
     float coc = abs(cocPrecalculation * ((focusDistance - pixelDistance) / pixelDistance));
     coc = clamp(coc, 0.0, 1.0);
     gl_FragColor = vec4(coc, coc, coc, 1.0);