Introduce try/await/unsafe macro lexical contexts with unfolded sequence handling

Sources/SwiftOperators/OperatorTable+Folding.swift

@@ -155,6 +155,8 @@ extension OperatorTable {
         )
       )
     }
+    // NOTE: If you add a new try/await/unsafe-like hoisting case here, make
+    // sure to also update `allMacroLexicalContexts` to handle it.
 
     // The form of the binary operation depends on the operator itself,
     // which will be one of the unresolved infix operators.

Sources/SwiftSyntaxMacros/Syntax+LexicalContext.swift

@@ -67,6 +67,20 @@ extension SyntaxProtocol {
     case let freestandingMacro as FreestandingMacroExpansionSyntax:
       return Syntax(freestandingMacro.detached) as Syntax
 
+    // Try and await are preserved: A freestanding expression macro preceded
+    // by try or await may need to know whether those keywords are present so it
+    // can propagate them to any expressions in its expansion which were passed
+    // as arguments to the macro. The expression of the try or await is replaced
+    // with a trivial placeholder, though.
+    case var tryExpr as TryExprSyntax:
+      tryExpr = tryExpr.detached
+      tryExpr.expression = ExprSyntax(TypeExprSyntax(type: IdentifierTypeSyntax(name: .wildcardToken())))
+      return Syntax(tryExpr)
+    case var awaitExpr as AwaitExprSyntax:
+      awaitExpr = awaitExpr.detached
+      awaitExpr.expression = ExprSyntax(TypeExprSyntax(type: IdentifierTypeSyntax(name: .wildcardToken())))
+      return Syntax(awaitExpr)
+
     default:
       return nil
     }
@@ -92,6 +106,39 @@ extension SyntaxProtocol {
       if let parentContext = parentNode.asMacroLexicalContext() {
         parentContexts.append(parentContext)
       }
+      if let sequence = parentNode.as(SequenceExprSyntax.self) {
+        var sequenceExprContexts: [Syntax] = []
+        for elt in sequence.elements {
+          if elt.range.contains(self.position) {
+            // `sequenceExprContexts` is built from the top-down, but we
+            // build the rest of the contexts bottom-up. Reverse for
+            // consistency.
+            parentContexts += sequenceExprContexts.reversed()
+            break
+          }
+          var elt = elt
+          while true {
+            if let tryElt = elt.as(TryExprSyntax.self) {
+              sequenceExprContexts.append(tryElt.asMacroLexicalContext()!)
+              elt = tryElt.expression
+              continue
+            }
+            if let awaitElt = elt.as(AwaitExprSyntax.self) {
+              sequenceExprContexts.append(awaitElt.asMacroLexicalContext()!)
+              elt = awaitElt.expression
+              continue
+            }
+            if let unsafeElt = elt.as(UnsafeExprSyntax.self) {
+              // No scope for this currently, but we need to look through it
+              // since it's similar to 'try' in that it's hoisted above a
+              // binary operator when appearing on the LHS.
+              elt = unsafeElt.expression
+              continue
+            }
+            break
+          }
+        }
+      }
 
       currentNode = parentNode
     }

Tests/SwiftSyntaxMacroExpansionTest/LexicalContextTests.swift

@@ -531,7 +531,7 @@ final class LexicalContextTests: XCTestCase {
         struct S {
           let arg: C
           var contextDescription: String {
-            #lexicalContextDescription
+            try await #lexicalContextDescription
           }
         }
         return S(arg: c)
@@ -542,7 +542,9 @@ final class LexicalContextTests: XCTestCase {
           struct S {
             let arg: C
             var contextDescription: String {
-              """
+              try await """
+              await _
+              try _
               contextDescription: String
               struct S {}
               { c in
@@ -551,7 +553,7 @@ final class LexicalContextTests: XCTestCase {
                 struct S {
                   let arg: C
                   var contextDescription: String {
-                    #lexicalContextDescription
+                    try await #lexicalContextDescription
                   }
                 }
                 return S(arg: c)
@@ -565,4 +567,165 @@ final class LexicalContextTests: XCTestCase {
       macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
     )
   }
+
+  func testTryAwaitInSequenceLexicalContext() {
+    // Valid cases.
+    assertMacroExpansion(
+      "try await #lexicalContextDescription + #lexicalContextDescription",
+      expandedSource: #"""
+        try await """
+        await _
+        try _
+        """ + """
+        await _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    assertMacroExpansion(
+      "try await 0 + 1 + foo(#lexicalContextDescription) + 2",
+      expandedSource: #"""
+        try await 0 + 1 + foo("""
+        await _
+        try _
+        """) + 2
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    assertMacroExpansion(
+      "x = try await 0 + 1 + foo(#lexicalContextDescription) + 2",
+      expandedSource: #"""
+        x = try await 0 + 1 + foo("""
+        await _
+        try _
+        """) + 2
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    // `try await` in the 'then' branch doesn't cover condition or else.
+    assertMacroExpansion(
+      "#lexicalContextDescription ? try await #lexicalContextDescription : #lexicalContextDescription",
+      expandedSource: #"""
+        """
+        """ ? try await """
+        await _
+        try _
+        """ : """
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    // Same for else.
+    assertMacroExpansion(
+      "#lexicalContextDescription ? #lexicalContextDescription : try await #lexicalContextDescription",
+      expandedSource: #"""
+        """
+        """ ? """
+        """ : try await """
+        await _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    // 'try await' in the condition here covers the entire expression
+    assertMacroExpansion(
+      "try await #lexicalContextDescription ? #lexicalContextDescription : #lexicalContextDescription ~~ #lexicalContextDescription",
+      expandedSource: #"""
+        try await """
+        await _
+        try _
+        """ ? """
+        await _
+        try _
+        """ : """
+        await _
+        try _
+        """ ~~ """
+        await _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    assertMacroExpansion(
+      "x = unsafe try try! await 0 + #lexicalContextDescription",
+      expandedSource: #"""
+        x = unsafe try try! await 0 + """
+        await _
+        try! _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+
+    // Invalid cases
+    assertMacroExpansion(
+      "0 + try await #lexicalContextDescription",
+      expandedSource: #"""
+        0 + try await """
+        await _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    // The 'try await' may not actually cover 'lexicalContextDescription' here,
+    // but this will be rejected by the compiler.
+    assertMacroExpansion(
+      "0 + try await 1 ^ #lexicalContextDescription",
+      expandedSource: #"""
+        0 + try await 1 ^ """
+        await _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    // Invalid if '^' has a lower precedence than '='.
+    assertMacroExpansion(
+      "x = try await 0 ^ #lexicalContextDescription",
+      expandedSource: #"""
+        x = try await 0 ^ """
+        await _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    // Unassignable
+    assertMacroExpansion(
+      "#lexicalContextDescription = try await 0 + 1",
+      expandedSource: #"""
+        """
+        """ = try await 0 + 1
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    assertMacroExpansion(
+      "try await #lexicalContextDescription = 0 + #lexicalContextDescription",
+      expandedSource: #"""
+        try await """
+        await _
+        try _
+        """ = 0 + """
+        await _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+    assertMacroExpansion(
+      "try await foo() ? 0 : 1 = #lexicalContextDescription",
+      expandedSource: #"""
+        try await foo() ? 0 : 1 = """
+        await _
+        try _
+        """
+        """#,
+      macros: ["lexicalContextDescription": LexicalContextDescriptionMacro.self]
+    )
+  }
 }