Hook up effectful keywords discovery from lexical context

Author: grynspan

Sources/TestingMacros/ConditionMacro.swift

@@ -131,6 +131,7 @@ extension ConditionMacro {
     var expandedFunctionName = TokenSyntax.identifier("__checkCondition")
     var checkArguments = [Argument]()
     var effectKeywordsToApply: Set<Keyword> = []
+    var effectKeywordsToApplyOverall: Set<Keyword> = []
     do {
       if let trailingClosureIndex {
         expandedFunctionName = .identifier("__checkClosureCall")
@@ -168,7 +169,10 @@ extension ConditionMacro {
 
       } else if let firstArgument = macroArguments.first {
         let originalArgumentExpr = firstArgument.expression
-        effectKeywordsToApply = findEffectKeywords(in: originalArgumentExpr, context: context)
+        let effectKeywordsFromNode = findEffectKeywords(in: originalArgumentExpr)
+        let effectKeywordsFromLexicalContext = findEffectKeywords(in: context)
+        effectKeywordsToApply = effectKeywordsFromNode.union(effectKeywordsFromLexicalContext)
+        effectKeywordsToApplyOverall = effectKeywordsToApply.subtracting(effectKeywordsFromLexicalContext)
 
         var useEscapeHatch = false
         if let asExpr = originalArgumentExpr.as(AsExprSyntax.self), asExpr.questionOrExclamationMark == nil {
@@ -278,12 +282,8 @@ extension ConditionMacro {
     } else {
       "\(call).__expected()"
     }
-    if effectKeywordsToApply.contains(.await) {
-      call = "await \(call)"
-    }
-    if !isThrowing && effectKeywordsToApply.contains(.try) {
-      call = "try \(call)"
-    }
+
+    call = applyEffectfulKeywords(effectKeywordsToApplyOverall, to: call, insertThunkCalls: false)
     return call
   }
 

Sources/TestingMacros/Support/EffectfulExpressionHandling.swift

@@ -14,61 +14,105 @@ import SwiftSyntaxMacros
 
 // MARK: - Finding effect keywords and expressions
 
+/// Get the effect keyword corresponding to a given syntax node, if any.
+///
+/// - Parameters:
+/// 	- expr: The syntax node that may represent an effectful expression.
+///
+/// - Returns: The effect keyword corresponding to `expr`, if any.
+private func _effectKeyword(for expr: ExprSyntax) -> Keyword? {
+  switch expr.kind {
+  case .tryExpr:
+    return .try
+  case .awaitExpr:
+    return .await
+  case .consumeExpr:
+    return .consume
+  case .borrowExpr:
+    return .borrow
+  case .unsafeExpr:
+    return .unsafe
+  default:
+    return nil
+  }
+}
+
+/// Determine how to descend further into a syntax node tree from a given node.
+///
+/// - Parameters:
+///   - node: The syntax node currently being walked.
+///
+/// - Returns: Whether or not to descend into `node` and visit its children.
+private func _continueKind(for node: Syntax) -> SyntaxVisitorContinueKind {
+  switch node.kind {
+  case .tryExpr, .awaitExpr, .consumeExpr, .borrowExpr, .unsafeExpr:
+    // If this node represents an effectful expression, look inside it for
+    // additional such expressions.
+    return .visitChildren
+  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
+}
+
 /// 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
-      }
+    if let expr = node.as(ExprSyntax.self), let keyword = _effectKeyword(for: expr) {
+      effectKeywords.insert(keyword)
     }
 
-    // Recurse into everything else.
-    return .visitChildren
+    return _continueKind(for: node)
   }
 }
 
 /// 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
+func findEffectKeywords(in node: some SyntaxProtocol) -> Set<Keyword> {
   let effectFinder = _EffectFinder(viewMode: .sourceAccurate)
   effectFinder.walk(node)
   return effectFinder.effectKeywords
 }
 
+/// Find effectful keywords in a macro's lexical context.
+///
+/// - Parameters:
+///   - context: The macro context in which the expression is being parsed.
+///
+/// - Returns: A set of effectful keywords such as `await` that are present in
+///   `context` and would apply to an expression macro during its expansion.
+func findEffectKeywords(in context: some MacroExpansionContext) -> Set<Keyword> {
+  let result = context.lexicalContext.reversed().lazy
+    .prefix { _continueKind(for: $0) == .visitChildren }
+    .compactMap { $0.as(ExprSyntax.self) }
+    .compactMap(_effectKeyword(for:))
+  return Set(result)
+}
+
 extension BidirectionalCollection<Syntax> {
   /// The suffix of syntax nodes in this collection which are effectful
   /// expressions, such as those for `try` or `await`.
@@ -117,10 +161,13 @@ private func _makeCallToEffectfulThunk(_ thunkName: TokenSyntax, passing expr: s
 /// - Parameters:
 ///   - effectfulKeywords: The effectful keywords to apply.
 ///   - expr: The expression to apply the keywords and thunk functions to.
+///   - insertThunkCalls: Whether or not to also insert calls to thunks to
+///   	ensure the inserted keywords do not generate warnings. If you aren't
+///     sure whether thunk calls are needed, pass `true`.
 ///
 /// - 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 {
+func applyEffectfulKeywords(_ effectfulKeywords: Set<Keyword>, to expr: some ExprSyntaxProtocol, insertThunkCalls: Bool = true) -> ExprSyntax {
   let originalExpr = expr
   var expr = ExprSyntax(expr.trimmed)
 
@@ -135,33 +182,35 @@ func applyEffectfulKeywords(_ effectfulKeywords: Set<Keyword>, to expr: some Exp
 #endif
 
   // First, add thunk function calls.
-  if needAwait {
-    expr = _makeCallToEffectfulThunk(.identifier("__requiringAwait"), passing: expr)
-  }
-  if needTry {
-    expr = _makeCallToEffectfulThunk(.identifier("__requiringTry"), passing: expr)
-  }
+  if insertThunkCalls {
+    if needAwait {
+      expr = _makeCallToEffectfulThunk(.identifier("__requiringAwait"), passing: expr)
+    }
+    if needTry {
+      expr = _makeCallToEffectfulThunk(.identifier("__requiringTry"), passing: expr)
+    }
 #if compiler(>=6.2)
-  if needUnsafe {
-    expr = _makeCallToEffectfulThunk(.identifier("__requiringUnsafe"), passing: expr)
-  }
+    if needUnsafe {
+      expr = _makeCallToEffectfulThunk(.identifier("__requiringUnsafe"), passing: expr)
+    }
 #endif
+  }
 
   // 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),
+        awaitKeyword: .keyword(.await, trailingTrivia: .space, presence: .present),
         expression: expr
       )
     )
   }
   if needTry {
     expr = ExprSyntax(
       TryExprSyntax(
-        tryKeyword: .keyword(.try).with(\.trailingTrivia, .space),
+        tryKeyword: .keyword(.try, trailingTrivia: .space, presence: .present),
         expression: expr
       )
     )
@@ -170,7 +219,7 @@ func applyEffectfulKeywords(_ effectfulKeywords: Set<Keyword>, to expr: some Exp
   if needUnsafe {
     expr = ExprSyntax(
       UnsafeExprSyntax(
-        unsafeKeyword: .keyword(.unsafe).with(\.trailingTrivia, .space),
+        unsafeKeyword: .keyword(.unsafe, trailingTrivia: .space, presence: .present),
         expression: expr
       )
     )