Support reusing the semantic model provided by the source generator.

src/ComputeSharp.D2D1.SourceGenerators/D2DPixelShaderDescriptorGenerator.HlslSource.cs

@@ -27,7 +27,7 @@ private static partial class HlslSource
         /// Gathers all necessary information on a transpiled HLSL source for a given shader type.
         /// </summary>
         /// <param name="diagnostics">The collection of produced <see cref="DiagnosticInfo"/> instances.</param>
-        /// <param name="compilation">The input <see cref="Compilation"/> object currently in use.</param>
+        /// <param name="semanticModelProvider">The <see cref="SemanticModelProvider"/> instance currently in use.</param>
         /// <param name="structDeclarationSymbol">The <see cref="INamedTypeSymbol"/> for the shader type.</param>
         /// <param name="shaderInterfaceType">The shader interface type implemented by the shader type.</param>
         /// <param name="inputCount">The number of inputs for the shader.</param>
@@ -37,7 +37,7 @@ private static partial class HlslSource
         /// <returns>The HLSL source for the shader.</returns>
         public static string GetHlslSource(
             ImmutableArrayBuilder<DiagnosticInfo> diagnostics,
-            Compilation compilation,
+            SemanticModelProvider semanticModelProvider,
             INamedTypeSymbol structDeclarationSymbol,
             INamedTypeSymbol shaderInterfaceType,
             int inputCount,
@@ -70,8 +70,6 @@ public static string GetHlslSource(
 
             token.ThrowIfCancellationRequested();
 
-            SemanticModelProvider semanticModelProvider = new(compilation);
-
             // Explore the syntax tree and extract the processed info
             (string entryPoint, ImmutableArray<HlslMethod> processedMethods) = GetProcessedMethods(
                 diagnostics,

src/ComputeSharp.D2D1.SourceGenerators/D2DPixelShaderDescriptorGenerator.cs

@@ -139,10 +139,12 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 
                     using ImmutableArrayBuilder<DiagnosticInfo> diagnostics = new();
 
+                    SemanticModelProvider semanticModelProvider = new(context.SemanticModel);
+
                     // Get HLSL source for HlslSource
                     string hlslSource = HlslSource.GetHlslSource(
                         diagnostics,
-                        context.SemanticModel.Compilation,
+                        semanticModelProvider,
                         typeSymbol,
                         shaderInterfaceType,
                         inputCount,

src/ComputeSharp.SourceGeneration.Hlsl/Helpers/SemanticModelProvider.cs

@@ -6,13 +6,15 @@ namespace ComputeSharp.SourceGeneration.Helpers;
 /// <summary>
 /// A type providing <see cref="SemanticModel"/> instances for nodes.
 /// </summary>
-/// <param name="compilation">The source <see cref="Compilation"/> instance.</param>
-internal sealed class SemanticModelProvider(Compilation compilation)
+/// <param name="baseSemanticModel">A <see cref="SemanticModel"/> whose <see cref="SemanticModel.Compilation"/>
+/// will be used to get <see cref="SemanticModel"/> instances for other syntax trees.</param>
+internal sealed class SemanticModelProvider(SemanticModel baseSemanticModel)
 {
     /// <summary>
-    /// The map of loaded <see cref="SemanticModel"/> instances.
+    /// The map of loaded <see cref="SemanticModel"/> instances for syntax trees other than
+    /// the one for <see cref="baseSemanticModel"/>.
     /// </summary>
-    private readonly Dictionary<SyntaxTree, SemanticModel> semanticModelsMap = [];
+    private Dictionary<SyntaxTree, SemanticModel>? additionalSemanticModels = null;
 
     /// <summary>
     /// Gets a <see cref="SemanticModel"/> instance with info on a given <see cref="SyntaxNode"/>.
@@ -21,11 +23,20 @@ internal sealed class SemanticModelProvider(Compilation compilation)
     /// <returns>A <see cref="SemanticModel"/> instance containing info on <paramref name="syntaxNode"/>.</returns>
     public SemanticModel For(SyntaxNode syntaxNode)
     {
-        if (!this.semanticModelsMap.TryGetValue(syntaxNode.SyntaxTree, out SemanticModel semanticModel))
+        // Reuse the base semantic model if the syntax node belongs to the same tree.
+        // This will avoid creating new semantic models if the entire type's definition is in the same file.
+        if (syntaxNode.SyntaxTree == baseSemanticModel.SyntaxTree)
         {
-            semanticModel = compilation.GetSemanticModel(syntaxNode.SyntaxTree);
+            return baseSemanticModel;
+        }
+
+        this.additionalSemanticModels ??= [];
+
+        if (!this.additionalSemanticModels.TryGetValue(syntaxNode.SyntaxTree, out SemanticModel semanticModel))
+        {
+            semanticModel = baseSemanticModel.Compilation.GetSemanticModel(syntaxNode.SyntaxTree);
 
-            this.semanticModelsMap.Add(syntaxNode.SyntaxTree, semanticModel);
+            this.additionalSemanticModels.Add(syntaxNode.SyntaxTree, semanticModel);
         }
 
         return semanticModel;

src/ComputeSharp.SourceGenerators/ComputeShaderDescriptorGenerator.HlslSource.cs

@@ -29,7 +29,7 @@ internal static partial class HlslSource
         /// Gathers all necessary information on a transpiled HLSL source for a given shader type.
         /// </summary>
         /// <param name="diagnostics">The collection of produced <see cref="DiagnosticInfo"/> instances.</param>
-        /// <param name="compilation">The input <see cref="Compilation"/> object currently in use.</param>
+        /// <param name="semanticModelProvider">The <see cref="SemanticModelProvider"/> instance currently in use.</param>
         /// <param name="structDeclarationSymbol">The <see cref="INamedTypeSymbol"/> for the shader type.</param>
         /// <param name="shaderInterfaceType">The shader interface type implemented by the shader type.</param>
         /// <param name="isPixelShaderLike">Whether <paramref name="structDeclarationSymbol"/> is a "pixel shader like" type.</param>
@@ -42,7 +42,7 @@ internal static partial class HlslSource
         /// <param name="hlslSource">The resulting HLSL source for the current shader.</param>
         public static void GetInfo(
             ImmutableArrayBuilder<DiagnosticInfo> diagnostics,
-            Compilation compilation,
+            SemanticModelProvider semanticModelProvider,
             INamedTypeSymbol structDeclarationSymbol,
             INamedTypeSymbol shaderInterfaceType,
             bool isPixelShaderLike,
@@ -89,8 +89,6 @@ public static void GetInfo(
 
             token.ThrowIfCancellationRequested();
 
-            SemanticModelProvider semanticModelProvider = new(compilation);
-
             (string entryPoint, ImmutableArray<HlslMethod> processedMethods, isSamplerUsed) = GetProcessedMethods(
                 diagnostics,
                 structDeclarationSymbol,

src/ComputeSharp.SourceGenerators/ComputeShaderDescriptorGenerator.cs

@@ -90,10 +90,12 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
 
                     using ImmutableArrayBuilder<DiagnosticInfo> diagnostics = new();
 
+                    SemanticModelProvider semanticModelProvider = new(context.SemanticModel);
+
                     // Transpiled HLSL source info
                     HlslSource.GetInfo(
                         diagnostics,
-                        context.SemanticModel.Compilation,
+                        semanticModelProvider,
                         typeSymbol,
                         shaderInterfaceType,
                         isPixelShaderLike,