Handle try and await in unfolded sequences

We can say that any `try`/`await` element also covers all elements to
the right of it in an unfolded sequence. Cases where this isn't true
will be rejected by the compiler, e.g:

```
0 * try foo() + bar()
_ = try foo() ~~~ bar() // Assuming `~~~` has lower precedence than `=`
```

rdar://109470248

Author: hamishknight

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

@@ -106,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

@@ -567,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]
+    )
+  }
 }