Capture x instead of &x as source for inout args, ensure dollar identifiers are rewritten as late as possible

Author: grynspan

Sources/TestingMacros/ConditionMacro.swift

@@ -189,30 +189,71 @@ extension ConditionMacro {
             argumentExpr = "try Testing.__requiringTry(\(argumentExpr))"
           }
 
-          // Replace any dollar identifiers we find.
-          let closureArguments = rewriteClosureArguments(in: argumentExpr)
-          if let closureArguments {
-            argumentExpr = closureArguments.rewrittenNode.cast(ExprSyntax.self)
-          }
-
-          // If we're inserting any additional code into the closure before the
-          // rewritten argument, we can't elide the return keyword for brevity.
-          var returnKeyword: TokenSyntax?
-          if !prefixCodeBlockItems.isEmpty {
-            returnKeyword = .keyword(.return)
-              .with(\.leadingTrivia, argumentExpr.leadingTrivia)
-            argumentExpr.leadingTrivia = .space
+          // Construct the body of the closure that we'll pass to the expanded
+          // function.
+          var codeBlockItems = CodeBlockItemListSyntax {
+            if prefixCodeBlockItems.isEmpty {
+              CodeBlockItemSyntax(item: .expr(argumentExpr))
+                .with(\.trailingTrivia, .newline)
+            } else {
+              prefixCodeBlockItems
+
+              // If we're inserting any additional code into the closure before
+              // the rewritten argument, we can't elide the return keyword.
+              CodeBlockItemSyntax(
+                item: .stmt(
+                  StmtSyntax(
+                    ReturnStmtSyntax(
+                      expression: argumentExpr
+                        .with(\.leadingTrivia, .space)
+                    )
+                  )
+                )
+              ).with(\.trailingTrivia, .newline)
+            }
           }
 
-          // Enclose the expression in a closure into which we pass our local
-          // context object.
-          argumentExpr = """
-          { \(closureArguments?.captureList) (\(expressionContextName): inout Testing.__ExpectationContext) in
-            \(prefixCodeBlockItems)\(returnKeyword)\(argumentExpr)
+          // Replace any dollar identifiers we find.
+          let closureArguments = rewriteClosureArguments(in: codeBlockItems)
+          if let closureArguments {
+            codeBlockItems = closureArguments.rewrittenNode.cast(CodeBlockItemListSyntax.self)
           }
-          """
 
-          checkArguments.append(Argument(expression: argumentExpr))
+          // Enclose the code block in the final closure.
+          let closureExpr = ClosureExprSyntax(
+            signature: ClosureSignatureSyntax(
+              capture: closureArguments?.captureList,
+              parameterClause: .parameterClause(
+                ClosureParameterClauseSyntax(
+                  parameters: ClosureParameterListSyntax {
+                    ClosureParameterSyntax(
+                      firstName: expressionContextName,
+                      colon: .colonToken().with(\.trailingTrivia, .space),
+                      type: TypeSyntax(
+                        AttributedTypeSyntax(
+                          specifiers: [
+                            .init(
+                              SimpleTypeSpecifierSyntax(specifier: .keyword(.inout))
+                                .with(\.trailingTrivia, .space)
+                            )
+                          ],
+                          baseType: MemberTypeSyntax(
+                            baseType: IdentifierTypeSyntax(name: .identifier("Testing")),
+                            name: .identifier("__ExpectationContext")
+                          )
+                        )
+                      )
+                    )
+                  }
+                )
+              ),
+              inKeyword: .keyword(.in)
+                .with(\.leadingTrivia, .space)
+                .with(\.trailingTrivia, .newline)
+            ),
+            statements: codeBlockItems
+          )
+          checkArguments.append(Argument(expression: closureExpr))
 
           // Sort the rewritten nodes. This isn't strictly necessary for
           // correctness but it does make the produced code more consistent.

Sources/TestingMacros/Support/ConditionArgumentParsing.swift

@@ -188,13 +188,16 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
   ///
   /// - Parameters:
   ///   - node: The node to rewrite.
+  ///   - functionName: If not `nil`, the name of the function to call (as a
+  ///     member function of the expression context.)
+  ///   - additionalArguments: Any additional arguments to pass to the function.
   ///
   /// - Returns: A rewritten copy of `node` that calls into the expression
   ///   context when it is evaluated at runtime.
   ///
   /// This function is equivalent to `_rewrite(node, originalWas: node)`.
-  private func _rewrite<E>(_ node: E) -> ExprSyntax where E: ExprSyntaxProtocol {
-    _rewrite(node, originalWas: node)
+  private func _rewrite<E>(_ node: E, calling functionName: TokenSyntax? = nil, passing additionalArguments: [Argument] = []) -> ExprSyntax where E: ExprSyntaxProtocol {
+    _rewrite(node, originalWas: node, calling: functionName, passing: additionalArguments)
   }
 
   /// Whether or not the parent node of the given node is capable of containing
@@ -355,13 +358,11 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
       return _rewrite(
         ClosureExprSyntax {
           InfixOperatorExprSyntax(
-            leftOperand: DeclReferenceExprSyntax(
-              baseName: .dollarIdentifier("$0")
-            ).with(\.trailingTrivia, .space),
+            leftOperand: DeclReferenceExprSyntax(baseName: .dollarIdentifier("$0"))
+              .with(\.trailingTrivia, .space),
             operator: BinaryOperatorExprSyntax(text: op),
-            rightOperand: DeclReferenceExprSyntax(
-              baseName: .dollarIdentifier("$1")
-            ).with(\.leadingTrivia, .space)
+            rightOperand: DeclReferenceExprSyntax(baseName: .dollarIdentifier("$1"))
+              .with(\.leadingTrivia, .space)
           )
         },
         originalWas: node,
@@ -391,11 +392,7 @@ private final class _ContextInserter<C, M>: SyntaxRewriter where C: MacroExpansi
 
     let teardownItem = CodeBlockItemSyntax(
       item: .expr(
-        _rewrite(
-          node.expression,
-          originalWas: node,
-          calling: .identifier("__inoutAfter")
-        )
+        _rewrite(node.expression, calling: .identifier("__inoutAfter"))
       )
     )
     teardownItems.append(teardownItem)
@@ -571,9 +568,7 @@ func insertCalls(
         item: .stmt(
           StmtSyntax(
             DeferStmtSyntax {
-              for teardownItem in contextInserter.teardownItems {
-                teardownItem
-              }
+              contextInserter.teardownItems
             }
           )
         )
@@ -667,6 +662,24 @@ private final class _DollarIdentifierReplacer: SyntaxRewriter {
   /// The dollar identifier tokens that have been rewritten.
   var dollarIdentifierTokens = Set<TokenSyntax>()
 
+  /// The node to treat as the root node when expanding expressions.
+  var effectiveRootNode: Syntax
+
+  init(rootedAt effectiveRootNode: Syntax) {
+    self.effectiveRootNode = effectiveRootNode
+  }
+
+  override func visitAny(_ node: Syntax) -> Syntax? {
+    // Do not recurse into closure expressions (except the root node) because
+    // they will have their own argument/capture lists that won't conflict with
+    // the enclosing scope's.
+    if node.is(ClosureExprSyntax.self) && node != effectiveRootNode {
+      return Syntax(node)
+    }
+
+    return nil
+  }
+
   override func visit(_ node: TokenSyntax) -> TokenSyntax {
     if case let .dollarIdentifier(id) = node.tokenKind, id.dropFirst().allSatisfy(\.isWholeNumber) {
       // This dollar identifier is numeric, so it's a closure argument.
@@ -676,12 +689,6 @@ private final class _DollarIdentifierReplacer: SyntaxRewriter {
 
     return node
   }
-
-  override func visit(_ node: ClosureExprSyntax) -> ExprSyntax {
-    // Do not recurse into closure expressions because they will have their own
-    // argument lists that won't conflict with the enclosing scope's.
-    return ExprSyntax(node)
-  }
 }
 
 /// Rewrite any implicit closure arguments (dollar identifiers such as `$0`) in
@@ -694,7 +701,7 @@ private final class _DollarIdentifierReplacer: SyntaxRewriter {
 ///   can be used to transform the original dollar identifiers to their
 ///   rewritten counterparts in a nested closure invocation.
 func rewriteClosureArguments(in node: some SyntaxProtocol) -> (rewrittenNode: Syntax, captureList: ClosureCaptureClauseSyntax)? {
-  let replacer = _DollarIdentifierReplacer()
+  let replacer = _DollarIdentifierReplacer(rootedAt: Syntax(node))
   let result = replacer.rewrite(node)
   if replacer.dollarIdentifierTokens.isEmpty {
     return nil

Tests/SubexpressionShowcase/SubexpressionShowcase.swift

@@ -30,7 +30,6 @@ struct T {
   static func j(_ d: Double) -> Bool { false }
 }
 
-
 @Test func runSubexpressionShowcase() async {
   await withKnownIssue {
     try await subexpressionShowcase()

Tests/TestingMacrosTests/ConditionMacroTests.swift

@@ -61,11 +61,11 @@ struct ConditionMacroTests {
       ##"#expect(a.b(c, d: e))"##:
         ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in __ec(__ec(a.self, "6").b(__ec(c, "700"), d: __ec(e, "12100")), "") }, sourceCode: ["": "a.b(c, d: e)", "6": "a", "700": "c", "12100": "e"], comments: [], isRequired: false, sourceLocation: Testing.SourceLocation.__here()).__expected()"##,
       ##"#expect(a.b(&c))"##:
-        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "700") } return __ec(__ec(a.self, "6").b(&c), "") }, sourceCode: ["": "a.b(&c)", "6": "a", "700": "&c"], comments: [], isRequired: false, sourceLocation: Testing.SourceLocation.__here()).__expected()"##,
+        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "1700") } return __ec(__ec(a.self, "6").b(&c), "") }, sourceCode: ["": "a.b(&c)", "6": "a", "1700": "c"], comments: [], isRequired: false, sourceLocation: Testing.SourceLocation.__here()).__expected()"##,
       ##"#expect(a.b(&c, &d.e))"##:
-        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "700") __ec.__inoutAfter(d.e, "18100") } return __ec(__ec(a.self, "6").b(&c, &d.e), "") }, sourceCode: ["": "a.b(&c, &d.e)", "6": "a", "700": "&c", "18100": "&d.e"], comments: [], isRequired: false, sourceLocation: Testing.SourceLocation.__here()).__expected()"##,
+        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "1700") __ec.__inoutAfter(d.e, "58100") } return __ec(__ec(a.self, "6").b(&c, &d.e), "") }, sourceCode: ["": "a.b(&c, &d.e)", "6": "a", "1700": "c", "58100": "d.e"], comments: [], isRequired: false, sourceLocation: Testing.SourceLocation.__here()).__expected()"##,
       ##"#expect(a.b(&c, d))"##:
-        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "700") } return  __ec(__ec(a.self, "6").b(&c, __ec(d, "18100")), "") }, sourceCode: ["": "a.b(&c, d)", "6": "a", "700": "&c", "18100": "d"], comments: [], isRequired: false, sourceLocation: Testing.SourceLocation.__here()).__expected()"##,
+        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "1700") } return __ec(__ec(a.self, "6").b(&c, __ec(d, "18100")), "") }, sourceCode: ["": "a.b(&c, d)", "6": "a", "1700": "c", "18100": "d"], comments: [], isRequired: false, sourceLocation: Testing.SourceLocation.__here()).__expected()"##,
       ##"#expect(a.b(try c()))"##:
         ##"try Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in try Testing.__requiringTry(__ec(__ec(a.self, "6").b(try __ec(c(), "1700")), "")) }, sourceCode: ["": "a.b(try c())", "6": "a", "1700": "c()"], comments: [], isRequired: false, sourceLocation: Testing.SourceLocation.__here()).__expected()"##,
       ##"#expect(a?.b(c))"##:
@@ -142,11 +142,11 @@ struct ConditionMacroTests {
       ##"#require(a.b(c, d: e))"##:
         ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in try Testing.__requiringTry(__ec(__ec(a.self, "6").b(__ec(c, "700"), d: __ec(e, "12100")), "")) }, sourceCode: ["": "a.b(c, d: e)", "6": "a", "700": "c", "12100": "e"], comments: [], isRequired: true, sourceLocation: Testing.SourceLocation.__here()).__required()"##,
       ##"#require(a.b(&c))"##:
-        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "700") } return try Testing.__requiringTry(__ec(__ec(a.self, "6").b(&c), "")) }, sourceCode: ["": "a.b(&c)", "6": "a", "700": "&c"], comments: [], isRequired: true, sourceLocation: Testing.SourceLocation.__here()).__required()"##,
+        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "1700") } return try Testing.__requiringTry(__ec(__ec(a.self, "6").b(&c), "")) }, sourceCode: ["": "a.b(&c)", "6": "a", "1700": "c"], comments: [], isRequired: true, sourceLocation: Testing.SourceLocation.__here()).__required()"##,
       ##"#require(a.b(&c, &d.e))"##:
-        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "700") __ec.__inoutAfter(d.e, "18100") } return try Testing.__requiringTry(__ec(__ec(a.self, "6").b(&c, &d.e), "")) }, sourceCode: ["": "a.b(&c, &d.e)", "6": "a", "700": "&c", "18100": "&d.e"], comments: [], isRequired: true, sourceLocation: Testing.SourceLocation.__here()).__required()"##,
+        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "1700") __ec.__inoutAfter(d.e, "58100") } return try Testing.__requiringTry(__ec(__ec(a.self, "6").b(&c, &d.e), "")) }, sourceCode: ["": "a.b(&c, &d.e)", "6": "a", "1700": "c", "58100": "d.e"], comments: [], isRequired: true, sourceLocation: Testing.SourceLocation.__here()).__required()"##,
       ##"#require(a.b(&c, d))"##:
-        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "700") } return try Testing.__requiringTry(__ec(__ec(a.self, "6").b(&c, __ec(d, "18100")), "")) }, sourceCode: ["": "a.b(&c, d)", "6": "a", "700": "&c", "18100": "d"], comments: [], isRequired: true, sourceLocation: Testing.SourceLocation.__here()).__required()"##,
+        ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in defer { __ec.__inoutAfter(c, "1700") } return try Testing.__requiringTry(__ec(__ec(a.self, "6").b(&c, __ec(d, "18100")), "")) }, sourceCode: ["": "a.b(&c, d)", "6": "a", "1700": "c", "18100": "d"], comments: [], isRequired: true, sourceLocation: Testing.SourceLocation.__here()).__required()"##,
       ##"#require(a.b(try c()))"##:
         ##"Testing.__checkCondition({ (__ec: inout Testing.__ExpectationContext) in try Testing.__requiringTry(__ec(__ec(a.self, "6").b(try __ec(c(), "1700")), "")) }, sourceCode: ["": "a.b(try c())", "6": "a", "1700": "c()"], comments: [], isRequired: true, sourceLocation: Testing.SourceLocation.__here()).__required()"##,
       ##"#require(a?.b(c))"##: