From 0098422b89be902cd7b0076dc1846b0f114ae9c1 Mon Sep 17 00:00:00 2001
From: Stuart Montgomery <smontgomery@apple.com>
Date: Tue, 10 Dec 2024 13:41:33 -0600
Subject: [PATCH] Simplify the logic for parsing and handling test arguments in
 the @Test macro (#834)

This is a minor refactor of the logic and data structures in the `@Test`
macro implementation responsible for parsing and handling test function
argument expressions.

### Motivation:

I discovered there was a way to simplify this logic while working on
#808, and since this isn't strictly pertinent to that PR I wanted to
extract it and land it separately here.

### Modifications:

- Remove the need to hardcode the parameter label `arguments:` in three
places.
- Remove the need to iterate a Collection using raw indices.
- Resolve one "TODO" comment.

### Checklist:

- [x] Code and documentation should follow the style of the [Style
Guide](https://github.com/apple/swift-testing/blob/main/Documentation/StyleGuide.md).
- [x] If public symbols are renamed or modified, DocC references should
be updated.
---
 .../Support/AttributeDiscovery.swift          | 53 +++++++++----------
 .../DiagnosticMessage+Diagnosing.swift        |  5 +-
 .../TestingMacros/TestDeclarationMacro.swift  |  3 +-
 3 files changed, 27 insertions(+), 34 deletions(-)

diff --git a/Sources/TestingMacros/Support/AttributeDiscovery.swift b/Sources/TestingMacros/Support/AttributeDiscovery.swift
index f1cfd665f..77b2b174e 100644
--- a/Sources/TestingMacros/Support/AttributeDiscovery.swift
+++ b/Sources/TestingMacros/Support/AttributeDiscovery.swift
@@ -66,17 +66,19 @@ struct AttributeInfo {
   /// The traits applied to the attribute, if any.
   var traits = [ExprSyntax]()
 
+  /// Test arguments passed to a parameterized test function, if any.
+  ///
+  /// When non-`nil`, the value of this property is an array beginning with the
+  /// argument passed to this attribute for the parameter labeled `arguments:`
+  /// followed by all of the remaining, unlabeled arguments.
+  var testFunctionArguments: [Argument]?
+
   /// Whether or not this attribute specifies arguments to the associated test
   /// function.
   var hasFunctionArguments: Bool {
-    otherArguments.lazy
-      .compactMap(\.label?.tokenKind)
-      .contains(.identifier("arguments"))
+    testFunctionArguments != nil
   }
 
-  /// Additional arguments passed to the attribute, if any.
-  var otherArguments = [Argument]()
-
   /// The source location of the attribute.
   ///
   /// When parsing, the testing library uses the start of the attribute's name
@@ -98,6 +100,7 @@ struct AttributeInfo {
   init(byParsing attribute: AttributeSyntax, on declaration: some SyntaxProtocol, in context: some MacroExpansionContext) {
     self.attribute = attribute
 
+    var nonDisplayNameArguments: [Argument] = []
     if let arguments = attribute.arguments, case let .argumentList(argumentList) = arguments {
       // If the first argument is an unlabelled string literal, it's the display
       // name of the test or suite. If it's anything else, including a nil
@@ -106,11 +109,11 @@ struct AttributeInfo {
         let firstArgumentHasLabel = (firstArgument.label != nil)
         if !firstArgumentHasLabel, let stringLiteral = firstArgument.expression.as(StringLiteralExprSyntax.self) {
           displayName = stringLiteral
-          otherArguments = argumentList.dropFirst().map(Argument.init)
+          nonDisplayNameArguments = argumentList.dropFirst().map(Argument.init)
         } else if !firstArgumentHasLabel, firstArgument.expression.is(NilLiteralExprSyntax.self) {
-          otherArguments = argumentList.dropFirst().map(Argument.init)
+          nonDisplayNameArguments = argumentList.dropFirst().map(Argument.init)
         } else {
-          otherArguments = argumentList.map(Argument.init)
+          nonDisplayNameArguments = argumentList.map(Argument.init)
         }
       }
     }
@@ -119,7 +122,7 @@ struct AttributeInfo {
     // See _SelfRemover for more information. Rewriting a syntax tree discards
     // location information from the copy, so only invoke the rewriter if the
     // `Self` keyword is present somewhere.
-    otherArguments = otherArguments.map { argument in
+    nonDisplayNameArguments = nonDisplayNameArguments.map { argument in
       var expr = argument.expression
       if argument.expression.tokens(viewMode: .sourceAccurate).map(\.tokenKind).contains(.keyword(.Self)) {
         let selfRemover = _SelfRemover(in: context)
@@ -131,15 +134,14 @@ struct AttributeInfo {
     // Look for any traits in the remaining arguments and slice them off. Traits
     // are the remaining unlabelled arguments. The first labelled argument (if
     // present) is the start of subsequent context-specific arguments.
-    if !otherArguments.isEmpty {
-      if let labelledArgumentIndex = otherArguments.firstIndex(where: { $0.label != nil }) {
+    if !nonDisplayNameArguments.isEmpty {
+      if let labelledArgumentIndex = nonDisplayNameArguments.firstIndex(where: { $0.label != nil }) {
         // There is an argument with a label, so splice there.
-        traits = otherArguments[otherArguments.startIndex ..< labelledArgumentIndex].map(\.expression)
-        otherArguments = Array(otherArguments[labelledArgumentIndex...])
+        traits = nonDisplayNameArguments[..<labelledArgumentIndex].map(\.expression)
+        testFunctionArguments = Array(nonDisplayNameArguments[labelledArgumentIndex...])
       } else {
         // No argument has a label, so all the remaining arguments are traits.
-        traits = otherArguments.map(\.expression)
-        otherArguments.removeAll(keepingCapacity: false)
+        traits = nonDisplayNameArguments.map(\.expression)
       }
     }
 
@@ -178,21 +180,16 @@ struct AttributeInfo {
       }
     }))
 
-    // Any arguments of the test declaration macro which specify test arguments
-    // need to be wrapped a closure so they may be evaluated lazily by the
-    // testing library at runtime. If any such arguments are present, they will
-    // begin with a labeled argument named `arguments:` and include all
-    // subsequent unlabeled arguments.
-    var otherArguments = self.otherArguments
-    if let argumentsIndex = otherArguments.firstIndex(where: { $0.label?.tokenKind == .identifier("arguments") }) {
-      for index in argumentsIndex ..< otherArguments.endIndex {
-        var argument = otherArguments[index]
-        argument.expression = .init(ClosureExprSyntax { argument.expression.trimmed })
-        otherArguments[index] = argument
+    // If there are any parameterized test function arguments, wrap each in a
+    // closure so they may be evaluated lazily at runtime.
+    if let testFunctionArguments {
+      arguments += testFunctionArguments.map { argument in
+        var copy = argument
+        copy.expression = .init(ClosureExprSyntax { argument.expression.trimmed })
+        return copy
       }
     }
 
-    arguments += otherArguments
     arguments.append(Argument(label: "sourceLocation", expression: sourceLocation))
 
     return LabeledExprListSyntax(arguments)
diff --git a/Sources/TestingMacros/Support/DiagnosticMessage+Diagnosing.swift b/Sources/TestingMacros/Support/DiagnosticMessage+Diagnosing.swift
index e0bd906c5..e6682dc8f 100644
--- a/Sources/TestingMacros/Support/DiagnosticMessage+Diagnosing.swift
+++ b/Sources/TestingMacros/Support/DiagnosticMessage+Diagnosing.swift
@@ -142,10 +142,7 @@ private func _diagnoseIssuesWithParallelizationTrait(_ traitExpr: MemberAccessEx
     return
   }
 
-  let hasArguments = attributeInfo.otherArguments.lazy
-    .compactMap(\.label?.textWithoutBackticks)
-    .contains("arguments")
-  if !hasArguments {
+  if !attributeInfo.hasFunctionArguments {
     // Serializing a non-parameterized test function has no effect.
     context.diagnose(.traitHasNoEffect(traitExpr, in: attributeInfo.attribute))
   }
diff --git a/Sources/TestingMacros/TestDeclarationMacro.swift b/Sources/TestingMacros/TestDeclarationMacro.swift
index 0726aa099..463412d2a 100644
--- a/Sources/TestingMacros/TestDeclarationMacro.swift
+++ b/Sources/TestingMacros/TestDeclarationMacro.swift
@@ -425,9 +425,8 @@ public struct TestDeclarationMacro: PeerMacro, Sendable {
       // case the availability checks fail below.
       let unavailableTestName = context.makeUniqueName(thunking: functionDecl)
 
-      // TODO: don't assume otherArguments is only parameterized function arguments
       var attributeInfo = attributeInfo
-      attributeInfo.otherArguments = []
+      attributeInfo.testFunctionArguments = nil
       result.append(
         """
         @available(*, deprecated, message: "This property is an implementation detail of the testing library. Do not use it directly.")