[Distributed][Macro] Handle more cases in distributed protocol macro (#72177)

Author: ktoso

lib/Macros/Sources/SwiftMacros/DistributedProtocolMacro.swift

@@ -19,6 +19,15 @@ import SwiftSyntaxBuilder
 /// - `distributed actor $MyDistributedActor<ActorSystem>: $MyDistributedActor, _DistributedActorStub where ...`
 /// - `extension MyDistributedActor where Self: _DistributedActorStub {}`
 public struct DistributedProtocolMacro: ExtensionMacro, PeerMacro {
+}
+
+// ===== -----------------------------------------------------------------------
+// MARK: Default Stub implementations Extension
+
+extension DistributedProtocolMacro {
+
+  /// Introduce the `extension MyDistributedActor` which contains default
+  /// implementations of the protocol's requirements.
   public static func expansion(
     of node: AttributeSyntax,
     attachedTo declaration: some DeclGroupSyntax,
@@ -27,25 +36,24 @@ public struct DistributedProtocolMacro: ExtensionMacro, PeerMacro {
     in context: some MacroExpansionContext
   ) throws -> [ExtensionDeclSyntax] {
     guard let proto = declaration.as(ProtocolDeclSyntax.self) else {
+      // we diagnose here, only once
+      try throwIllegalTargetDecl(node: node, declaration)
+    }
+
+    guard !proto.memberBlock.members.isEmpty else {
+      // ok, the protocol has no requirements so we no-op it
       return []
     }
 
+    let accessModifiers: String = proto.accessModifiersString
+
     let requirements =
       proto.memberBlock.members.map { member in
         member.trimmed
       }
     let requirementStubs = requirements
-      .map { req in
-        """
-        \(req) {
-            if #available(SwiftStdlib 6.0, *) {
-              Distributed._distributedStubFatalError()
-            } else {
-              fatalError()
-            }
-        }
-        """
-      }.joined(separator: "\n    ")
+      .map { stubMethod(access: accessModifiers, $0) }
+      .joined(separator: "\n    ")
 
     let extensionDecl: DeclSyntax =
       """
@@ -56,30 +64,236 @@ public struct DistributedProtocolMacro: ExtensionMacro, PeerMacro {
     return [extensionDecl.cast(ExtensionDeclSyntax.self)]
   }
 
+  static func stubMethod(access: String, _ requirementDeclaration: MemberBlockItemListSyntax.Element) -> String {
+    """
+    \(access)\(requirementDeclaration) {
+      \(stubFunctionBody())
+    }
+    """
+  }
+
+  static func stubFunctionBody() -> DeclSyntax {
+    """
+    if #available(SwiftStdlib 6.0, *) {
+      Distributed._distributedStubFatalError()
+    } else {
+      fatalError()
+    }
+    """
+  }
+}
+
+// ===== -----------------------------------------------------------------------
+// MARK: Distributed Actor Stub type
+
+extension DistributedProtocolMacro {
+
+  /// Introduce the `distributed actor` stub type.
   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 {
+      // don't diagnose here (again),
+      // we'll already report an error here from the other macro role
       return []
     }
 
-    // FIXME must detect this off the protocol
-    let serializationRequirementType =
-      "Codable"
+    var isGenericStub = false
+    var specificActorSystemRequirement: TypeSyntax?
 
-    let stubActorDecl: DeclSyntax =
-      """
-      distributed actor $\(proto.name.trimmed)<ActorSystem>: \(proto.name.trimmed), 
-        Distributed._DistributedActorStub
-        where ActorSystem: DistributedActorSystem<any \(raw: serializationRequirementType)>, 
-          ActorSystem.ActorID: \(raw: serializationRequirementType) 
-      { }
-      """
+    if proto.genericWhereClause == nil {
+      throw DiagnosticsError(
+        syntax: node,
+        message: """
+                 Distributed protocol must declare actor system with SerializationRequirement, for example:
+                    protocol Greeter<ActorSystem>: DistributedActor where ActorSystem: DistributedActorSystem<any Codable>
+                 """, id: .invalidApplication)
+    }
+
+    let accessModifiers = proto.accessModifiersString
+
+    for req in proto.genericWhereClause?.requirements ?? [] {
+      switch req.requirement {
+      case .conformanceRequirement(let conformanceReq)
+           where conformanceReq.leftType.isActorSystem:
+        specificActorSystemRequirement = conformanceReq.rightType.trimmed
+        isGenericStub = true
+
+      case .sameTypeRequirement(let sameTypeReq)
+           where sameTypeReq.leftType.isActorSystem:
+        specificActorSystemRequirement = sameTypeReq.rightType.trimmed
+        isGenericStub = false
+
+      default:
+        continue
+      }
+    }
+
+    if isGenericStub, let specificActorSystemRequirement {
+      return [
+        """
+        \(proto.modifiers) distributed actor $\(proto.name.trimmed)<ActorSystem>: \(proto.name.trimmed), 
+          Distributed._DistributedActorStub
+          where ActorSystem: \(specificActorSystemRequirement) 
+        { }
+        """
+      ]
+    } else if let specificActorSystemRequirement {
+      return [
+        """
+        \(proto.modifiers) distributed actor $\(proto.name.trimmed): \(proto.name.trimmed), 
+          Distributed._DistributedActorStub
+        { 
+          \(typealiasActorSystem(access: accessModifiers, proto, specificActorSystemRequirement)) 
+        }
+        """
+      ]
+    } else {
+      // there may be no `where` clause specifying an actor system,
+      // but perhaps there is a typealias (or extension with a typealias),
+      // specifying a concrete actor system so we let this synthesize
+      // an empty `$Greeter` -- this may fail, or succeed depending on
+      // surrounding code using a default distributed actor system,
+      // or extensions providing it.
+      return [
+        """
+        \(proto.modifiers) distributed actor $\(proto.name.trimmed): \(proto.name.trimmed), 
+          Distributed._DistributedActorStub
+        {
+        }
+        """
+      ]
+    }
+  }
 
-    // return [extensionDecl, stubActorDecl]
-    return [stubActorDecl]
+  private static func typealiasActorSystem(access: String, _ proto: ProtocolDeclSyntax, _ type: TypeSyntax) -> DeclSyntax {
+    "\(raw: access)typealias ActorSystem = \(type)"
   }
+}
 
+// ===== -----------------------------------------------------------------------
+// MARK: Convenience Extensions
+
+extension ProtocolDeclSyntax {
+  var accessModifiersString: String {
+    let modifiers = modifiers.filter { modifier in
+        modifier.isAccessControl
+      }
+
+    guard !modifiers.isEmpty else {
+      return ""
+    }
+
+    let string = modifiers
+      .map { "\($0.trimmed)" }
+      .joined(separator: " ")
+    return "\(string) "
+  }
+}
+
+extension TypeSyntax {
+  fileprivate var isActorSystem: Bool {
+    self.trimmedDescription == "ActorSystem"
+  }
+}
+
+extension DeclSyntaxProtocol {
+  var isClass: Bool {
+    return self.is(ClassDeclSyntax.self)
+  }
+
+  var isActor: Bool {
+    return self.is(ActorDeclSyntax.self)
+  }
+
+  var isEnum: Bool {
+    return self.is(EnumDeclSyntax.self)
+  }
+
+  var isStruct: Bool {
+    return self.is(StructDeclSyntax.self)
+  }
+}
+
+extension DeclModifierSyntax {
+  var isAccessControl: Bool {
+    switch self.name.tokenKind {
+    case .keyword(.private): fallthrough
+    case .keyword(.fileprivate): fallthrough
+    case .keyword(.internal): fallthrough
+    case .keyword(.package): fallthrough
+    case .keyword(.public):
+      return true
+    default:
+      return false
+    }
+  }
+}
+
+// ===== -----------------------------------------------------------------------
+// MARK: DistributedProtocol macro errors
+
+extension DistributedProtocolMacro {
+  static func throwIllegalTargetDecl(node: AttributeSyntax, _ declaration: some DeclSyntaxProtocol) throws -> Never {
+    let kind: String
+    if declaration.isClass {
+      kind = "class"
+    } else if declaration.isActor {
+      kind = "actor"
+    } else if declaration.isStruct {
+      kind = "struct"
+    } else if declaration.isStruct {
+      kind = "enum"
+    } else {
+      kind = "\(declaration.kind)"
+    }
+
+    throw DiagnosticsError(
+      syntax: node,
+      message: "'@DistributedProtocol' can only be applied to 'protocol', but was attached to '\(kind)'", id: .invalidApplication)
+  }
+}
+
+struct DistributedProtocolMacroDiagnostic: DiagnosticMessage {
+  enum ID: String {
+    case invalidApplication = "invalid type"
+    case missingInitializer = "missing initializer"
+  }
+
+  var message: String
+  var diagnosticID: MessageID
+  var severity: DiagnosticSeverity
+
+  init(message: String, diagnosticID: SwiftDiagnostics.MessageID, severity: SwiftDiagnostics.DiagnosticSeverity = .error) {
+    self.message = message
+    self.diagnosticID = diagnosticID
+    self.severity = severity
+  }
+
+  init(message: String, domain: String, id: ID, severity: SwiftDiagnostics.DiagnosticSeverity = .error) {
+    self.message = message
+    self.diagnosticID = MessageID(domain: domain, id: id.rawValue)
+    self.severity = severity
+  }
+}
+
+extension DiagnosticsError {
+  init<S: SyntaxProtocol>(
+    syntax: S,
+    message: String,
+    domain: String = "Distributed",
+    id: DistributedProtocolMacroDiagnostic.ID,
+    severity: SwiftDiagnostics.DiagnosticSeverity = .error) {
+    self.init(diagnostics: [
+      Diagnostic(
+        node: Syntax(syntax),
+        message: DistributedProtocolMacroDiagnostic(
+          message: message,
+          domain: domain,
+          id: id,
+          severity: severity))
+    ])
+  }
 }

lib/Sema/TypeCheckDeclPrimary.cpp

@@ -3332,9 +3332,8 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
       }
     }
 
-    if (CD->isDistributedActor()) {
-      TypeChecker::checkDistributedActor(SF, CD);
-    }
+    // Check distributed actors
+    TypeChecker::checkDistributedActor(SF, CD);
 
     // Force lowering of stored properties.
     (void) CD->getStoredProperties();
@@ -3945,8 +3944,7 @@ class DeclChecker : public DeclVisitor<DeclChecker> {
 
     checkExplicitAvailability(ED);
 
-    if (nominal->isDistributedActor())
-      TypeChecker::checkDistributedActor(SF, nominal);
+    TypeChecker::checkDistributedActor(SF, nominal);
 
     diagnoseIncompatibleProtocolsForMoveOnlyType(ED);
     diagnoseExtensionOfMarkerProtocol(ED);

lib/Sema/TypeCheckDistributed.cpp

@@ -686,7 +686,7 @@ void swift::checkDistributedActorProperties(const NominalTypeDecl *decl) {
 // ==== ------------------------------------------------------------------------
 
 void TypeChecker::checkDistributedActor(SourceFile *SF, NominalTypeDecl *nominal) {
-  if (!nominal)
+  if (!nominal || !nominal->isDistributedActor())
     return;
 
   // ==== Ensure the Distributed module is available,

stdlib/public/Distributed/DistributedMacros.swift

@@ -15,8 +15,17 @@
 import Swift
 import _Concurrency
 
-#if $Macros
+// Macros are disabled when Swift is built without swift-syntax.
+#if $Macros && hasAttribute(attached)
 
+/// Enables the attached to protocol to be resolved as remote distributed
+/// actor reference.
+///
+/// ### Requirements
+///
+/// The attached to type must be a protocol that refines the `DistributedActor`
+/// protocol. It must either specify a concrete `ActorSystem` or constrain it
+/// in such way that the system's `SerializationRequirement` is statically known.
 @attached(peer, names: prefixed(`$`)) // provides $Greeter concrete stub type
 @attached(extension, names: arbitrary) // provides extension for Greeter & _DistributedActorStub
 public macro _DistributedProtocol() =

test/Distributed/Macros/distributed_macro_expansion_DistributedProtocol_errors.swift

@@ -0,0 +1,33 @@
+// REQUIRES: swift_swift_parser, asserts
+//
+// UNSUPPORTED: back_deploy_concurrency
+// REQUIRES: concurrency
+// REQUIRES: distributed
+//
+// RUN: %empty-directory(%t)
+// RUN: %empty-directory(%t-scratch)
+
+// RUN: %target-swift-frontend-emit-module -emit-module-path %t/FakeDistributedActorSystems.swiftmodule -module-name FakeDistributedActorSystems -disable-availability-checking %S/../Inputs/FakeDistributedActorSystems.swift
+// RUN: %target-swift-frontend -typecheck -verify -disable-availability-checking -plugin-path %swift-plugin-dir -parse-as-library -I %t %S/../Inputs/FakeDistributedActorSystems.swift -dump-macro-expansions %s 2>&1
+
+import Distributed
+
+@_DistributedProtocol // expected-error{{'@DistributedProtocol' can only be applied to 'protocol', but was attached to 'struct' (from macro '_DistributedProtocol')}}
+struct Struct {}
+
+@_DistributedProtocol // expected-error{{'@DistributedProtocol' can only be applied to 'protocol', but was attached to 'class' (from macro '_DistributedProtocol')}}
+class Clazz {}
+
+@_DistributedProtocol // expected-error{{'@DistributedProtocol' can only be applied to 'protocol', but was attached to 'actor' (from macro '_DistributedProtocol')}}
+actor Act {}
+
+@_DistributedProtocol // expected-error{{'@DistributedProtocol' can only be applied to 'protocol', but was attached to 'actor' (from macro '_DistributedProtocol')}}
+distributed actor Caplin {
+  typealias ActorSystem = FakeActorSystem
+}
+
+@_DistributedProtocol // expected-error{{Distributed protocol must declare actor system with SerializationRequirement}}
+protocol Fail: DistributedActor {
+  distributed func method() -> String
+}
+