[VFX] Missing Includes with CustomHLSL & SG

Issue introduced at https://github.cds.internal.unity3d.com/unity/unity/pull/40913

The support of `CustomHLSL` was partial and didn't consider custom HLSL and includes workflow.

The exact repro was about usage in update leaking into output generated but the issue was reproducible actually using SG in output.

➕  Additionally, the `defines` wasn't applied in SG (but it doesn't have any effect, so far, only [FillAabbBuffer](https://github.cds.internal.unity3d.com/unity/unity/blob/05bbb1aacd661908468d5fe51f2a83f96c0649f2/Packages/com.unity.visualeffectgraph/Shaders/VFXBoundsUtils.hlsl#L194) and it isn't called in any vertex or pixel shader. I introduced the define declaration for consistency.

➕ Move the `VFXCommonOutput.hlsl` in case of SG because it wasn't consistent with builtin generated code.

Author: PaulDemeulenaere

Packages/com.unity.visualeffectgraph/Editor/Compiler/VFXCodeGenerator.cs

@@ -401,12 +401,19 @@ internal static Dictionary<VFXExpression, string> BuildExpressionToName(VFXConte
 
         internal static void BuildContextBlocks(VFXContext context, VFXTaskCompiledData taskData, Dictionary<VFXExpression, string> expressionToName,
             out VFXShaderWriter blockFunction,
-            out VFXShaderWriter blockCallFunction)
+            out VFXShaderWriter blockCallFunction,
+            out VFXShaderWriter blockIncludes,
+            out VFXShaderWriter blockDefines)
         {
             //< Block processor
             blockFunction = new VFXShaderWriter();
             blockCallFunction = new VFXShaderWriter();
+            blockIncludes = new VFXShaderWriter();
+            blockDefines = new VFXShaderWriter();
+
             var blockDeclared = new HashSet<string>();
+            var includesProcessed = new HashSet<string>();
+            var defineProcessed = new HashSet<string>();
 
             int cpt = 0;
             foreach (var current in context.activeFlattenedChildrenWithImplicit)
@@ -415,14 +422,42 @@ internal static void BuildContextBlocks(VFXContext context, VFXTaskCompiledData
                 if (current is IHLSLCodeHolder hlslCodeHolder)
                 {
                     blockFunction.Write(hlslCodeHolder.customCode);
+                    foreach (var includePath in hlslCodeHolder.includes)
+                    {
+                        if (includesProcessed.Add(includePath))
+                        {
+                            blockIncludes.WriteLine($"#include \"{includePath}\"");
+                        }
+                    }
+                }
+
+                foreach (var define in current.defines)
+                {
+                    if (defineProcessed.Add(define))
+                    {
+                        blockDefines.WriteLineFormat("#define {0}{1}", define, define.Contains(' ') ? "" : " 1");
+                    }
                 }
                 BuildBlock(taskData, blockFunction, blockCallFunction, blockDeclared, expressionToName, current, ref cpt);
             }
 
             // Custom HLSL Operators
-            foreach (var group in taskData.hlslCodeHolders.GroupBy(x => x.customCode.GetHashCode()))
+            var customCodeProcessed = new HashSet<string>();
+            foreach (var hlslCodeHolder in taskData.hlslCodeHolders)
             {
-                blockFunction.Write(group.First().customCode);
+                var customCode = hlslCodeHolder.customCode;
+                if (customCodeProcessed.Add(customCode))
+                {
+                    blockFunction.Write(customCode);
+                }
+
+                foreach (var includePath in hlslCodeHolder.includes)
+                {
+                    if (includesProcessed.Add(includePath))
+                    {
+                        blockIncludes.WriteLine($"#include \"{includePath}\"");
+                    }
+                }
             }
         }
 
@@ -681,7 +716,7 @@ static private StringBuilder Build(
             globalDeclaration.WriteEventBuffers(eventListOutName, taskData.linkedEventOut.Length);
 
             var expressionToName = BuildExpressionToName(context, taskData);
-            BuildContextBlocks(context, taskData, expressionToName, out var blockFunction, out var blockCallFunction);
+            BuildContextBlocks(context, taskData, expressionToName, out var blockFunction, out var blockCallFunction, out var blockIncludes, out var blockDefines);
 
             //< Final composition
             var globalIncludeContent = new VFXShaderWriter();
@@ -743,25 +778,8 @@ static private StringBuilder Build(
             {
                 perPassIncludeContent.WriteLine("#include \"Packages/com.unity.visualeffectgraph/Shaders/VFXCommonOutput.hlsl\"");
             }
-
-            // Per-block defines
-            var defines = Enumerable.Empty<string>();
-            foreach (var block in context.activeFlattenedChildrenWithImplicit)
-                defines = defines.Concat(block.defines);
-            var uniqueDefines = new HashSet<string>(defines);
-            foreach (var define in uniqueDefines)
-                globalIncludeContent.WriteLineFormat("#define {0}{1}", define, define.Contains(' ') ? "" : " 1");
-
-            // Per-block includes
-            var includes = Enumerable.Empty<string>();
-            foreach (var block in context.activeFlattenedChildrenWithImplicit.OfType<IHLSLCodeHolder>())
-                includes = includes.Concat(block.includes);
-            foreach (var hlslHolder in taskData.hlslCodeHolders)
-                includes = includes.Concat(hlslHolder.includes);
-            var uniqueIncludes = new HashSet<string>(includes);
-            foreach (var includePath in uniqueIncludes)
-                perPassIncludeContent.WriteLine(string.Format("#include \"{0}\"", includePath));
-
+            globalIncludeContent.Write(blockDefines.builder.ToString());
+            perPassIncludeContent.Write(blockIncludes.builder.ToString());
 
             ReplaceMultiline(stringBuilder, "${VFXGlobalInclude}", globalIncludeContent.builder);
             ReplaceMultiline(stringBuilder, "${VFXGlobalDeclaration}", globalDeclaration.builder);

Packages/com.unity.visualeffectgraph/Editor/ShaderGraph/Templates/VFXConfig.template.hlsl

@@ -56,14 +56,17 @@ UNITY_INSTANCING_BUFFER_END(PerInstance)
     #define VFX_GET_INSTANCE_ID(i)      input.instanceID
 #endif
 
+$splice(VFXPerBlockDefines)
+
 $splice(VFXSRPCommonInclude)
 #include "Packages/com.unity.visualeffectgraph/Shaders/VFXCommon.hlsl"
+$splice(VFXPerBlockIncludes)
+#include "Packages/com.unity.visualeffectgraph/Shaders/VFXCommonOutput.hlsl"
 
 $splice(VFXParameterBuffer)
 
 $splice(VFXGeneratedBlockFunction)
 
-#include "Packages/com.unity.visualeffectgraph/Shaders/VFXCommonOutput.hlsl"
 
 struct AttributesElement
 {

Packages/com.unity.visualeffectgraph/Editor/ShaderGraph/VFXSubTarget.cs

@@ -140,6 +140,8 @@ static StructDescriptor GenerateVFXAttributesStruct(VFXContext context, VFXAttri
 
         static void GenerateVFXAdditionalCommands(VFXContext context, VFXSRPBinder srp, VFXSRPBinder.ShaderGraphBinder shaderGraphBinder, VFXTaskCompiledData taskData,
             out AdditionalCommandDescriptor srpCommonInclude,
+            out AdditionalCommandDescriptor perBlockDefines,
+            out AdditionalCommandDescriptor perBlockIncludes,
             out AdditionalCommandDescriptor loadAttributeDescriptor,
             out AdditionalCommandDescriptor blockFunctionDescriptor,
             out AdditionalCommandDescriptor blockCallFunctionDescriptor,
@@ -177,7 +179,14 @@ static void GenerateVFXAdditionalCommands(VFXContext context, VFXSRPBinder srp,
 
             // Graph Blocks
             var expressionToName = VFXCodeGenerator.BuildExpressionToName(context, taskData);
-            VFXCodeGenerator.BuildContextBlocks(context, taskData, expressionToName, out var blockFunction, out var blockCallFunction);
+            VFXCodeGenerator.BuildContextBlocks(context, taskData, expressionToName,
+                out var blockFunction,
+                out var blockCallFunction,
+                out var blockIncludes,
+                out var blockDefines);
+
+            perBlockDefines = new AdditionalCommandDescriptor("VFXPerBlockDefines", blockDefines.builder.ToString());
+            perBlockIncludes = new AdditionalCommandDescriptor("VFXPerBlockIncludes", blockIncludes.builder.ToString());
 
             blockFunctionDescriptor = new AdditionalCommandDescriptor("VFXGeneratedBlockFunction", blockFunction.builder.ToString());
             blockCallFunctionDescriptor = new AdditionalCommandDescriptor("VFXProcessBlocks", blockCallFunction.builder.ToString());
@@ -452,6 +461,8 @@ internal static SubShaderDescriptor PostProcessSubShader(SubShaderDescriptor sub
             GenerateVFXAdditionalCommands(
                 context, srp, shaderGraphSRPInfo, data,
                 out var srpCommonInclude,
+                out var perBlockDefines,
+                out var perBlockIncludes,
                 out var loadAttributeDescriptor,
                 out var blockFunctionDescriptor,
                 out var blockCallFunctionDescriptor,
@@ -512,6 +523,8 @@ out var fragInputsDescriptor
                 passDescriptor.additionalCommands = new AdditionalCommandCollection
                 {
                     srpCommonInclude,
+                    perBlockDefines,
+                    perBlockIncludes,
                     loadAttributeDescriptor,
                     blockFunctionDescriptor,
                     blockCallFunctionDescriptor,

Tests/SRPTests/Projects/VisualEffectGraph_HDRP/Assets/AllTests/Editor/Tests/CustomHLSLOperatorTest.cs

@@ -105,6 +105,67 @@ public void Check_CustomHLSL_Operator_Use_Shader_File()
             Assert.AreEqual(expectedGeneratedCode, hlslExpression.customCode);
         }
 
+        [UnityTest, Description("Regression for UUM-69735")]
+        public IEnumerator Check_CustomHLSL_Operator_Use_Shader_File_And_ShaderGraph()
+        {
+            var hlslCode = @"
+float3 GetCustomColor(in float3 vec)
+{
+    return float3(0.1f, 0.2f, 0.3f);
+}";
+            var shaderInclude = CustomHLSLBlockTest.CreateShaderFile(hlslCode, out var shaderIncludePath);
+
+            var hlslOperator = ScriptableObject.CreateInstance<CustomHLSL>();
+            hlslOperator.SetSettingValue("m_HLSLCode", hlslCode);
+            hlslOperator.SetSettingValue("m_ShaderFile", shaderInclude);
+            hlslOperator.SetSettingValue("m_OperatorName", "GetCustomColor");
+            MakeSimpleGraphWithCustomHLSL(hlslOperator, out var view, out var graph);
+
+            var vfxAsset = graph.GetResource().asset;
+            var outputToReplace = graph.children.OfType<VFXContext>().First(o => o.contextType.HasFlag(VFXContextType.Output));
+            var shaderGraphVariant = VFXLibrary.GetContexts().First(o => o.model is VFXComposedParticleOutput);
+
+            var newShaderGraphOutput = shaderGraphVariant.CreateInstance();
+            graph.AddChild(newShaderGraphOutput);
+            Assert.IsTrue(newShaderGraphOutput.GetSetting("shaderGraph").valid);
+
+            var tracker = "Find_Me_In_Generated_Name";
+            newShaderGraphOutput.label = tracker;
+            newShaderGraphOutput.LinkFrom(outputToReplace.inputFlowSlot[0].link.First().context);
+            outputToReplace.UnlinkAll();
+            graph.RemoveChild(outputToReplace);
+
+            var blockSetColor = ScriptableObject.CreateInstance<Block.SetAttribute>();
+            blockSetColor.SetSettingValue("attribute", "color");
+            newShaderGraphOutput.AddChild(blockSetColor);
+            Assert.IsTrue(hlslOperator.outputSlots[0].Link(blockSetColor.inputSlots[0]));
+
+            var blockOrientCameraPlane = ScriptableObject.CreateInstance<Block.Orient>();
+            blockOrientCameraPlane.SetSettingValue("faceRay", true);
+            blockOrientCameraPlane.SetSettingValue("mode", Block.Orient.Mode.FaceCameraPlane);
+            newShaderGraphOutput.AddChild(blockOrientCameraPlane);
+            var defineToFind = blockOrientCameraPlane.defines.First();
+
+            AssetDatabase.ImportAsset(AssetDatabase.GetAssetPath(vfxAsset));
+            yield return null;
+
+            string source = null;
+            for (int shaderIndex = 0; shaderIndex < graph.GetResource().GetShaderSourceCount(); shaderIndex++)
+            {
+                var name = graph.GetResource().GetShaderSourceName(shaderIndex);
+                if (name.Contains(tracker))
+                {
+                    source = graph.GetResource().GetShaderSource(shaderIndex);
+                    break;
+                }
+            }
+            Assert.IsNotNull(source);
+            Assert.IsTrue(source.Contains(shaderIncludePath, StringComparison.Ordinal));
+            Assert.IsTrue(source.Contains(defineToFind, StringComparison.Ordinal));
+
+            yield return null;
+        }
+
         [UnityTest]
         public IEnumerator Check_CustomHLSL_Operator_Return_Type_Is_Void()
         {