Add unsafe keyword handling to macro expansions.

Package.swift

@@ -96,7 +96,7 @@ let package = Package(
   }(),
 
   dependencies: [
-    .package(url: "https://github.com/swiftlang/swift-syntax.git", from: "601.0.0-latest"),
+    .package(url: "https://github.com/swiftlang/swift-syntax.git", from: "602.0.0-latest"),
   ],
 
   targets: [

Sources/Testing/Test+Macro.swift

@@ -537,6 +537,15 @@ extension Test {
   value
 }
 
+/// A function that abstracts away whether or not the `unsafe` keyword is needed
+/// on an expression.
+///
+/// - Warning: This function is used to implement the `@Test` macro. Do not use
+///   it directly.
+@unsafe @inlinable public func __requiringUnsafe<T>(_ value: consuming T) throws -> T where T: ~Copyable {
+  value
+}
+
 /// The current default isolation context.
 ///
 /// - Warning: This property is used to implement the `@Test` macro. Do not call

Sources/TestingMacros/CMakeLists.txt

@@ -31,7 +31,7 @@ if(SwiftTesting_BuildMacrosAsExecutables)
   set(FETCHCONTENT_BASE_DIR ${CMAKE_BINARY_DIR}/_d)
   FetchContent_Declare(SwiftSyntax
     GIT_REPOSITORY https://github.com/swiftlang/swift-syntax
-    GIT_TAG 1cd35348b089ff8966588742c69727205d99f8ed) # 601.0.0-prerelease-2024-11-18
+    GIT_TAG 340f8400262d494c7c659cd838223990195d7fed) # 602.0.0-prerelease-2025-04-10
   FetchContent_MakeAvailable(SwiftSyntax)
 endif()
 
@@ -101,6 +101,7 @@ target_sources(TestingMacros PRIVATE
   Support/ConditionArgumentParsing.swift
   Support/DiagnosticMessage.swift
   Support/DiagnosticMessage+Diagnosing.swift
+  Support/EffectfulExpressionHandling.swift
   Support/SHA256.swift
   Support/SourceCodeCapturing.swift
   Support/SourceLocationGeneration.swift

Sources/TestingMacros/ConditionMacro.swift

@@ -117,7 +117,6 @@ extension ConditionMacro {
     var checkArguments = [Argument]()
     do {
       if let trailingClosureIndex {
-
         // Include all arguments other than the "comment" and "sourceLocation"
         // arguments here.
         checkArguments += macroArguments.indices.lazy
@@ -458,7 +457,7 @@ extension ExitTestConditionMacro {
     decls.append(
       """
       @Sendable func \(bodyThunkName)() async throws -> Swift.Void {
-        return try await Testing.__requiringTry(Testing.__requiringAwait(\(bodyArgumentExpr.trimmed)))()
+        return \(applyEffectfulKeywords([.try, .await, .unsafe], to: bodyArgumentExpr))()
       }
       """
     )

Sources/TestingMacros/Support/ConditionArgumentParsing.swift

@@ -472,17 +472,6 @@ private func _parseCondition(from expr: ExprSyntax, for macro: some Freestanding
     return _parseCondition(from: closureExpr, for: macro, in: context)
   }
 
-  // If the condition involves the `try` or `await` keywords, assume we cannot
-  // expand it. This check cannot handle expressions like
-  // `try #expect(a.b(c))` where `b()` is throwing because the `try` keyword is
-  // outside the macro expansion. SEE: rdar://109470248
-  let containsTryOrAwait = expr.tokens(viewMode: .sourceAccurate).lazy
-    .map(\.tokenKind)
-    .contains { $0 == .keyword(.try) || $0 == .keyword(.await) }
-  if containsTryOrAwait {
-    return Condition(expression: expr)
-  }
-
   if let infixOperator = expr.as(InfixOperatorExprSyntax.self),
      let op = infixOperator.operator.as(BinaryOperatorExprSyntax.self) {
     return _parseCondition(from: expr, leftOperand: infixOperator.leftOperand, operator: op, rightOperand: infixOperator.rightOperand, for: macro, in: context)
@@ -527,6 +516,15 @@ private func _parseCondition(from expr: ExprSyntax, for macro: some Freestanding
 ///
 /// - Returns: An instance of ``Condition`` describing `expr`.
 func parseCondition(from expr: ExprSyntax, for macro: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext) -> Condition {
+  // If the condition involves the `unsafe`, `try`, or `await` keywords, assume
+  // we cannot expand it. This check cannot handle expressions like
+  // `try #expect(a.b(c))` where `b()` is throwing because the `try` keyword is
+  // outside the macro expansion. SEE: rdar://109470248
+  let effectKeywordsToApply = findEffectKeywords(in: expr, context: context)
+  guard effectKeywordsToApply.intersection([.unsafe, .try, .await]).isEmpty else {
+    return Condition(expression: expr)
+  }
+
   _diagnoseTrivialBooleanValue(from: expr, for: macro, in: context)
   let result = _parseCondition(from: expr, for: macro, in: context)
   return result

Sources/TestingMacros/Support/EffectfulExpressionHandling.swift

@@ -0,0 +1,159 @@
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for Swift project authors
+//
+
+import SwiftSyntax
+import SwiftSyntaxBuilder
+import SwiftSyntaxMacros
+
+// MARK: - Finding effect keywords
+
+/// A syntax visitor class that looks for effectful keywords in a given
+/// expression.
+private final class _EffectFinder: SyntaxAnyVisitor {
+  /// The effect keywords discovered so far.
+  var effectKeywords: Set<Keyword> = []
+
+  override func visitAny(_ node: Syntax) -> SyntaxVisitorContinueKind {
+    switch node.kind {
+    case .tryExpr:
+      effectKeywords.insert(.try)
+    case .awaitExpr:
+      effectKeywords.insert(.await)
+    case .consumeExpr:
+      effectKeywords.insert(.consume)
+    case .borrowExpr:
+      effectKeywords.insert(.borrow)
+    case .unsafeExpr:
+      effectKeywords.insert(.unsafe)
+    case .closureExpr, .functionDecl:
+      // Do not delve into closures or function declarations.
+      return .skipChildren
+    case .variableDecl:
+      // Delve into variable declarations.
+      return .visitChildren
+    default:
+      // Do not delve into declarations other than variables.
+      if node.isProtocol((any DeclSyntaxProtocol).self) {
+        return .skipChildren
+      }
+    }
+
+    // Recurse into everything else.
+    return .visitChildren
+  }
+}
+
+/// Find effectful keywords in a syntax node.
+///
+/// - Parameters:
+///   - node: The node to inspect.
+///   - context: The macro context in which the expression is being parsed.
+///
+/// - Returns: A set of effectful keywords such as `await` that are present in
+///   `node`.
+///
+/// This function does not descend into function declarations or closure
+/// expressions because they represent distinct lexical contexts and their
+/// effects are uninteresting in the context of `node` unless they are called.
+func findEffectKeywords(in node: some SyntaxProtocol, context: some MacroExpansionContext) -> Set<Keyword> {
+  // TODO: gather any effects from the lexical context once swift-syntax-#3037 and related PRs land
+  let effectFinder = _EffectFinder(viewMode: .sourceAccurate)
+  effectFinder.walk(node)
+  return effectFinder.effectKeywords
+}
+
+// MARK: - Inserting effect keywords/thunks
+
+/// Make a function call expression to an effectful thunk function provided by
+/// the testing library.
+///
+/// - Parameters:
+///   - thunkName: The unqualified name of the thunk function to call. This
+///     token must be the name of a function in the `Testing` module.
+///   - expr: The expression to thunk.
+///
+/// - Returns: An expression representing a call to the function named
+///   `thunkName`, passing `expr`.
+private func _makeCallToEffectfulThunk(_ thunkName: TokenSyntax, passing expr: some ExprSyntaxProtocol) -> ExprSyntax {
+  ExprSyntax(
+    FunctionCallExprSyntax(
+      calledExpression: MemberAccessExprSyntax(
+        base: DeclReferenceExprSyntax(baseName: .identifier("Testing")),
+        declName: DeclReferenceExprSyntax(baseName: thunkName)
+      ),
+      leftParen: .leftParenToken(),
+      rightParen: .rightParenToken()
+    ) {
+      LabeledExprSyntax(expression: expr.trimmed)
+    }
+  )
+}
+
+/// Apply the given effectful keywords (i.e. `try` and `await`) to an expression
+/// using thunk functions provided by the testing library.
+///
+/// - Parameters:
+///   - effectfulKeywords: The effectful keywords to apply.
+///   - expr: The expression to apply the keywords and thunk functions to.
+///
+/// - Returns: A copy of `expr` if no changes are needed, or an expression that
+///   adds the keywords in `effectfulKeywords` to `expr`.
+func applyEffectfulKeywords(_ effectfulKeywords: Set<Keyword>, to expr: some ExprSyntaxProtocol) -> ExprSyntax {
+  let originalExpr = expr
+  var expr = ExprSyntax(expr)
+
+  let needAwait = effectfulKeywords.contains(.await) && !expr.is(AwaitExprSyntax.self)
+  let needTry = effectfulKeywords.contains(.try) && !expr.is(TryExprSyntax.self)
+  let needUnsafe = effectfulKeywords.contains(.unsafe) && !expr.is(UnsafeExprSyntax.self)
+
+  // First, add thunk function calls.
+  if needAwait {
+    expr = _makeCallToEffectfulThunk(.identifier("__requiringAwait"), passing: expr)
+  }
+  if needTry {
+    expr = _makeCallToEffectfulThunk(.identifier("__requiringTry"), passing: expr)
+  }
+  if needUnsafe {
+    expr = _makeCallToEffectfulThunk(.identifier("__requiringUnsafe"), passing: expr)
+  }
+
+  // Then add keyword expressions. (We do this separately so we end up writing
+  // `try await __r(__r(self))` instead of `try __r(await __r(self))` which is
+  // less accepted by the compiler.)
+  if needAwait {
+    expr = ExprSyntax(
+      AwaitExprSyntax(
+        awaitKeyword: .keyword(.await).with(\.trailingTrivia, .space),
+        expression: expr
+      )
+    )
+  }
+  if needTry {
+    expr = ExprSyntax(
+      TryExprSyntax(
+        tryKeyword: .keyword(.try).with(\.trailingTrivia, .space),
+        expression: expr
+      )
+    )
+  }
+  if needUnsafe {
+    expr = ExprSyntax(
+      UnsafeExprSyntax(
+        unsafeKeyword: .keyword(.unsafe).with(\.trailingTrivia, .space),
+        expression: expr
+      )
+    )
+  }
+
+  expr.leadingTrivia = originalExpr.leadingTrivia
+  expr.trailingTrivia = originalExpr.trailingTrivia
+
+  return expr
+}

Sources/TestingMacros/TestDeclarationMacro.swift

@@ -246,17 +246,17 @@ public struct TestDeclarationMacro: PeerMacro, Sendable {
     // detecting isolation to other global actors.
     lazy var isMainActorIsolated = !functionDecl.attributes(named: "MainActor", inModuleNamed: "_Concurrency").isEmpty
     var forwardCall: (ExprSyntax) -> ExprSyntax = {
-      "try await Testing.__requiringTry(Testing.__requiringAwait(\($0)))"
+      applyEffectfulKeywords([.try, .await, .unsafe], to: $0)
     }
     let forwardInit = forwardCall
     if functionDecl.noasyncAttribute != nil {
       if isMainActorIsolated {
         forwardCall = {
-          "try await MainActor.run { try Testing.__requiringTry(\($0)) }"
+          "try await MainActor.run { \(applyEffectfulKeywords([.try, .unsafe], to: $0)) }"
         }
       } else {
         forwardCall = {
-          "try { try Testing.__requiringTry(\($0)) }()"
+          "try { \(applyEffectfulKeywords([.try, .unsafe], to: $0)) }()"
         }
       }
     }