Match explicit entry points for DX12 shaders

src/ComputeSharp.SourceGeneration/Extensions/ITypeSymbolExtensions.cs

@@ -9,6 +9,25 @@ namespace ComputeSharp.SourceGeneration.Extensions;
 /// </summary>
 internal static class ITypeSymbolExtensions
 {
+    /// <summary>
+    /// Gets the method of this symbol that have a particular name.
+    /// </summary>
+    /// <param name="symbol">The input <see cref="ITypeSymbol"/> instance to check.</param>
+    /// <param name="name">The name of the method to find.</param>
+    /// <returns>The target method, if present.</returns>
+    public static IMethodSymbol? GetMethod(this ITypeSymbol symbol, string name)
+    {
+        foreach (IMethodSymbol methodSymbol in symbol.GetMembers(name))
+        {
+            if (methodSymbol.Name == name)
+            {
+                return methodSymbol;
+            }
+        }
+
+        return null;
+    }
+
     /// <summary>
     /// Checks whether or not a given type symbol has a specified fully qualified metadata name.
     /// </summary>
@@ -28,7 +47,7 @@ public static bool HasFullyQualifiedMetadataName(this ITypeSymbol symbol, string
     /// Checks whether or not a given <see cref="ITypeSymbol"/> implements an interface of a specified type.
     /// </summary>
     /// <param name="typeSymbol">The target <see cref="ITypeSymbol"/> instance to check.</param>
-    /// <param name="interfaceSymbol">The <see cref="ITypeSymbol"/> instane to check for inheritance from.</param>
+    /// <param name="interfaceSymbol">The <see cref="ITypeSymbol"/> instance to check for inheritance from.</param>
     /// <returns>Whether or not <paramref name="typeSymbol"/> has an interface of type <paramref name="interfaceSymbol"/>.</returns>
     public static bool HasInterfaceWithType(this ITypeSymbol typeSymbol, ITypeSymbol interfaceSymbol)
     {

src/ComputeSharp.SourceGenerators/ComputeShaderDescriptorGenerator.Helpers.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics.CodeAnalysis;
 using Microsoft.CodeAnalysis;
 
 namespace ComputeSharp.SourceGenerators;
@@ -10,9 +11,14 @@ partial class ComputeShaderDescriptorGenerator
     /// </summary>
     /// <param name="typeSymbol">The input <see cref="INamedTypeSymbol"/> instance to check.</param>
     /// <param name="compilation">The <see cref="Compilation"/> instance currently in use.</param>
+    /// <param name="shaderInterfaceType">The (constructed) shader interface type implemented by the shader type.</param>
     /// <param name="isPixelShaderLike">Whether <paramref name="typeSymbol"/> is a "pixel shader like" type.</param>
     /// <returns>Whether <paramref name="typeSymbol"/> is a compute shader type at all.</returns>
-    private static bool TryGetIsPixelShaderLike(INamedTypeSymbol typeSymbol, Compilation compilation, out bool isPixelShaderLike)
+    private static bool TryGetIsPixelShaderLike(
+        INamedTypeSymbol typeSymbol,
+        Compilation compilation,
+        [NotNullWhen(true)] out INamedTypeSymbol? shaderInterfaceType,
+        out bool isPixelShaderLike)
     {
         INamedTypeSymbol computeShaderSymbol = compilation.GetTypeByMetadataName("ComputeSharp.IComputeShader")!;
         INamedTypeSymbol pixelShaderSymbol = compilation.GetTypeByMetadataName("ComputeSharp.IComputeShader`1")!;
@@ -21,18 +27,21 @@ private static bool TryGetIsPixelShaderLike(INamedTypeSymbol typeSymbol, Compila
         {
             if (SymbolEqualityComparer.Default.Equals(interfaceSymbol, computeShaderSymbol))
             {
+                shaderInterfaceType = interfaceSymbol;
                 isPixelShaderLike = false;
 
                 return true;
             }
             else if (SymbolEqualityComparer.Default.Equals(interfaceSymbol.ConstructedFrom, pixelShaderSymbol))
             {
+                shaderInterfaceType = interfaceSymbol;
                 isPixelShaderLike = true;
 
                 return true;
             }
         }
 
+        shaderInterfaceType = null;
         isPixelShaderLike = false;
 
         return false;

src/ComputeSharp.SourceGenerators/ComputeShaderDescriptorGenerator.HlslSource.cs

@@ -31,6 +31,8 @@ internal static partial class HlslSource
         /// <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="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>
         /// <param name="threadsX">The thread ids value for the X axis.</param>
         /// <param name="threadsY">The thread ids value for the Y axis.</param>
         /// <param name="threadsZ">The thread ids value for the Z axis.</param>
@@ -42,6 +44,8 @@ public static void GetInfo(
             ImmutableArrayBuilder<DiagnosticInfo> diagnostics,
             Compilation compilation,
             INamedTypeSymbol structDeclarationSymbol,
+            INamedTypeSymbol shaderInterfaceType,
+            bool isPixelShaderLike,
             int threadsX,
             int threadsY,
             int threadsZ,
@@ -53,6 +57,8 @@ public static void GetInfo(
             // Detect any invalid properties
             HlslDefinitionsSyntaxProcessor.DetectAndReportInvalidPropertyDeclarations(diagnostics, structDeclarationSymbol);
 
+            token.ThrowIfCancellationRequested();
+
             // We need to sets to track all discovered custom types and static methods
             HashSet<INamedTypeSymbol> discoveredTypes = new(SymbolEqualityComparer.Default);
             Dictionary<IMethodSymbol, MethodDeclarationSyntax> staticMethods = new(SymbolEqualityComparer.Default);
@@ -62,9 +68,8 @@ public static void GetInfo(
             Dictionary<IFieldSymbol, HlslStaticField> staticFieldDefinitions = new(SymbolEqualityComparer.Default);
 
             // Setup the semantic model and basic properties
-            INamedTypeSymbol? pixelShaderSymbol = structDeclarationSymbol.AllInterfaces.FirstOrDefault(static interfaceSymbol => interfaceSymbol is { IsGenericType: true, Name: "IComputeShader" });
-            bool isComputeShader = pixelShaderSymbol is null;
-            string? implicitTextureType = isComputeShader ? null : HlslKnownTypes.GetMappedNameForPixelShaderType(pixelShaderSymbol!);
+            bool isComputeShader = !isPixelShaderLike;
+            string? implicitTextureType = HlslKnownTypes.GetMappedNameForPixelShaderType(shaderInterfaceType);
 
             token.ThrowIfCancellationRequested();
 
@@ -90,6 +95,7 @@ public static void GetInfo(
             (string entryPoint, ImmutableArray<HlslMethod> processedMethods, isSamplerUsed) = GetProcessedMethods(
                 diagnostics,
                 structDeclarationSymbol,
+                shaderInterfaceType,
                 semanticModelProvider,
                 discoveredTypes,
                 staticMethods,
@@ -360,6 +366,7 @@ private static ImmutableArray<HlslSharedBuffer> GetSharedBuffers(
         /// </summary>
         /// <param name="diagnostics">The collection of produced <see cref="DiagnosticInfo"/> instances.</param>
         /// <param name="structDeclarationSymbol">The type symbol for the shader type.</param>
+        /// <param name="shaderInterfaceType">The shader interface type implemented by the shader type.</param>
         /// <param name="semanticModel">The <see cref="SemanticModelProvider"/> instance for the type to process.</param>
         /// <param name="discoveredTypes">The collection of currently discovered types.</param>
         /// <param name="staticMethods">The set of discovered and processed static methods.</param>
@@ -373,6 +380,7 @@ private static ImmutableArray<HlslSharedBuffer> GetSharedBuffers(
         private static (string EntryPoint, ImmutableArray<HlslMethod> Methods, bool IsSamplerUser) GetProcessedMethods(
             ImmutableArrayBuilder<DiagnosticInfo> diagnostics,
             INamedTypeSymbol structDeclarationSymbol,
+            INamedTypeSymbol shaderInterfaceType,
             SemanticModelProvider semanticModel,
             ICollection<INamedTypeSymbol> discoveredTypes,
             IDictionary<IMethodSymbol, MethodDeclarationSyntax> staticMethods,
@@ -385,6 +393,7 @@ private static (string EntryPoint, ImmutableArray<HlslMethod> Methods, bool IsSa
         {
             using ImmutableArrayBuilder<HlslMethod> methods = new();
 
+            IMethodSymbol entryPointInterfaceMethod = shaderInterfaceType.GetMethod("Execute")!;
             string? entryPoint = null;
             bool isSamplerUsed = false;
 
@@ -396,22 +405,15 @@ private static (string EntryPoint, ImmutableArray<HlslMethod> Methods, bool IsSa
                     continue;
                 }
 
+                // Ensure that we have accessible source information
                 if (!methodSymbol.TryGetSyntaxNode(token, out MethodDeclarationSyntax? methodDeclaration))
                 {
                     continue;
                 }
 
-                bool isShaderEntryPoint =
-                    (isComputeShader &&
-                     methodSymbol.Name == "Execute" &&
-                     methodSymbol.ReturnsVoid &&
-                     methodSymbol.TypeParameters.Length == 0 &&
-                     methodSymbol.Parameters.Length == 0) ||
-                    (!isComputeShader &&
-                     methodSymbol.Name == "Execute" &&
-                     methodSymbol.ReturnType is not null && // TODO: match for pixel type
-                     methodSymbol.TypeParameters.Length == 0 &&
-                     methodSymbol.Parameters.Length == 0);
+                // Check whether the current method is the entry point (ie. it's implementing 'Execute'). We use
+                // 'FindImplementationForInterfaceMember' to handle explicit interface implementations as well.
+                bool isShaderEntryPoint = structDeclarationSymbol.FindImplementationForInterfaceMember(entryPointInterfaceMethod) is not null;
 
                 // Except for the entry point, ignore explicit interface implementations
                 if (!isShaderEntryPoint && !methodSymbol.ExplicitInterfaceImplementations.IsDefaultOrEmpty)

src/ComputeSharp.SourceGenerators/ComputeShaderDescriptorGenerator.cs

@@ -54,7 +54,11 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                     }
 
                     // Check whether type is a compute shader, and if so, if it's pixel shader like
-                    if (!TryGetIsPixelShaderLike(typeSymbol, context.SemanticModel.Compilation, out bool isPixelShaderLike))
+                    if (!TryGetIsPixelShaderLike(
+                        typeSymbol,
+                        context.SemanticModel.Compilation,
+                        out INamedTypeSymbol? shaderInterfaceType,
+                        out bool isPixelShaderLike))
                     {
                         return default;
                     }
@@ -91,6 +95,8 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
                         diagnostics,
                         context.SemanticModel.Compilation,
                         typeSymbol,
+                        shaderInterfaceType,
+                        isPixelShaderLike,
                         threadsX,
                         threadsY,
                         threadsZ,

src/ComputeSharp.SourceGenerators/Mappings/HlslKnownTypes.cs

@@ -208,11 +208,18 @@ public static partial string GetMappedName(INamedTypeSymbol typeSymbol)
     /// <summary>
     /// Gets the mapped HLSL-compatible type name for the output texture of a pixel shader.
     /// </summary>
-    /// <param name="typeSymbol">The pixel shader type to map.</param>
+    /// <param name="typeSymbol">The shader type to map.</param>
     /// <returns>The HLSL-compatible type name that can be used in an HLSL shader.</returns>
-    public static string GetMappedNameForPixelShaderType(INamedTypeSymbol typeSymbol)
+    public static string? GetMappedNameForPixelShaderType(INamedTypeSymbol typeSymbol)
     {
-        string genericArgumentName = ((INamedTypeSymbol)typeSymbol.TypeArguments.First()).GetFullyQualifiedMetadataName();
+        // If the shader type is not a pixel shader type (ie. it has a type argument), stop here.
+        // At this point the input is guaranteed to either be 'IComputeShader' or 'IComputeShader<TPixel>'.
+        if (typeSymbol.TypeArguments is not [INamedTypeSymbol pixelShaderType])
+        {
+            return null;
+        }
+
+        string genericArgumentName = pixelShaderType.GetFullyQualifiedMetadataName();
 
         // If the current type is a custom type, format it as needed
         if (!KnownHlslTypeMetadataNames.TryGetValue(genericArgumentName, out string? mappedElementType))