Add unsafe as a lexical context too

Match the behavior of `try` and `await`.

Author: hamishknight

Sources/SwiftSyntaxMacros/Syntax+LexicalContext.swift

@@ -67,11 +67,11 @@ 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.
+    // `try`, `await`, and `unsafe` are preserved: A freestanding expression
+    // macro 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 sub-expression is replaced with a trivial
+    // placeholder, though.
     case var tryExpr as TryExprSyntax:
       tryExpr = tryExpr.detached
       tryExpr.expression = ExprSyntax(TypeExprSyntax(type: IdentifierTypeSyntax(name: .wildcardToken())))
@@ -80,6 +80,10 @@ extension SyntaxProtocol {
       awaitExpr = awaitExpr.detached
       awaitExpr.expression = ExprSyntax(TypeExprSyntax(type: IdentifierTypeSyntax(name: .wildcardToken())))
       return Syntax(awaitExpr)
+    case var unsafeExpr as UnsafeExprSyntax:
+      unsafeExpr = unsafeExpr.detached
+      unsafeExpr.expression = ExprSyntax(TypeExprSyntax(type: IdentifierTypeSyntax(name: .wildcardToken())))
+      return Syntax(unsafeExpr)
 
     default:
       return nil
@@ -106,6 +110,12 @@ extension SyntaxProtocol {
       if let parentContext = parentNode.asMacroLexicalContext() {
         parentContexts.append(parentContext)
       }
+      // Unfolded sequence expressions require special handling - effect marker
+      // nodes like `try`, `await`, and `unsafe` are treated as lexical contexts
+      // for all the nodes on their right. Cases where they don't end up
+      // covering nodes to their right in the folded tree are invalid and will
+      // be diagnosed by the compiler. This matches the compiler's ASTScope
+      // handling logic.
       if let sequence = parentNode.as(SequenceExprSyntax.self) {
         var sequenceExprContexts: [Syntax] = []
         for elt in sequence.elements {
@@ -129,9 +139,7 @@ extension SyntaxProtocol {
               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.
+              sequenceExprContexts.append(unsafeElt.asMacroLexicalContext()!)
               elt = unsafeElt.expression
               continue
             }

Tests/SwiftSyntaxMacroExpansionTest/LexicalContextTests.swift

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