[Macros] Reproduce issue with peer+extension macro extension's methods not being checked as witnesses (#71717)

Co-authored-by: Pavel Yaskevich <pyaskevich@apple.com>

lib/AST/NameLookup.cpp

@@ -1844,13 +1844,18 @@ PotentialMacroExpansions PotentialMacroExpansionsInContextRequest::evaluate(
   auto containerDecl = container.getAsDecl();
   forEachPotentialAttachedMacro(containerDecl, MacroRole::Member, nameTracker);
 
-  // If the container is an extension that was created from an extension macro,
-  // look at the nominal declaration to find any extension macros.
-  if (auto ext = dyn_cast<ExtensionDecl>(containerDecl)) {
-    if (auto nominal = nominalForExpandedExtensionDecl(ext)) {
-      forEachPotentialAttachedMacro(
-          nominal, MacroRole::Extension, nameTracker);
-    }
+  // Extension macros on the type or extension.
+  {
+    NominalTypeDecl *nominal = nullptr;
+    // If the container is an extension that was created from an extension
+    // macro, look at the nominal declaration to find any extension macros.
+    if (auto ext = dyn_cast<ExtensionDecl>(containerDecl))
+      nominal = nominalForExpandedExtensionDecl(ext);
+    else
+      nominal = container.getBaseNominal();
+
+    if (nominal)
+      forEachPotentialAttachedMacro(nominal, MacroRole::Extension, nameTracker);
   }
 
   // Peer and freestanding declaration macros.
@@ -1911,18 +1916,23 @@ populateLookupTableEntryFromMacroExpansions(ASTContext &ctx,
   // names match.
   {
     MacroIntroducedNameTracker nameTracker;
-    if (auto ext = dyn_cast<ExtensionDecl>(container.getAsDecl())) {
-      if (auto nominal = nominalForExpandedExtensionDecl(ext)) {
-        forEachPotentialAttachedMacro(nominal, MacroRole::Extension, nameTracker);
+    NominalTypeDecl *nominal = nullptr;
+    // If the container is an extension that was created from an extension
+    // macro, look at the nominal declaration to find any extension macros.
+    if (auto ext = dyn_cast<ExtensionDecl>(container.getAsDecl()))
+      nominal = nominalForExpandedExtensionDecl(ext);
+    else
+      nominal = container.getBaseNominal();
+
+    if (nominal) {
+      forEachPotentialAttachedMacro(nominal,
+                                  MacroRole::Extension, nameTracker);
       if (nameTracker.shouldExpandForName(name)) {
-          (void)evaluateOrDefault(
-              ctx.evaluator,
-              ExpandExtensionMacros{nominal},
+        (void)evaluateOrDefault(ctx.evaluator, ExpandExtensionMacros{nominal},
                                 false);
       }
     }
   }
-  }
 
   auto dc = container.getAsDeclContext();
   auto *module = dc->getParentModule();

lib/Macros/Sources/SwiftMacros/DistributedProtocolMacro.swift

@@ -16,7 +16,8 @@ import SwiftOperators
 import SwiftSyntaxBuilder
 
 /// Introduces:
-/// - `distributed actor $MyDistributedActor<ActorSystem>`
+/// - `distributed actor $MyDistributedActor<ActorSystem>: $MyDistributedActor, _DistributedActorStub where ...`
+/// - `extension MyDistributedActor where Self: _DistributedActorStub {}`
 public struct DistributedProtocolMacro: ExtensionMacro, PeerMacro {
   public static func expansion(
     of node: AttributeSyntax,
@@ -48,7 +49,7 @@ public struct DistributedProtocolMacro: ExtensionMacro, PeerMacro {
 
     let extensionDecl: DeclSyntax =
       """
-      extension \(proto.name) {
+      extension \(proto.name.trimmed) where Self: Distributed._DistributedActorStub {
         \(raw: requirementStubs)
       }
       """
@@ -68,33 +69,10 @@ public struct DistributedProtocolMacro: ExtensionMacro, PeerMacro {
     let serializationRequirementType =
       "Codable"
 
-    let requirements =
-      proto.memberBlock.members.map { member in
-        member.trimmed
-      }
-    let requirementStubs = requirements
-      .map { req in
-        """
-        \(req) {
-            if #available(SwiftStdlib 5.11, *) {
-              Distributed._distributedStubFatalError()
-            } else {
-              fatalError()
-            }
-        }
-        """
-      }.joined(separator: "\n    ")
-
-    let extensionDecl: DeclSyntax =
-      """
-      extension \(proto.name) where Self: _DistributedActorStub {
-        \(raw: requirementStubs)
-      }
-      """
-    
     let stubActorDecl: DeclSyntax =
       """
-      distributed actor $\(proto.name)<ActorSystem>: \(proto.name), _DistributedActorStub
+      distributed actor $\(proto.name.trimmed)<ActorSystem>: \(proto.name.trimmed), 
+        Distributed._DistributedActorStub
         where ActorSystem: DistributedActorSystem<any \(raw: serializationRequirementType)>, 
           ActorSystem.ActorID: \(raw: serializationRequirementType) 
       { }

test/Distributed/Macros/distributed_macro_expansion_DistributedProtocol_inheritance.swift

@@ -7,10 +7,7 @@
 // RUN: %empty-directory(%t)
 // RUN: %empty-directory(%t-scratch)
 
-// RUN: %target-swift-frontend -typecheck -verify -disable-availability-checking -plugin-path %swift-plugin-dir -I %t -dump-macro-expansions %s -dump-macro-expansions 2>&1 | %FileCheck %s
-
-// FIXME: inheritance tests limited because cannot refer to any generated macro from the same module...
-// XFAIL: *
+// RUN: %target-swift-frontend -typecheck -verify -disable-availability-checking -plugin-path %swift-plugin-dir -I %t -dump-macro-expansions %s 2>&1 | %FileCheck %s
 
 import Distributed
 
@@ -19,62 +16,48 @@ typealias System = LocalTestingDistributedActorSystem
 @_DistributedProtocol
 protocol EmptyBase {}
 
-// TODO: allow this?
-//@_DistributedProtocol
-//extension EmptyBase {}
-
 // @_DistributedProtocol ->
 //
-// CHECK: @freestanding(declaration)
-// CHECK: macro _distributed_stubs_EmptyBase() =
-// CHECK:   #distributedStubs(
-// CHECK:     module: "main", protocolName: "EmptyBase",
-// CHECK:     stubProtocols: []
-// CHECK:   )
-//
-// CHECK: // distributed actor $EmptyBase <ActorSystem>: EmptyBase  where SerializationRequirement == any Codable {
-// CHECK: distributed actor $EmptyBase : EmptyBase  {
-// CHECK:   typealias ActorSystem = LocalTestingDistributedActorSystem // FIXME: remove this
-//
-// CHECK:   #distributedStubs(
-// CHECK:     module: "main", protocolName: "EmptyBase",
-// CHECK:     stubProtocols: []
-// CHECK:   )
+// CHECK: distributed actor $EmptyBase<ActorSystem>: EmptyBase,
+// CHECK:   Distributed._DistributedActorStub
+// CHECK:   where ActorSystem: DistributedActorSystem<any Codable>,
+// CHECK:         ActorSystem.ActorID: Codable
+// CHECK:  {
 // CHECK:  }
 
+// CHECK: extension EmptyBase where Self: Distributed._DistributedActorStub {
+// CHECK: }
+
+// ==== ------------------------------------------------------------------------
+
 @_DistributedProtocol
-protocol G3: DistributedActor, EmptyBase where SerializationRequirement == any Codable {
+protocol G3<ActorSystem>: DistributedActor, EmptyBase where ActorSystem: DistributedActorSystem<any Codable> {
   distributed func get() -> String
   distributed func greet(name: String) -> String
 }
 
 // @_DistributedProtocol ->
-//
-// Since we have also the EmptyBase we don't know what names it will introduce,
-// so this stubs macro must be "names: arbitrary":
-// CHECK: @freestanding(declaration, names: arbitrary)
-// CHECK: macro _distributed_stubs_G3() =
-// CHECK:   #distributedStubs(
-// CHECK:     module: "main", protocolName: "G3",
-// CHECK:     stubProtocols: ["EmptyBase"],
-// CHECK:     "distributed func get() -> String",
-// CHECK:     "distributed func greet(name: String) -> String"
-// CHECK:   )
-//
-// TODO: distributed actor $G3<ActorSystem>: Greeter where SerializationRequirement == any Codable {
-// CHECK: distributed actor $G3: G3, EmptyBase {
-// TODO:    Preferably, we could refer to our own macro like this: #_distributed_stubs_G3
-// WORKAROUND:
-// CHECK:   #distributedStubs(
-// CHECK:      module: "main", protocolName: "G3",
-// CHECK:      stubProtocols: ["EmptyBase"],
-// CHECK:      "distributed func get() -> String",
-// CHECK:      "distributed func greet(name: String) -> String"
-// CHECK:    )
-// CHECK:
-// FIXME: the below cannot find the macro because it's form the same module
-// CHECK:    // stub inherited members
-// CHECK:    #_distributed_stubs_EmptyBase
+// CHECK: distributed actor $G3<ActorSystem>: G3
+// CHECK:   Distributed._DistributedActorStub
+// CHECK:   where ActorSystem: DistributedActorSystem<any Codable>,
+// CHECK:         ActorSystem.ActorID: Codable
+// CHECK: {
+// CHECK: }
+
+// CHECK: extension G3 where Self: Distributed._DistributedActorStub {
+// CHECK:   func get() -> String {
+// CHECK:     if #available (SwiftStdlib 5.11, *) {
+// CHECK:       Distributed._distributedStubFatalError()
+// CHECK:     } else {
+// CHECK:       fatalError()
+// CHECK:     }
+// CHECK:   }
+// CHECK:   distributed func greet(name: String) -> String {
+// CHECK:     if #available (SwiftStdlib 5.11, *) {
+// CHECK:       Distributed._distributedStubFatalError()
+// CHECK:     } else {
+// CHECK:       fatalError()
+// CHECK:     }
+// CHECK:   }
 // CHECK: }
 
-// ==== ------------------------------------------------------------------------

test/Distributed/Macros/distributed_macro_expansion_DistributedProtocol_simple.swift

@@ -11,29 +11,26 @@
 
 import Distributed
 
-// FIXME: the errors below are bugs: the added methods should be considered witnesses rdar://123012943
-// expected-note@+1{{in expansion of macro '_DistributedProtocol' on protocol 'Greeter' here}}
 @_DistributedProtocol
 protocol Greeter: DistributedActor where ActorSystem: DistributedActorSystem<any Codable> {
-  // FIXME: this is a bug
-  // expected-note@+1{{protocol requires function 'greet(name:)' with type '(String) -> String'}}
   distributed func greet(name: String) -> String
 }
 
 // @_DistributedProtocol ->
 
-// CHECK: distributed actor $Greeter<ActorSystem>: Greeter, _DistributedActorStub
-// CHECK:   where ActorSystem: DistributedActorSystem<any Codable>,
-// CHECK:   ActorSystem.ActorID: Codable
-// CHECK: {
-// CHECK: }
+// CHECK: distributed actor $Greeter<ActorSystem>: Greeter,
+// CHECK-NEXT: Distributed._DistributedActorStub
+// CHECK-NEXT: where ActorSystem: DistributedActorSystem<any Codable>,
+// CHECK-NEXT: ActorSystem.ActorID: Codable
+// CHECK-NEXT: {
+// CHECK-NEXT: }
 
-// CHECK: extension Greeter {
-// CHECK:   distributed func greet(name: String) -> String {
-// CHECK:     if #available (SwiftStdlib 5.11, *) {
-// CHECK:       Distributed._distributedStubFatalError()
-// CHECK:     } else {
-// CHECK:       fatalError()
-// CHECK:     }
-// CHECK:   }
-// CHECK: }
+// CHECK: extension Greeter where Self: Distributed._DistributedActorStub {
+// CHECK-NEXT:   distributed func greet(name: String) -> String {
+// CHECK-NEXT:     if #available (SwiftStdlib 5.11, *) {
+// CHECK-NEXT:       Distributed._distributedStubFatalError()
+// CHECK-NEXT:     } else {
+// CHECK-NEXT:       fatalError()
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }

test/Macros/Inputs/syntax_macro_definitions.swift

@@ -2199,7 +2199,55 @@ public struct SingleMemberStubMacro: DeclarationMacro {
   }
 }
 
-public struct FakeCodeItemMacro: DeclarationMacro, PeerMacro {
+public struct GenerateStubsForProtocolRequirementsMacro: PeerMacro, ExtensionMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    attachedTo declaration: some DeclGroupSyntax,
+    providingExtensionsOf type: some TypeSyntaxProtocol,
+    conformingTo protocols: [TypeSyntax],
+    in context: some MacroExpansionContext
+  ) throws -> [ExtensionDeclSyntax] {
+    guard let proto = declaration.as(ProtocolDeclSyntax.self) else {
+      return []
+    }
+
+    let requirements =
+      proto.memberBlock.members.map { member in member.trimmed }
+    let requirementStubs = requirements
+      .map { req in
+        "\(req) { fatalError() }"
+      }
+      .joined(separator: "\n    ")
+
+    let extensionDecl: DeclSyntax =
+      """
+      extension \(proto.name) where Self: _TestStub {
+        \(raw: requirementStubs)
+      }
+      """
+    return [extensionDecl.cast(ExtensionDeclSyntax.self)]
+  }
+
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    guard let proto = declaration.as(ProtocolDeclSyntax.self) else {
+      return []
+    }
+
+    return [
+      """
+      struct __\(proto.name): \(proto.name), _TestStub {
+        init() {} 
+      }
+      """
+    ]
+  }
+}
+
+  public struct FakeCodeItemMacro: DeclarationMacro, PeerMacro {
   public static func expansion(
     of node: some FreestandingMacroExpansionSyntax,
     in context: some MacroExpansionContext

test/Macros/macro_expand_synthesized_members.swift

@@ -157,3 +157,21 @@ struct NestedMacroExpansion {}
 func callNestedExpansionMember() {
   NestedMacroExpansion.member()
 }
+
+@attached(peer, names: prefixed(`__`)) // introduces `__GenerateStubsForProtocolRequirements
+@attached(extension, names: arbitrary) // introduces `extension GenerateStubsForProtocolRequirements`
+macro GenerateStubsForProtocolRequirements() = #externalMacro(module: "MacroDefinition", type: "GenerateStubsForProtocolRequirementsMacro")
+
+protocol _TestStub {} // used by 'GenerateStubsForProtocolRequirements'
+
+@GenerateStubsForProtocolRequirements
+protocol MacroExpansionRequirements {
+  func hello(name: String) -> String
+}
+// struct __MacroExpansionRequirements: _TestStub where ...
+// extension MacroExpansionRequirements where Self: _TestStub ...
+
+func testWitnessStub() {
+  let stub: any MacroExpansionRequirements = __MacroExpansionRequirements()
+  _ = stub.hello(name: "Caplin")
+}