Apply suite diagnostics to the entire lexical context. (#336)

grynspan · web-flow · commit 948b59df055e · 2024-04-09T05:01:00.000-04:00
diff --git a/Sources/TestingMacros/SuiteDeclarationMacro.swift b/Sources/TestingMacros/SuiteDeclarationMacro.swift
@@ -50,7 +50,7 @@ public struct SuiteDeclarationMacro: MemberMacro, PeerMacro, Sendable {
   /// - Returns: Whether or not macro expansion should continue (i.e. stopping
   ///   if a fatal error was diagnosed.)
   private static func _diagnoseIssues(
-    with declaration: some SyntaxProtocol,
+    with declaration: some DeclSyntaxProtocol,
     suiteAttribute: AttributeSyntax,
     in context: some MacroExpansionContext
   ) -> Bool {
@@ -59,43 +59,18 @@ public struct SuiteDeclarationMacro: MemberMacro, PeerMacro, Sendable {
       context.diagnose(diagnostics)
     }
 
-    // The @Suite attribute is only supported on type declarations, all of which
-    // are DeclGroupSyntax types.
-    guard let declaration = declaration.asProtocol((any DeclGroupSyntax).self) else {
-      diagnostics.append(.attributeNotSupported(suiteAttribute, on: declaration))
-      return false
-    }
-
 #if canImport(SwiftSyntax600)
     // Check if the lexical context is appropriate for a suite or test.
-    diagnostics += diagnoseIssuesWithLexicalContext(containing: declaration, attribute: suiteAttribute, in: context)
+    diagnostics += diagnoseIssuesWithLexicalContext(context.lexicalContext, containing: declaration, attribute: suiteAttribute)
 #endif
-
-    // Generic suites are not supported.
-    if let genericClause = declaration.asProtocol((any WithGenericParametersSyntax).self)?.genericParameterClause {
-      diagnostics.append(.genericDeclarationNotSupported(declaration, whenUsing: suiteAttribute, becauseOf: genericClause))
-    } else if let whereClause = declaration.genericWhereClause {
-      diagnostics.append(.genericDeclarationNotSupported(declaration, whenUsing: suiteAttribute, becauseOf: whereClause))
-    }
+    diagnostics += diagnoseIssuesWithLexicalContext(declaration, containing: declaration, attribute: suiteAttribute)
 
     // Suites inheriting from XCTestCase are not supported.
-    if declaration.inherits(fromTypeNamed: "XCTestCase", inModuleNamed: "XCTest") {
+    if let declaration = declaration.asProtocol((any DeclGroupSyntax).self),
+       declaration.inherits(fromTypeNamed: "XCTestCase", inModuleNamed: "XCTest") {
       diagnostics.append(.xcTestCaseNotSupported(declaration, whenUsing: suiteAttribute))
     }
 
-    // Suites that are classes must be final.
-    if let classDecl = declaration.as(ClassDeclSyntax.self) {
-      if !classDecl.modifiers.lazy.map(\.name.tokenKind).contains(.keyword(.final)) {
-        diagnostics.append(.nonFinalClassNotSupported(classDecl, whenUsing: suiteAttribute))
-      }
-    }
-
-    // Suites cannot be protocols (there's nowhere to put most of the
-    // declarations we generate.)
-    if let protocolDecl = declaration.as(ProtocolDeclSyntax.self) {
-      diagnostics.append(.attributeNotSupported(suiteAttribute, on: protocolDecl))
-    }
-
     // @Suite cannot be applied to a type extension (although a type extension
     // can still contain test functions and test suites.)
     if let extensionDecl = declaration.as(ExtensionDeclSyntax.self) {
@@ -106,24 +81,10 @@ public struct SuiteDeclarationMacro: MemberMacro, PeerMacro, Sendable {
     // impossible to reach this point if the declaration can't have attributes.
     if let attributedDecl = declaration.asProtocol((any WithAttributesSyntax).self) {
       // Only one @Suite attribute is supported.
-      let suiteAttributes = attributedDecl.attributes(named: "Suite", in: context)
+      let suiteAttributes = attributedDecl.attributes(named: "Suite")
       if suiteAttributes.count > 1 {
         diagnostics.append(.multipleAttributesNotSupported(suiteAttributes, on: declaration))
       }
-
-      // Availability is not supported on suites (we need semantic availability
-      // to correctly understand the availability of a suite.)
-      let availabilityAttributes = attributedDecl.availabilityAttributes
-      if !availabilityAttributes.isEmpty {
-        // Diagnose all @available attributes.
-        for availabilityAttribute in availabilityAttributes {
-          diagnostics.append(.availabilityAttributeNotSupported(availabilityAttribute, on: declaration, whenUsing: suiteAttribute))
-        }
-      } else if let noasyncAttribute = attributedDecl.noasyncAttribute {
-        // No @available attributes, but we do have an @_unavailableFromAsync
-        // attribute and we still need to diagnose that.
-        diagnostics.append(.availabilityAttributeNotSupported(noasyncAttribute, on: declaration, whenUsing: suiteAttribute))
-      }
     }
 
     return !diagnostics.lazy.map(\.severity).contains(.error)
diff --git a/Sources/TestingMacros/Support/Additions/WithAttributesSyntaxAdditions.swift b/Sources/TestingMacros/Support/Additions/WithAttributesSyntaxAdditions.swift
@@ -121,11 +121,10 @@ extension WithAttributesSyntax {
   ///   - name: The name of the attribute to look for.
   ///   - moduleName: The name of the module that declares the attribute named
   ///     `name`.
-  ///   - context: The macro context in which the expression is being parsed.
   ///
   /// - Returns: An array of `AttributeSyntax` corresponding to the attached
   ///   `@Test` attributes, or the empty array if none is attached.
-  func attributes(named name: String, inModuleNamed moduleName: String = "Testing", in context: some MacroExpansionContext) -> [AttributeSyntax] {
+  func attributes(named name: String, inModuleNamed moduleName: String = "Testing") -> [AttributeSyntax] {
     attributes.lazy.compactMap { attribute in
       if case let .attribute(attribute) = attribute {
         return attribute
diff --git a/Sources/TestingMacros/Support/DiagnosticMessage+Diagnosing.swift b/Sources/TestingMacros/Support/DiagnosticMessage+Diagnosing.swift
@@ -88,37 +88,104 @@ func diagnoseIssuesWithTags(in traitExprs: [ExprSyntax], addedTo attribute: Attr
   }
 }
 
-#if canImport(SwiftSyntax600)
-/// Diagnose issues with the lexical context containing a declaration.
+/// Diagnose issues with a synthesized suite (one without an `@Suite` attribute)
+/// containing a declaration.
 ///
 /// - Parameters:
+///   - lexicalContext: The single lexical context to inspect.
 ///   - decl: The declaration to inspect.
-///   - testAttribute: The `@Test` attribute applied to `decl`.
-///   - context: The macro context in which the expression is being parsed.
+///   - attribute: The `@Test` or `@Suite` attribute applied to `decl`.
 ///
 /// - Returns: An array of zero or more diagnostic messages related to the
 ///   lexical context containing `decl`.
+///
+/// This function is also used by ``SuiteDeclarationMacro`` for a number of its
+/// own diagnostics. The implementation substitutes different diagnostic
+/// messages when `suiteDecl` and `decl` are the same syntax node on the
+/// assumption that a suite is self-diagnosing.
 func diagnoseIssuesWithLexicalContext(
+  _ lexicalContext: some SyntaxProtocol,
   containing decl: some DeclSyntaxProtocol,
-  attribute: AttributeSyntax,
-  in context: some MacroExpansionContext
+  attribute: AttributeSyntax
 ) -> [DiagnosticMessage] {
   var diagnostics = [DiagnosticMessage]()
 
-  for lexicalContext in context.lexicalContext {
-    if !lexicalContext.isProtocol((any DeclGroupSyntax).self) {
+  // Functions, closures, etc. are not supported as enclosing lexical contexts.
+  guard let lexicalContext = lexicalContext.asProtocol((any DeclGroupSyntax).self) else {
+    if Syntax(lexicalContext) == Syntax(decl) {
+      diagnostics.append(.attributeNotSupported(attribute, on: lexicalContext))
+    } else {
       diagnostics.append(.containingNodeUnsupported(lexicalContext, whenUsing: attribute, on: decl))
     }
+    return diagnostics
+  }
+
+  // Generic suites are not supported.
+  if let genericClause = lexicalContext.asProtocol((any WithGenericParametersSyntax).self)?.genericParameterClause {
+    diagnostics.append(.genericDeclarationNotSupported(decl, whenUsing: attribute, becauseOf: genericClause))
+  } else if let whereClause = lexicalContext.genericWhereClause {
+    diagnostics.append(.genericDeclarationNotSupported(decl, whenUsing: attribute, becauseOf: whereClause))
+  }
 
-    if let classDecl = lexicalContext.as(ClassDeclSyntax.self) {
-      if !classDecl.modifiers.lazy.map(\.name.tokenKind).contains(.keyword(.final)) {
+  // Suites that are classes must be final.
+  if let classDecl = lexicalContext.as(ClassDeclSyntax.self) {
+    if !classDecl.modifiers.lazy.map(\.name.tokenKind).contains(.keyword(.final)) {
+      if Syntax(classDecl) == Syntax(decl) {
+        diagnostics.append(.nonFinalClassNotSupported(classDecl, whenUsing: attribute))
+      } else {
         diagnostics.append(.containingNodeUnsupported(classDecl, whenUsing: attribute, on: decl))
       }
-    } else if let protocolDecl = lexicalContext.as(ProtocolDeclSyntax.self) {
+    }
+  }
+
+  // Suites cannot be protocols (there's nowhere to put most of the
+  // declarations we generate.)
+  if let protocolDecl = lexicalContext.as(ProtocolDeclSyntax.self) {
+    if Syntax(protocolDecl) == Syntax(decl) {
+      diagnostics.append(.attributeNotSupported(attribute, on: protocolDecl))
+    } else {
       diagnostics.append(.containingNodeUnsupported(protocolDecl, whenUsing: attribute, on: decl))
     }
   }
 
+  // Check other attributes on the declaration. Note that it should be
+  // impossible to reach this point if the declaration can't have attributes.
+  if let attributedDecl = lexicalContext.asProtocol((any WithAttributesSyntax).self) {
+    // Availability is not supported on suites (we need semantic availability
+    // to correctly understand the availability of a suite.)
+    let availabilityAttributes = attributedDecl.availabilityAttributes
+    if !availabilityAttributes.isEmpty {
+      // Diagnose all @available attributes.
+      for availabilityAttribute in availabilityAttributes {
+        diagnostics.append(.availabilityAttributeNotSupported(availabilityAttribute, on: decl, whenUsing: attribute))
+      }
+    } else if let noasyncAttribute = attributedDecl.noasyncAttribute {
+      // No @available attributes, but we do have an @_unavailableFromAsync
+      // attribute and we still need to diagnose that.
+      diagnostics.append(.availabilityAttributeNotSupported(noasyncAttribute, on: decl, whenUsing: attribute))
+    }
+  }
+
   return diagnostics
 }
+
+#if canImport(SwiftSyntax600)
+/// Diagnose issues with the lexical context containing a declaration.
+///
+/// - Parameters:
+///   - decl: The declaration to inspect.
+///   - attribute: The `@Test` or `@Suite` attribute applied to `decl`.
+///   - context: The macro context in which the expression is being parsed.
+///
+/// - Returns: An array of zero or more diagnostic messages related to the
+///   lexical context containing `decl`.
+func diagnoseIssuesWithLexicalContext(
+  _ lexicalContext: [Syntax],
+  containing decl: some DeclSyntaxProtocol,
+  attribute: AttributeSyntax
+) -> [DiagnosticMessage] {
+  lexicalContext.lazy
+    .map { diagnoseIssuesWithLexicalContext($0, containing: decl, attribute: attribute) }
+    .reduce(into: [], +=)
+}
 #endif
diff --git a/Sources/TestingMacros/Support/DiagnosticMessage.swift b/Sources/TestingMacros/Support/DiagnosticMessage.swift
@@ -188,8 +188,11 @@ struct DiagnosticMessage: SwiftDiagnostics.DiagnosticMessage {
   ///
   /// - Returns: A diagnostic message.
   static func genericDeclarationNotSupported(_ decl: some SyntaxProtocol, whenUsing attribute: AttributeSyntax, becauseOf genericClause: some SyntaxProtocol) -> Self {
-    Self(
-      syntax: Syntax(genericClause),
+    // Avoid using a syntax node from a lexical context (it won't have source
+    // location information.)
+    let syntax = (genericClause.root != decl.root) ? Syntax(decl) : Syntax(genericClause)
+    return Self(
+      syntax: syntax,
       message: "Attribute \(_macroName(attribute)) cannot be applied to a generic \(_kindString(for: decl))",
       severity: .error
     )
@@ -209,8 +212,11 @@ struct DiagnosticMessage: SwiftDiagnostics.DiagnosticMessage {
   ///   semantic availability and fully-qualified names for types at macro
   ///   expansion time. ([104081994](rdar://104081994))
   static func availabilityAttributeNotSupported(_ availabilityAttribute: AttributeSyntax, on decl: some SyntaxProtocol, whenUsing attribute: AttributeSyntax) -> Self {
-    Self(
-      syntax: Syntax(availabilityAttribute),
+    // Avoid using a syntax node from a lexical context (it won't have source
+    // location information.)
+    let syntax = (availabilityAttribute.root != decl.root) ? Syntax(decl) : Syntax(availabilityAttribute)
+    return Self(
+      syntax: syntax,
       message: "Attribute \(_macroName(attribute)) cannot be applied to this \(_kindString(for: decl)) because it has been marked '\(availabilityAttribute.trimmed)'",
       severity: .error
     )
@@ -331,7 +337,6 @@ struct DiagnosticMessage: SwiftDiagnostics.DiagnosticMessage {
     )
   }
 
-#if canImport(SwiftSyntax600)
   /// Create a diagnostic message stating that the given attribute cannot be
   /// used within a lexical context.
   ///
@@ -375,6 +380,7 @@ struct DiagnosticMessage: SwiftDiagnostics.DiagnosticMessage {
     }
   }
 
+#if canImport(SwiftSyntax600)
   /// Create a diagnostic message stating that the given attribute cannot be
   /// applied to the given declaration outside the scope of an extension to
   /// `Tag`.
diff --git a/Sources/TestingMacros/TestDeclarationMacro.swift b/Sources/TestingMacros/TestDeclarationMacro.swift
@@ -58,11 +58,11 @@ public struct TestDeclarationMacro: PeerMacro, Sendable {
 
 #if canImport(SwiftSyntax600)
     // Check if the lexical context is appropriate for a suite or test.
-    diagnostics += diagnoseIssuesWithLexicalContext(containing: declaration, attribute: testAttribute, in: context)
+    diagnostics += diagnoseIssuesWithLexicalContext(context.lexicalContext, containing: declaration, attribute: testAttribute)
 #endif
 
     // Only one @Test attribute is supported.
-    let suiteAttributes = function.attributes(named: "Test", in: context)
+    let suiteAttributes = function.attributes(named: "Test")
     if suiteAttributes.count > 1 {
       diagnostics.append(.multipleAttributesNotSupported(suiteAttributes, on: declaration))
     }
@@ -290,7 +290,7 @@ public struct TestDeclarationMacro: PeerMacro, Sendable {
     // If the function is noasync *and* main-actor-isolated, we'll call through
     // MainActor.run to invoke it. We do not have a general mechanism for
     // detecting isolation to other global actors.
-    lazy var isMainActorIsolated = !functionDecl.attributes(named: "MainActor", inModuleNamed: "Swift", in: context).isEmpty
+    lazy var isMainActorIsolated = !functionDecl.attributes(named: "MainActor", inModuleNamed: "Swift").isEmpty
     var forwardCall: (ExprSyntax) -> ExprSyntax = {
       "try await (\($0), Testing.__requiringTry, Testing.__requiringAwait).0"
     }
diff --git a/Tests/TestingMacrosTests/TestDeclarationMacroTests.swift b/Tests/TestingMacrosTests/TestDeclarationMacroTests.swift
@@ -123,6 +123,10 @@ struct TestDeclarationMacroTests {
         "Attribute 'Test' cannot be applied to a function within function 'f()'",
       "struct S { func f(x: Int) { @Suite struct S { } } }":
         "Attribute 'Suite' cannot be applied to a structure within function 'f(x:)'",
+      "struct S<T> { @Test func f() {} }":
+        "Attribute 'Test' cannot be applied to a generic function",
+      "struct S<T> { @Suite struct S {} }":
+        "Attribute 'Suite' cannot be applied to a generic structure",
       "class C { @Test func f() {} }":
         "Attribute 'Test' cannot be applied to a function within non-final class 'C'",
       "class C { @Suite struct S {} }":
@@ -135,6 +139,10 @@ struct TestDeclarationMacroTests {
         "Attribute 'Test' cannot be applied to a function within a closure",
       "{ _ in @Suite struct S {} }":
         "Attribute 'Suite' cannot be applied to a structure within a closure",
+      "@available(*, noasync) struct S { @Test func f() {} }":
+        "Attribute 'Test' cannot be applied to this function because it has been marked '@available(*, noasync)'",
+      "@available(*, noasync) struct S { @Suite struct S {} }":
+        "Attribute 'Suite' cannot be applied to this structure because it has been marked '@available(*, noasync)'",
     ]
   )
   func invalidLexicalContext(input: String, expectedMessage: String) throws {
