Add type to shader material constructor (#14908)

* Add type to shader material constructor

* Put everything into one object

* Fix one CI failure

* Move IShaderPath

* Fixes

* Make vertexToken have a higher priority

* Tweaks

* Add safety to compute shader as well

packages/dev/core/src/Compute/computeEffect.ts

@@ -10,6 +10,31 @@ import { ShaderLanguage } from "../Materials/shaderLanguage";
 
 import type { Engine } from "../Engines/engine";
 
+/**
+ * Defines the route to the shader code. The priority is as follows:
+ *  * object: `{ computeSource: "compute shader code string"}` for directly passing the shader code
+ *  * object: `{ computeElement: "vertexShaderCode" }`, used with shader code in script tags
+ *  * object: `{ compute: "custom" }`, used with `Effect.ShadersStore["customVertexShader"]` and `Effect.ShadersStore["customFragmentShader"]`
+ *  * string: `"./COMMON_NAME"`, used with external files COMMON_NAME.vertex.fx and COMMON_NAME.fragment.fx in index.html folder.
+ */
+export type IComputeShaderPath =
+    | {
+          /**
+           * Directly pass the shader code
+           */
+          computeSource?: string;
+          /**
+           * Used with Effect.ShadersStore. If the `vertex` is set to `"custom`, then
+           * Babylon.js will read from Effect.ShadersStore["customVertexShader"]
+           */
+          compute?: string;
+          /**
+           * Used with shader code in script tags
+           */
+          computeElement?: string;
+      }
+    | string;
+
 /**
  * Options to be used when creating a compute effect.
  */
@@ -49,7 +74,7 @@ export class ComputeEffect {
     /**
      * Name of the effect.
      */
-    public name: any = null;
+    public name: IComputeShaderPath;
     /**
      * String container all the define statements that should be set on the shader.
      */
@@ -110,7 +135,7 @@ export class ComputeEffect {
      * @param engine The engine the effect is created for
      * @param key Effect Key identifying uniquely compiled shader variants
      */
-    constructor(baseName: any, options: IComputeEffectCreationOptions, engine: Engine, key = "") {
+    constructor(baseName: IComputeShaderPath, options: IComputeEffectCreationOptions, engine: Engine, key = "") {
         this.name = baseName;
         this._key = key;
 
@@ -126,18 +151,16 @@ export class ComputeEffect {
         this._shaderRepository = ShaderStore.GetShadersRepository(this._shaderLanguage);
         this._includeShaderStore = ShaderStore.GetIncludesShadersStore(this._shaderLanguage);
 
-        let computeSource: any;
+        let computeSource: IComputeShaderPath | HTMLElement;
 
         const hostDocument = IsWindowObjectExist() ? this._engine.getHostDocument() : null;
 
-        if (baseName.computeSource) {
+        if (typeof baseName === "string") {
+            computeSource = baseName;
+        } else if (baseName.computeSource) {
             computeSource = "source:" + baseName.computeSource;
         } else if (baseName.computeElement) {
-            computeSource = hostDocument ? hostDocument.getElementById(baseName.computeElement) : null;
-
-            if (!computeSource) {
-                computeSource = baseName.computeElement;
-            }
+            computeSource = hostDocument?.getElementById(baseName.computeElement) || baseName.computeElement;
         } else {
             computeSource = baseName.compute || baseName;
         }

packages/dev/core/src/Compute/computeShader.ts

@@ -5,7 +5,7 @@ import type { Nullable } from "../types";
 import { serialize } from "../Misc/decorators";
 import { SerializationHelper } from "../Misc/decorators.serialization";
 import { RegisterClass } from "../Misc/typeStore";
-import type { ComputeEffect, IComputeEffectCreationOptions } from "./computeEffect";
+import type { ComputeEffect, IComputeEffectCreationOptions, IComputeShaderPath } from "./computeEffect";
 import type { ComputeBindingMapping } from "../Engines/Extensions/engine.computeShader";
 import { ComputeBindingType } from "../Engines/Extensions/engine.computeShader";
 import type { BaseTexture } from "../Materials/Textures/baseTexture";
@@ -55,7 +55,7 @@ type ComputeBindingListInternal = { [key: string]: { type: ComputeBindingType; o
  */
 export class ComputeShader {
     private _engine: ThinEngine;
-    private _shaderPath: any;
+    private _shaderPath: IComputeShaderPath;
     private _options: IComputeShaderOptions;
     private _effect: ComputeEffect;
     private _cachedDefines: string;
@@ -116,14 +116,14 @@ export class ComputeShader {
      * Instantiates a new compute shader.
      * @param name Defines the name of the compute shader in the scene
      * @param engine Defines the engine the compute shader belongs to
-     * @param shaderPath Defines  the route to the shader code in one of three ways:
+     * @param shaderPath Defines the route to the shader code in one of three ways:
      *  * object: \{ compute: "custom" \}, used with ShaderStore.ShadersStoreWGSL["customComputeShader"]
      *  * object: \{ computeElement: "HTMLElementId" \}, used with shader code in script tags
      *  * object: \{ computeSource: "compute shader code string" \}, where the string contains the shader code
      *  * string: try first to find the code in ShaderStore.ShadersStoreWGSL[shaderPath + "ComputeShader"]. If not, assumes it is a file with name shaderPath.compute.fx in index.html folder.
      * @param options Define the options used to create the shader
      */
-    constructor(name: string, engine: ThinEngine, shaderPath: any, options: Partial<IComputeShaderOptions> = {}) {
+    constructor(name: string, engine: ThinEngine, shaderPath: IComputeShaderPath, options: Partial<IComputeShaderOptions> = {}) {
         this.name = name;
         this._engine = engine;
         this.uniqueId = UniqueIdGenerator.UniqueId;

packages/dev/core/src/Engines/Extensions/engine.computeShader.ts

@@ -1,5 +1,5 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
-import type { ComputeEffect, IComputeEffectCreationOptions } from "../../Compute/computeEffect";
+import type { ComputeEffect, IComputeEffectCreationOptions, IComputeShaderPath } from "../../Compute/computeEffect";
 import type { IComputeContext } from "../../Compute/IComputeContext";
 import type { IComputePipelineContext } from "../../Compute/IComputePipelineContext";
 import { ThinEngine } from "../../Engines/thinEngine";
@@ -41,7 +41,15 @@ declare module "../../Engines/thinEngine" {
          * @param options Options used to create the effect
          * @returns The new compute effect
          */
-        createComputeEffect(baseName: any, options: IComputeEffectCreationOptions): ComputeEffect;
+        createComputeEffect(
+            baseName: IComputeShaderPath & {
+                /**
+                 * @internal
+                 */
+                computeToken?: string;
+            },
+            options: IComputeEffectCreationOptions
+        ): ComputeEffect;
 
         /**
          * Creates a new compute pipeline context
@@ -111,7 +119,7 @@ declare module "../../Engines/thinEngine" {
     }
 }
 
-ThinEngine.prototype.createComputeEffect = function (baseName: any, options: IComputeEffectCreationOptions): ComputeEffect {
+ThinEngine.prototype.createComputeEffect = function (baseName: IComputeShaderPath & { computeToken?: string }, options: IComputeEffectCreationOptions): ComputeEffect {
     throw new Error("createComputeEffect: This engine does not support compute shaders!");
 };
 

packages/dev/core/src/Engines/WebGPU/Extensions/engine.computeShader.ts

@@ -1,5 +1,5 @@
 import { Logger } from "core/Misc/logger";
-import type { IComputeEffectCreationOptions } from "../../../Compute/computeEffect";
+import type { IComputeEffectCreationOptions, IComputeShaderPath } from "../../../Compute/computeEffect";
 import { ComputeEffect } from "../../../Compute/computeEffect";
 import type { IComputeContext } from "../../../Compute/IComputeContext";
 import type { IComputePipelineContext } from "../../../Compute/IComputePipelineContext";
@@ -24,8 +24,8 @@ WebGPUEngine.prototype.createComputeContext = function (): IComputeContext | und
     return new WebGPUComputeContext(this._device, this._cacheSampler);
 };
 
-WebGPUEngine.prototype.createComputeEffect = function (baseName: any, options: IComputeEffectCreationOptions): ComputeEffect {
-    const compute = baseName.computeElement || baseName.compute || baseName.computeToken || baseName.computeSource || baseName;
+WebGPUEngine.prototype.createComputeEffect = function (baseName: IComputeShaderPath & { computeToken?: string }, options: IComputeEffectCreationOptions): ComputeEffect {
+    const compute = typeof baseName === "string" ? baseName : baseName.computeToken || baseName.computeSource || baseName.computeElement || baseName.compute;
 
     const name = compute + "@" + options.defines;
     if (this._compiledComputeEffects[name]) {

packages/dev/core/src/Engines/thinEngine.ts

@@ -1,7 +1,7 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
 import { EngineStore } from "./engineStore";
 import type { IInternalTextureLoader } from "../Materials/Textures/internalTextureLoader";
-import type { IEffectCreationOptions } from "../Materials/effect";
+import type { IEffectCreationOptions, IShaderPath } from "../Materials/effect";
 import { Effect } from "../Materials/effect";
 import { _WarnImport } from "../Misc/devTools";
 import type { IShaderProcessor } from "./Processors/iShaderProcessor";
@@ -2924,7 +2924,7 @@ export class ThinEngine {
      * @returns the new Effect
      */
     public createEffect(
-        baseName: any,
+        baseName: IShaderPath & { vertexToken?: string; fragmentToken?: string },
         attributesNamesOrOptions: string[] | IEffectCreationOptions,
         uniformsNamesOrEngine: string[] | ThinEngine,
         samplers?: string[],
@@ -2935,8 +2935,8 @@ export class ThinEngine {
         indexParameters?: any,
         shaderLanguage = ShaderLanguage.GLSL
     ): Effect {
-        const vertex = baseName.vertexElement || baseName.vertex || baseName.vertexToken || baseName.vertexSource || baseName;
-        const fragment = baseName.fragmentElement || baseName.fragment || baseName.fragmentToken || baseName.fragmentSource || baseName;
+        const vertex = typeof baseName === "string" ? baseName : baseName.vertexToken || baseName.vertexSource || baseName.vertexElement || baseName.vertex;
+        const fragment = typeof baseName === "string" ? baseName : baseName.fragmentToken || baseName.fragmentSource || baseName.fragmentElement || baseName.fragment;
         const globalDefines = this._getGlobalDefines()!;
 
         let fullDefines = defines ?? (<IEffectCreationOptions>attributesNamesOrOptions).defines ?? "";

packages/dev/core/src/Engines/webgpuEngine.ts

@@ -4,7 +4,7 @@ import type { Nullable, DataArray, IndicesArray, Immutable, FloatArray } from ".
 import { Color4 } from "../Maths/math";
 import { Engine } from "../Engines/engine";
 import { InternalTexture, InternalTextureSource } from "../Materials/Textures/internalTexture";
-import type { IEffectCreationOptions } from "../Materials/effect";
+import type { IEffectCreationOptions, IShaderPath } from "../Materials/effect";
 import { Effect } from "../Materials/effect";
 import type { EffectFallbacks } from "../Materials/effectFallbacks";
 import { Constants } from "./constants";
@@ -1766,7 +1766,7 @@ export class WebGPUEngine extends Engine {
      * @returns the new Effect
      */
     public createEffect(
-        baseName: any,
+        baseName: IShaderPath & { vertexToken?: string; fragmentToken?: string },
         attributesNamesOrOptions: string[] | IEffectCreationOptions,
         uniformsNamesOrEngine: string[] | Engine,
         samplers?: string[],
@@ -1777,8 +1777,8 @@ export class WebGPUEngine extends Engine {
         indexParameters?: any,
         shaderLanguage = ShaderLanguage.GLSL
     ): Effect {
-        const vertex = baseName.vertexElement || baseName.vertex || baseName.vertexToken || baseName.vertexSource || baseName;
-        const fragment = baseName.fragmentElement || baseName.fragment || baseName.fragmentToken || baseName.fragmentSource || baseName;
+        const vertex = typeof baseName === "string" ? baseName : baseName.vertexToken || baseName.vertexSource || baseName.vertexElement || baseName.vertex;
+        const fragment = typeof baseName === "string" ? baseName : baseName.fragmentToken || baseName.fragmentSource || baseName.fragmentElement || baseName.fragment;
         const globalDefines = this._getGlobalDefines()!;
 
         let fullDefines = defines ?? (<IEffectCreationOptions>attributesNamesOrOptions).defines ?? "";
@@ -2025,7 +2025,7 @@ export class WebGPUEngine extends Engine {
             this._counters.numEnableEffects++;
             if (this.dbgLogIfNotDrawWrapper) {
                 Logger.Warn(
-                    `enableEffect has been called with an Effect and not a Wrapper! effect.uniqueId=${effect.uniqueId}, effect.name=${effect.name}, effect.name.vertex=${effect.name.vertex}, effect.name.fragment=${effect.name.fragment}`,
+                    `enableEffect has been called with an Effect and not a Wrapper! effect.uniqueId=${effect.uniqueId}, effect.name=${effect.name}, effect.name.vertex=${typeof effect.name === "string" ? "" : effect.name.vertex}, effect.name.fragment=${typeof effect.name === "string" ? "" : effect.name.fragment}`,
                     10
                 );
             }

packages/dev/core/src/Materials/effect.ts

@@ -21,6 +21,44 @@ import type { ThinTexture } from "../Materials/Textures/thinTexture";
 import type { RenderTargetTexture } from "../Materials/Textures/renderTargetTexture";
 import type { PostProcess } from "../PostProcesses/postProcess";
 
+/**
+ * Defines the route to the shader code. The priority is as follows:
+ *  * object: `{ vertexSource: "vertex shader code string", fragmentSource: "fragment shader code string" }` for directly passing the shader code
+ *  * object: `{ vertexElement: "vertexShaderCode", fragmentElement: "fragmentShaderCode" }`, used with shader code in script tags
+ *  * object: `{ vertex: "custom", fragment: "custom" }`, used with `Effect.ShadersStore["customVertexShader"]` and `Effect.ShadersStore["customFragmentShader"]`
+ *  * string: `"./COMMON_NAME"`, used with external files COMMON_NAME.vertex.fx and COMMON_NAME.fragment.fx in index.html folder.
+ */
+export type IShaderPath =
+    | {
+          /**
+           * Directly pass the shader code
+           */
+          vertexSource?: string;
+          /**
+           * Directly pass the shader code
+           */
+          fragmentSource?: string;
+          /**
+           * Used with Effect.ShadersStore. If the `vertex` is set to `"custom`, then
+           * Babylon.js will read from Effect.ShadersStore["customVertexShader"]
+           */
+          vertex?: string;
+          /**
+           * Used with Effect.ShadersStore. If the `fragment` is set to `"custom`, then
+           * Babylon.js will read from Effect.ShadersStore["customFragmentShader"]
+           */
+          fragment?: string;
+          /**
+           * Used with shader code in script tags
+           */
+          vertexElement?: string;
+          /**
+           * Used with shader code in script tags
+           */
+          fragmentElement?: string;
+      }
+    | string;
+
 /**
  * Options to be used when creating an effect.
  */
@@ -107,7 +145,7 @@ export class Effect implements IDisposable {
     /**
      * Name of the effect.
      */
-    public name: any = null;
+    public name: IShaderPath;
     /**
      * String container all the define statements that should be set on the shader.
      */
@@ -237,7 +275,7 @@ export class Effect implements IDisposable {
      * @param shaderLanguage the language the shader is written in (default: GLSL)
      */
     constructor(
-        baseName: any,
+        baseName: IShaderPath,
         attributesNamesOrOptions: string[] | IEffectCreationOptions,
         uniformsNamesOrEngine: string[] | ThinEngine,
         samplers: Nullable<string[]> = null,
@@ -303,32 +341,27 @@ export class Effect implements IDisposable {
 
     /** @internal */
     public _processShaderCode(shaderProcessor: Nullable<IShaderProcessor> = null, keepExistingPipelineContext = false) {
-        let vertexSource: any;
-        let fragmentSource: any;
+        let vertexSource: string | HTMLElement | IShaderPath;
+        let fragmentSource: string | HTMLElement | IShaderPath;
 
         const baseName = this.name;
         const hostDocument = IsWindowObjectExist() ? this._engine.getHostDocument() : null;
 
-        if (baseName.vertexSource) {
+        if (typeof baseName === "string") {
+            vertexSource = baseName;
+        } else if (baseName.vertexSource) {
             vertexSource = "source:" + baseName.vertexSource;
         } else if (baseName.vertexElement) {
-            vertexSource = hostDocument ? hostDocument.getElementById(baseName.vertexElement) : null;
-
-            if (!vertexSource) {
-                vertexSource = baseName.vertexElement;
-            }
+            vertexSource = hostDocument?.getElementById(baseName.vertexElement) || baseName.vertexElement;
         } else {
             vertexSource = baseName.vertex || baseName;
         }
-
-        if (baseName.fragmentSource) {
+        if (typeof baseName === "string") {
+            fragmentSource = baseName;
+        } else if (baseName.fragmentSource) {
             fragmentSource = "source:" + baseName.fragmentSource;
         } else if (baseName.fragmentElement) {
-            fragmentSource = hostDocument ? hostDocument.getElementById(baseName.fragmentElement) : null;
-
-            if (!fragmentSource) {
-                fragmentSource = baseName.fragmentElement;
-            }
+            fragmentSource = hostDocument?.getElementById(baseName.fragmentElement) || baseName.fragmentElement;
         } else {
             fragmentSource = baseName.fragment || baseName;
         }