feat(codegen): support custom grain-call return types

src/Orleans.Analyzers/AnalyzerReleases.Unshipped.md

@@ -4,6 +4,7 @@
 
 Rule ID | Category | Severity | Notes
 --------|----------|----------|-------
+ORLEANS0026 | Usage | Error | Invalid invokable base type mapping
 ORLEANS0014 | Usage | Warning | ConfigureAwaitAnalyzer, Grain code should not use ConfigureAwait(false) or ConfigureAwait without ContinueOnCapturedContext
 ORLEANS0016 | Orleans.Versioning | Warning | GrainInterfaceVersionAnalyzer, Grain interface not declared in OrleansContracts.txt
 ORLEANS0017 | Orleans.Versioning | Warning | GrainInterfaceVersionAnalyzer, Grain interface version mismatch between code and file

src/Orleans.Analyzers/GrainInterfaceMethodReturnTypeDiagnosticAnalyzer.cs

@@ -1,23 +1,32 @@
 using System;
+using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Linq;
 using Microsoft.CodeAnalysis;
-using Microsoft.CodeAnalysis.CSharp;
-using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Diagnostics;
+using Orleans.CodeGenerator;
 
 namespace Orleans.Analyzers
 {
     [DiagnosticAnalyzer(LanguageNames.CSharp)]
     public class GrainInterfaceMethodReturnTypeDiagnosticAnalyzer : DiagnosticAnalyzer
     {
         public const string DiagnosticId = "ORLEANS0009";
-        public const string Title = "Grain interfaces methods must return a compatible type";
-        public const string MessageFormat = $"Grain interfaces methods must return a compatible type, such as Task, Task<T>, ValueTask, ValueTask<T>, or void";
+        public const string Title = "Grain interface methods must return a registered grain-call return type";
+        public const string MessageFormat = "Grain interface methods must return a registered grain-call return type";
         public const string Category = "Usage";
+        public const string InvalidMappingDiagnosticId = "ORLEANS0026";
 
-        private static readonly DiagnosticDescriptor Rule = new DiagnosticDescriptor(DiagnosticId, Title, MessageFormat, Category, DiagnosticSeverity.Error, isEnabledByDefault: true);
+        private static readonly DiagnosticDescriptor Rule = new(DiagnosticId, Title, MessageFormat, Category, DiagnosticSeverity.Error, isEnabledByDefault: true);
+        private static readonly DiagnosticDescriptor InvalidMappingRule = new(
+            InvalidMappingDiagnosticId,
+            "Invalid invokable base type mapping",
+            "{0}",
+            Category,
+            DiagnosticSeverity.Error,
+            isEnabledByDefault: true);
 
-        public override ImmutableArray<DiagnosticDescriptor> SupportedDiagnostics { get; } = [Rule];
+        public override ImmutableArray<DiagnosticDescriptor> SupportedDiagnostics { get; } = [Rule, InvalidMappingRule];
 
         public override void Initialize(AnalysisContext context)
         {
@@ -30,53 +39,173 @@ public override void Initialize(AnalysisContext context)
                     return;
                 }
 
-                var builder = ImmutableHashSet.CreateBuilder<ITypeSymbol>(SymbolEqualityComparer.Default);
-
-                AddIfNotNull(builder, context.Compilation.GetTypeByMetadataName("System.Threading.Tasks.Task"));
-                AddIfNotNull(builder, context.Compilation.GetTypeByMetadataName("System.Threading.Tasks.Task`1"));
-                AddIfNotNull(builder, context.Compilation.GetTypeByMetadataName("System.Threading.Tasks.ValueTask"));
-                AddIfNotNull(builder, context.Compilation.GetTypeByMetadataName("System.Threading.Tasks.ValueTask`1"));
-                AddIfNotNull(builder, context.Compilation.GetTypeByMetadataName("System.Collections.Generic.IAsyncEnumerable`1"));
-                AddIfNotNull(builder, context.Compilation.GetSpecialType(SpecialType.System_Void));
-
-                context.RegisterSymbolAction(context => AnalyzeMethod(context, baseInterface, builder.ToImmutable()), SymbolKind.Method);
+                var generateMethodSerializersAttribute = context.Compilation.GetTypeByMetadataName("Orleans.GenerateMethodSerializersAttribute");
+                var proxyContexts = GetProxyContexts(
+                    context.Compilation.Assembly.GlobalNamespace,
+                    baseInterface,
+                    generateMethodSerializersAttribute);
+                var resolver = new InvokableBaseTypeResolver(context.Compilation);
+                context.RegisterSymbolAction(
+                    context => AnalyzeMethod(context, proxyContexts, resolver),
+                    SymbolKind.Method);
             });
+        }
 
+        private static void AnalyzeMethod(
+            SymbolAnalysisContext context,
+            ImmutableArray<ProxyContext> proxyContexts,
+            InvokableBaseTypeResolver resolver)
+        {
+            var symbol = (IMethodSymbol)context.Symbol;
+            if (symbol.ContainingType.TypeKind != TypeKind.Interface || symbol.IsStatic)
+            {
+                return;
+            }
 
-            static void AddIfNotNull(ImmutableHashSet<ITypeSymbol>.Builder builder, INamedTypeSymbol? symbol)
+            ResolverDiagnostic? diagnostic = null;
+            foreach (var proxyContext in proxyContexts)
             {
-                if (symbol is not null)
+                if (!TryGetContextMethod(proxyContext.InterfaceType, symbol, out var contextMethod))
+                {
+                    continue;
+                }
+
+                if (resolver.TryResolve(
+                    proxyContext.ProxyBaseType,
+                    contextMethod,
+                    proxyContext.InterfaceType,
+                    out _,
+                    out diagnostic))
                 {
-                    builder.Add(symbol);
+                    continue;
                 }
+
+                break;
             }
+
+            if (diagnostic is null)
+            {
+                return;
+            }
+
+            var syntaxReference = symbol.DeclaringSyntaxReferences[0];
+            if (diagnostic.Kind == ResolverDiagnosticKind.InvalidMapping)
+            {
+                context.ReportDiagnostic(Diagnostic.Create(
+                    InvalidMappingRule,
+                    diagnostic.Location ?? Location.Create(syntaxReference.SyntaxTree, syntaxReference.Span),
+                    diagnostic.Message));
+                return;
+            }
+
+            context.ReportDiagnostic(Diagnostic.Create(
+                Rule,
+                Location.Create(syntaxReference.SyntaxTree, syntaxReference.Span)));
         }
 
-        private static void AnalyzeMethod(SymbolAnalysisContext context, INamedTypeSymbol baseInterface, ImmutableHashSet<ITypeSymbol> supportedTypes)
+        private static bool TryGetContextMethod(
+            INamedTypeSymbol proxyInterface,
+            IMethodSymbol method,
+            out IMethodSymbol contextMethod)
         {
-            var symbol = (IMethodSymbol)context.Symbol;
+            if (SymbolEqualityComparer.Default.Equals(
+                proxyInterface.OriginalDefinition,
+                method.ContainingType.OriginalDefinition))
+            {
+                contextMethod = method;
+                return true;
+            }
 
-            if (symbol.ContainingType.TypeKind != TypeKind.Interface) return;
+            foreach (var inheritedInterface in proxyInterface.AllInterfaces)
+            {
+                if (!SymbolEqualityComparer.Default.Equals(
+                    inheritedInterface.OriginalDefinition,
+                    method.ContainingType.OriginalDefinition))
+                {
+                    continue;
+                }
 
-            // allow static interface methods to return any type
-            if (symbol.IsStatic)
-                return;
+                foreach (var candidate in inheritedInterface.GetMembers(method.Name).OfType<IMethodSymbol>())
+                {
+                    if (SymbolEqualityComparer.Default.Equals(candidate.OriginalDefinition, method.OriginalDefinition))
+                    {
+                        contextMethod = candidate;
+                        return true;
+                    }
+                }
+            }
+
+            contextMethod = null!;
+            return false;
+        }
 
-            var isIAddressableInterface = false;
-            foreach (var implementedInterface in symbol.ContainingType.AllInterfaces)
+        private static ImmutableArray<ProxyContext> GetProxyContexts(
+            INamespaceSymbol globalNamespace,
+            INamedTypeSymbol baseInterface,
+            INamedTypeSymbol? generateMethodSerializersAttribute)
+        {
+            if (generateMethodSerializersAttribute is null)
             {
-                if (implementedInterface.Equals(baseInterface, SymbolEqualityComparer.Default))
+                return [];
+            }
+
+            var result = new List<ProxyContext>();
+            AddNamespace(globalNamespace);
+            return [.. result
+                .OrderBy(static entry => entry.SourceOrderGroup)
+                .ThenBy(static entry => entry.FilePath, StringComparer.Ordinal)
+                .ThenBy(static entry => entry.Position)
+                .ThenBy(
+                    static entry => entry.InterfaceType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat),
+                    StringComparer.Ordinal)];
+
+            void AddNamespace(INamespaceSymbol @namespace)
+            {
+                foreach (var member in @namespace.GetMembers())
                 {
-                    isIAddressableInterface = true;
-                    break;
+                    if (member is INamespaceSymbol childNamespace)
+                    {
+                        AddNamespace(childNamespace);
+                    }
+                    else if (member is INamedTypeSymbol type)
+                    {
+                        AddType(type);
+                    }
                 }
             }
 
-            if (!isIAddressableInterface || supportedTypes.Contains(symbol.ReturnType.OriginalDefinition))
-                return;
+            void AddType(INamedTypeSymbol type)
+            {
+                if (type.TypeKind == TypeKind.Interface
+                    && type.AllInterfaces.Any(implemented =>
+                        SymbolEqualityComparer.Default.Equals(implemented, baseInterface))
+                    && InvokableBaseTypeResolver.TryGetProxyBaseType(
+                        type,
+                        generateMethodSerializersAttribute,
+                        out var proxyBaseType,
+                        out _))
+                {
+                    var location = type.Locations.FirstOrDefault(static candidate => candidate.IsInSource);
+                    result.Add(new ProxyContext(
+                        type,
+                        proxyBaseType,
+                        location is null ? 1 : 0,
+                        location?.SourceTree?.FilePath ?? string.Empty,
+                        location?.SourceSpan.Start ?? int.MaxValue));
+                }
 
-            var syntaxReference = symbol.DeclaringSyntaxReferences;
-            context.ReportDiagnostic(Diagnostic.Create(Rule, Location.Create(syntaxReference[0].SyntaxTree, syntaxReference[0].Span)));
+                foreach (var nestedType in type.GetTypeMembers())
+                {
+                    AddType(nestedType);
+                }
+            }
         }
+
+        private sealed record ProxyContext(
+            INamedTypeSymbol InterfaceType,
+            INamedTypeSymbol ProxyBaseType,
+            int SourceOrderGroup,
+            string FilePath,
+            int Position);
     }
 }

src/Orleans.Analyzers/Orleans.Analyzers.csproj

@@ -17,6 +17,10 @@
     <PackageReference Update="NETStandard.Library" PrivateAssets="all" />
   </ItemGroup>
 
+  <ItemGroup>
+    <Compile Include="..\Orleans.CodeGenerator.Shared\InvokableBaseTypeResolver.cs" Link="InvokableBaseTypeResolver.cs" />
+  </ItemGroup>
+
   <ItemGroup>
     <None Include="$(OutputPath)\$(AssemblyName).dll" Pack="true" PackagePath="analyzers/dotnet/cs" Visible="false" />
     <AdditionalFiles Include="AnalyzerReleases.Shipped.md" />