Merge pull request #120 from CommunityToolkit/dev/fix-messenger-generator-duplicates

Fix IRecipient generator for partial recipient declarations

Author: Sergio0694

CommunityToolkit.Mvvm.SourceGenerators/Messaging/IMessengerRegisterAllGenerator.cs

@@ -23,12 +23,21 @@ public sealed partial class IMessengerRegisterAllGenerator : IIncrementalGenerat
     /// <inheritdoc/>
     public void Initialize(IncrementalGeneratorInitializationContext context)
     {
-        // Get all class declarations
+        // Get all class declarations. This pipeline step also needs to filter out duplicate recipient
+        // definitions (it might happen if a recipient has partial declarations). To do this, all pairs
+        // of class declarations and associated symbols are gathered, and then only the pair where the
+        // class declaration is the first syntax reference for the associated symbol is kept.
         IncrementalValuesProvider<INamedTypeSymbol> typeSymbols =
             context.SyntaxProvider
             .CreateSyntaxProvider(
                 static (node, _) => node is ClassDeclarationSyntax,
-                static (context, _) => (INamedTypeSymbol)context.SemanticModel.GetDeclaredSymbol(context.Node)!);
+                static (context, _) => (context.Node, Symbol: (INamedTypeSymbol)context.SemanticModel.GetDeclaredSymbol(context.Node)!))
+            .Where(static item =>
+                item.Symbol.DeclaringSyntaxReferences.Length > 0 &&
+                item.Symbol.DeclaringSyntaxReferences[0] is SyntaxReference syntaxReference &&
+                syntaxReference.SyntaxTree == item.Node.SyntaxTree &&
+                syntaxReference.Span == item.Node.Span)
+            .Select(static (item, _) => item.Symbol);
 
         // Get the target IRecipient<TMessage> interfaces and filter out other types
         IncrementalValuesProvider<(INamedTypeSymbol Type, ImmutableArray<INamedTypeSymbol> Interfaces)> typeAndInterfaceSymbols =

tests/CommunityToolkit.Mvvm.UnitTests/Test_IRecipientGenerator.cs

@@ -16,11 +16,11 @@ public partial class Test_IRecipientGenerator
     [TestMethod]
     public void Test_IRecipientGenerator_GeneratedRegistration()
     {
-        StrongReferenceMessenger? messenger = new();
-        RecipientWithSomeMessages? recipient = new();
+        StrongReferenceMessenger messenger = new();
+        RecipientWithSomeMessages recipient = new();
 
-        MessageA? messageA = new();
-        MessageB? messageB = new();
+        MessageA messageA = new();
+        MessageB messageB = new();
 
         Action<IMessenger, object, int> registrator = Messaging.__Internals.__IMessengerExtensions.CreateAllMessagesRegistratorWithToken<int>(recipient);
 
@@ -43,6 +43,15 @@ public void Test_IRecipientGenerator_GeneratedRegistration()
         Assert.AreSame(recipient.B, messageB);
     }
 
+    [TestMethod]
+    public void Test_IRecipientGenerator_TypeWithMultipleClassDeclarations()
+    {
+        RecipientWithMultipleClassDeclarations recipient = new();
+
+        // This test really just needs to verify this compiles and executes normally
+        _ = Messaging.__Internals.__IMessengerExtensions.CreateAllMessagesRegistratorWithToken<int>(recipient);
+    }
+
     public sealed class RecipientWithSomeMessages :
         IRecipient<MessageA>,
         IRecipient<MessageB>
@@ -69,4 +78,17 @@ public sealed class MessageA
     public sealed class MessageB
     {
     }
+
+    public sealed partial class RecipientWithMultipleClassDeclarations : IRecipient<MessageA>
+    {
+        public void Receive(MessageA message)
+        {
+        }
+    }
+
+    // This empty partial type declarations needs to be present to ensure the generator
+    // correctly handles cases where the source type has multiple class declarations.
+    partial class RecipientWithMultipleClassDeclarations
+    {
+    }
 }