Introducing loop support for NME

packages/dev/core/src/Materials/Node/Blocks/Dual/imageSourceBlock.ts

@@ -13,6 +13,7 @@ import { NodeMaterial } from "../../nodeMaterial";
 import type { Scene } from "../../../../scene";
 import { NodeMaterialConnectionPointCustomObject } from "../../nodeMaterialConnectionPointCustomObject";
 import { EngineStore } from "../../../../Engines/engineStore";
+import { ShaderLanguage } from "core/Materials/shaderLanguage";
 /**
  * Block used to provide an image for a TextureBlock
  */
@@ -68,6 +69,8 @@ export class ImageSourceBlock extends NodeMaterialBlock {
             NodeMaterialBlockTargets.VertexAndFragment,
             new NodeMaterialConnectionPointCustomObject("source", this, NodeMaterialConnectionPointDirection.Output, ImageSourceBlock, "ImageSourceBlock")
         );
+
+        this.registerOutput("dimensions", NodeMaterialBlockConnectionPointTypes.Vector2);
     }
 
     public override bind(effect: Effect) {
@@ -101,6 +104,13 @@ export class ImageSourceBlock extends NodeMaterialBlock {
         return this._outputs[0];
     }
 
+    /**
+     * Gets the dimension component
+     */
+    public get dimensions(): NodeMaterialConnectionPoint {
+        return this._outputs[1];
+    }
+
     protected override _buildBlock(state: NodeMaterialBuildState) {
         super._buildBlock(state);
 
@@ -113,6 +123,17 @@ export class ImageSourceBlock extends NodeMaterialBlock {
             state.sharedData.bindableBlocks.push(this);
         }
 
+        if (this.dimensions.isConnected) {
+            let affect: string = "";
+            if (state.shaderLanguage === ShaderLanguage.WGSL) {
+                affect = `vec2f(textureDimensions(${this._samplerName}, 0).xy)`;
+            } else {
+                affect = `vec2(textureSize(${this._samplerName}, 0).xy)`;
+            }
+
+            state.compilationString += `${state._declareOutput(this.dimensions)} = ${affect};\n`;
+        }
+
         state._emit2DSampler(this._samplerName);
 
         return this;

packages/dev/core/src/Materials/Node/Blocks/Dual/textureBlock.ts

@@ -271,7 +271,37 @@ export class TextureBlock extends NodeMaterialBlock {
         return this._outputs[6];
     }
 
-    public override get target() {
+    private _isTiedToFragment(input: NodeMaterialConnectionPoint) {
+        if (input.target === NodeMaterialBlockTargets.Fragment) {
+            return true;
+        }
+
+        if (input.target === NodeMaterialBlockTargets.Vertex) {
+            return false;
+        }
+
+        if (input.target === NodeMaterialBlockTargets.Neutral || input.target === NodeMaterialBlockTargets.VertexAndFragment) {
+            const parentBlock = input.ownerBlock;
+
+            if (parentBlock.target === NodeMaterialBlockTargets.Fragment) {
+                return true;
+            }
+
+            for (const input of parentBlock.inputs) {
+                if (!input.isConnected) {
+                    continue;
+                }
+                if (this._isTiedToFragment(input.connectedPoint!)) {
+                    return true;
+                }
+            }
+        }
+
+        return false;
+    }
+
+    private _cachedTarget: Nullable<NodeMaterialBlockTargets> = null;
+    private _getEffectiveTarget() {
         if (this._fragmentOnly) {
             return NodeMaterialBlockTargets.Fragment;
         }
@@ -286,35 +316,19 @@ export class TextureBlock extends NodeMaterialBlock {
             return NodeMaterialBlockTargets.VertexAndFragment;
         }
 
-        let parent = this.uv.connectedPoint;
-
-        while (parent) {
-            if (parent.target === NodeMaterialBlockTargets.Fragment) {
-                return NodeMaterialBlockTargets.Fragment;
-            }
-
-            if (parent.target === NodeMaterialBlockTargets.Vertex) {
-                return NodeMaterialBlockTargets.VertexAndFragment;
-            }
-
-            if (parent.target === NodeMaterialBlockTargets.Neutral || parent.target === NodeMaterialBlockTargets.VertexAndFragment) {
-                const parentBlock = parent.ownerBlock;
+        if (this._isTiedToFragment(this.uv.connectedPoint!)) {
+            return NodeMaterialBlockTargets.Fragment;
+        }
 
-                if (parentBlock.target === NodeMaterialBlockTargets.Fragment) {
-                    return NodeMaterialBlockTargets.Fragment;
-                }
+        return NodeMaterialBlockTargets.VertexAndFragment;
+    }
 
-                parent = null;
-                for (const input of parentBlock.inputs) {
-                    if (input.connectedPoint) {
-                        parent = input.connectedPoint;
-                        break;
-                    }
-                }
-            }
+    public override get target() {
+        if (!this._cachedTarget) {
+            this._cachedTarget = this._getEffectiveTarget();
         }
 
-        return NodeMaterialBlockTargets.VertexAndFragment;
+        return this._cachedTarget;
     }
 
     public override set target(value: NodeMaterialBlockTargets) {}
@@ -576,6 +590,7 @@ export class TextureBlock extends NodeMaterialBlock {
     }
 
     protected override _buildBlock(state: NodeMaterialBuildState) {
+        this._cachedTarget = null;
         super._buildBlock(state);
 
         if (this.source.isConnected) {

packages/dev/core/src/Materials/Node/Blocks/index.ts

@@ -63,3 +63,6 @@ export * from "./matrixTransposeBlock";
 export * from "./meshAttributeExistsBlock";
 export * from "./curveBlock";
 export * from "./colorConverterBlock";
+export * from "./loopBlock";
+export * from "./storageReadBlock";
+export * from "./storageWriteBlock";

packages/dev/core/src/Materials/Node/Blocks/loopBlock.ts

@@ -0,0 +1,143 @@
+import { NodeMaterialBlock } from "../nodeMaterialBlock";
+import { NodeMaterialBlockConnectionPointTypes } from "../Enums/nodeMaterialBlockConnectionPointTypes";
+import type { NodeMaterialBuildState } from "../nodeMaterialBuildState";
+import { NodeMaterialConnectionPointDirection, type NodeMaterialConnectionPoint } from "../nodeMaterialBlockConnectionPoint";
+import { NodeMaterialBlockTargets } from "../Enums/nodeMaterialBlockTargets";
+import { RegisterClass } from "../../../Misc/typeStore";
+import { editableInPropertyPage, PropertyTypeForEdition } from "core/Decorators/nodeDecorator";
+import type { Scene } from "core/scene";
+import { ShaderLanguage } from "core/Materials/shaderLanguage";
+import { NodeMaterialConnectionPointCustomObject } from "../nodeMaterialConnectionPointCustomObject";
+/**
+ * Block used to repeat code
+ */
+export class LoopBlock extends NodeMaterialBlock {
+    /**
+     * Gets or sets the source range
+     */
+    @editableInPropertyPage("Iterations", PropertyTypeForEdition.Int)
+    public iterations = 4;
+
+    /**
+     * Creates a new LoopBlock
+     * @param name defines the block name
+     */
+    public constructor(name: string) {
+        super(name, NodeMaterialBlockTargets.Neutral);
+
+        this.registerInput("input", NodeMaterialBlockConnectionPointTypes.AutoDetect);
+        this.registerInput("iterations", NodeMaterialBlockConnectionPointTypes.Float, true);
+        this.registerOutput("output", NodeMaterialBlockConnectionPointTypes.BasedOnInput);
+        this.registerOutput("index", NodeMaterialBlockConnectionPointTypes.Float, NodeMaterialBlockTargets.Fragment);
+        this.registerOutput(
+            "loopID",
+            NodeMaterialBlockConnectionPointTypes.Object,
+            undefined,
+            new NodeMaterialConnectionPointCustomObject("loopID", this, NodeMaterialConnectionPointDirection.Output, LoopBlock, "LoopBlock")
+        );
+
+        this._outputs[0]._typeConnectionSource = this._inputs[0];
+        this._outputs[0]._forPostBuild = true;
+
+        this._outputs[2]._redirectedSource = this._inputs[0];
+
+        this._outputs[1]._preventBubbleUp = true;
+        this._outputs[2]._preventBubbleUp = true;
+    }
+
+    /**
+     * Gets the current class name
+     * @returns the class name
+     */
+    public override getClassName() {
+        return "LoopBlock";
+    }
+
+    /**
+     * Gets the main input component
+     */
+    public get input(): NodeMaterialConnectionPoint {
+        return this._inputs[0];
+    }
+
+    /**
+     * Gets the iterations input component
+     */
+    public get iterationsInput(): NodeMaterialConnectionPoint {
+        return this._inputs[1];
+    }
+
+    /**
+     * Gets the output component
+     */
+    public get output(): NodeMaterialConnectionPoint {
+        return this._outputs[0];
+    }
+
+    /**
+     * Gets the index component which will be incremented at each iteration
+     */
+    public get index(): NodeMaterialConnectionPoint {
+        return this._outputs[1];
+    }
+
+    /**
+     * Gets the loop ID component
+     */
+    public get loopID(): NodeMaterialConnectionPoint {
+        return this._outputs[2];
+    }
+
+    protected override _buildBlock(state: NodeMaterialBuildState) {
+        super._buildBlock(state);
+
+        const output = this._outputs[0];
+        const index = this._outputs[1];
+
+        const indexName = state._getFreeVariableName("index");
+
+        const decl = state.shaderLanguage === ShaderLanguage.WGSL ? "var" : "int";
+        const castFloat = state.shaderLanguage === ShaderLanguage.WGSL ? "f32" : "float";
+        const castInt = state.shaderLanguage === ShaderLanguage.WGSL ? "u32" : "int";
+
+        // Declare storage variable and store initial value
+        state.compilationString += state._declareOutput(output) + ` = ${this.input.associatedVariableName};\n`;
+
+        // Iterations
+        const iterations = this.iterationsInput.isConnected ? `${castInt}(${this.iterationsInput.associatedVariableName})` : this.iterations;
+
+        // Loop
+        state.compilationString += `for (${decl} ${indexName} = 0; ${indexName} < ${iterations}; ${indexName}++){\n`;
+        state.compilationString += `${state._declareOutput(index)} = ${castFloat}(${indexName});\n`;
+
+        return this;
+    }
+
+    protected override _postBuildBlock(state: NodeMaterialBuildState) {
+        super._postBuildBlock(state);
+
+        state.compilationString += `}\n`;
+
+        return this;
+    }
+
+    protected override _dumpPropertiesCode() {
+        return super._dumpPropertiesCode() + `${this._codeVariableName}.iterations = ${this.iterations};\n`;
+    }
+
+    public override serialize(): any {
+        const serializationObject = super.serialize();
+
+        serializationObject.iterations = this.iterations;
+
+        return serializationObject;
+    }
+
+    public override _deserialize(serializationObject: any, scene: Scene, rootUrl: string) {
+        super._deserialize(serializationObject, scene, rootUrl);
+
+        this.iterations = serializationObject.iterations;
+    }
+}
+
+RegisterClass("BABYLON.LoopBlock", LoopBlock);

packages/dev/core/src/Materials/Node/Blocks/storageReadBlock.ts

@@ -0,0 +1,71 @@
+import { NodeMaterialBlock } from "../nodeMaterialBlock";
+import { NodeMaterialBlockConnectionPointTypes } from "../Enums/nodeMaterialBlockConnectionPointTypes";
+import type { NodeMaterialBuildState } from "../nodeMaterialBuildState";
+import { NodeMaterialConnectionPointDirection, type NodeMaterialConnectionPoint } from "../nodeMaterialBlockConnectionPoint";
+import { NodeMaterialBlockTargets } from "../Enums/nodeMaterialBlockTargets";
+import { RegisterClass } from "../../../Misc/typeStore";
+import { LoopBlock } from "./loopBlock";
+import { NodeMaterialConnectionPointCustomObject } from "../nodeMaterialConnectionPointCustomObject";
+/**
+ * Block used to read from a variable within a loop
+ */
+export class StorageReadBlock extends NodeMaterialBlock {
+    /**
+     * Creates a new StorageReadBlock
+     * @param name defines the block name
+     */
+    public constructor(name: string) {
+        super(name, NodeMaterialBlockTargets.Neutral);
+
+        this.registerInput(
+            "loopID",
+            NodeMaterialBlockConnectionPointTypes.Object,
+            false,
+            undefined,
+            new NodeMaterialConnectionPointCustomObject("loopID", this, NodeMaterialConnectionPointDirection.Input, LoopBlock, "LoopBlock")
+        );
+        this.registerOutput("value", NodeMaterialBlockConnectionPointTypes.AutoDetect);
+
+        this._outputs[0]._linkedConnectionSource = this._inputs[0];
+    }
+
+    /**
+     * Gets the current class name
+     * @returns the class name
+     */
+    public override getClassName() {
+        return "StorageReadBlock";
+    }
+
+    /**
+     * Gets the loop link component
+     */
+    public get loopID(): NodeMaterialConnectionPoint {
+        return this._inputs[0];
+    }
+
+    /**
+     * Gets the value component
+     */
+    public get value(): NodeMaterialConnectionPoint {
+        return this._outputs[0];
+    }
+
+    protected override _buildBlock(state: NodeMaterialBuildState) {
+        super._buildBlock(state);
+
+        const value = this._outputs[0];
+
+        if (!this.loopID.isConnected) {
+            return this;
+        }
+
+        const loopBlock = this.loopID.connectedPoint!.ownerBlock as LoopBlock;
+
+        state.compilationString += state._declareOutput(value) + ` = ${loopBlock.output.associatedVariableName};\n`;
+
+        return this;
+    }
+}
+
+RegisterClass("BABYLON.StorageReadBlock", StorageReadBlock);