Fold operators using the standard operator table in MacroSystem

The compiler folds operators in attributes and freestanding macro nodes but `MacroSystem` wasn’t doing that. But it should to match the compiler behavior.

rdar://114786803
Fixes #2128

Author: ahoppen

Sources/SwiftCompilerPluginMessageHandling/PluginMacroExpansionContext.swift

@@ -85,7 +85,7 @@ class SourceManager {
     case .attribute:
       node = Syntax(AttributeSyntax.parse(from: &parser))
     }
-    if let operatorTable = operatorTable {
+    if let operatorTable {
       node = operatorTable.foldAll(node, errorHandler: { _ in /*ignore*/ })
     }
 

Sources/SwiftSyntaxMacroExpansion/BasicMacroExpansionContext.swift

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import SwiftDiagnostics
+import SwiftOperators
 import SwiftSyntax
 import SwiftSyntaxMacros
 
@@ -77,6 +78,20 @@ extension BasicMacroExpansionContext {
     detachedNodes[Syntax(detached)] = Syntax(node)
     return detached
   }
+
+  /// Fold all operators in `node` and associated the ``KnownSourceFile``
+  /// information of `node` with the original new, folded tree.
+  func foldAllOperators(of node: some SyntaxProtocol, with operatorTable: OperatorTable) -> Syntax {
+    let folded = operatorTable.foldAll(node, errorHandler: { _ in /*ignore*/ })
+    if let originalSourceFile = node.root.as(SourceFileSyntax.self),
+      let newSourceFile = folded.root.as(SourceFileSyntax.self)
+    {
+      // Folding operators doesn't change the source file and its associated locations
+      // Record the `KnownSourceFile` information for the folded tree.
+      sourceFiles[newSourceFile] = sourceFiles[originalSourceFile]
+    }
+    return folded
+  }
 }
 
 extension String {

Sources/SwiftSyntaxMacroExpansion/MacroSystem.swift

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import SwiftDiagnostics
+import SwiftOperators
 import SwiftSyntax
 import SwiftSyntaxBuilder
 @_spi(MacroExpansion) import SwiftParser
@@ -55,7 +56,7 @@ private func expandFreestandingMemberDeclList(
     let expanded = try expandFreestandingMacro(
       definition: definition,
       macroRole: inferFreestandingMacroRole(definition: definition),
-      node: node.detach(in: context),
+      node: node.detach(in: context, foldingWith: .standardOperators),
       in: context,
       indentationWidth: indentationWidth
     )
@@ -80,7 +81,7 @@ private func expandFreestandingCodeItemList(
     let expanded = try expandFreestandingMacro(
       definition: definition,
       macroRole: inferFreestandingMacroRole(definition: definition),
-      node: node.detach(in: context),
+      node: node.detach(in: context, foldingWith: .standardOperators),
       in: context,
       indentationWidth: indentationWidth
     )
@@ -108,7 +109,7 @@ private func expandFreestandingExpr(
     let expanded = expandFreestandingMacro(
       definition: definition,
       macroRole: .expression,
-      node: node.detach(in: context),
+      node: node.detach(in: context, foldingWith: .standardOperators),
       in: context,
       indentationWidth: indentationWidth
     )
@@ -134,7 +135,7 @@ private func expandMemberMacro(
     let expanded = expandAttachedMacro(
       definition: definition,
       macroRole: .member,
-      attributeNode: attributeNode.detach(in: context),
+      attributeNode: attributeNode.detach(in: context, foldingWith: .standardOperators),
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: nil,
@@ -163,7 +164,7 @@ private func expandMemberAttributeMacro(
     let expanded = expandAttachedMacro(
       definition: definition,
       macroRole: .memberAttribute,
-      attributeNode: attributeNode.detach(in: context),
+      attributeNode: attributeNode.detach(in: context, foldingWith: .standardOperators),
       declarationNode: member.detach(in: context),
       parentDeclNode: declaration.detach(in: context),
       extendedType: nil,
@@ -191,7 +192,7 @@ private func expandPeerMacroMember(
     let expanded = expandAttachedMacro(
       definition: definition,
       macroRole: .peer,
-      attributeNode: attributeNode.detach(in: context),
+      attributeNode: attributeNode.detach(in: context, foldingWith: .standardOperators),
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: nil,
@@ -219,7 +220,7 @@ private func expandPeerMacroCodeItem(
     let expanded = expandAttachedMacro(
       definition: definition,
       macroRole: .peer,
-      attributeNode: attributeNode.detach(in: context),
+      attributeNode: attributeNode.detach(in: context, foldingWith: .standardOperators),
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: nil,
@@ -251,7 +252,7 @@ private func expandAccessorMacroWithoutExistingAccessors(
     let expanded = expandAttachedMacro(
       definition: definition,
       macroRole: .accessor,
-      attributeNode: attributeNode.detach(in: context),
+      attributeNode: attributeNode.detach(in: context, foldingWith: .standardOperators),
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: nil,
@@ -285,7 +286,7 @@ private func expandAccessorMacroWithExistingAccessors(
     let expanded = expandAttachedMacro(
       definition: definition,
       macroRole: .accessor,
-      attributeNode: attributeNode.detach(in: context),
+      attributeNode: attributeNode.detach(in: context, foldingWith: .standardOperators),
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: nil,
@@ -322,7 +323,7 @@ private func expandExtensionMacro(
     let expanded = expandAttachedMacro(
       definition: definition,
       macroRole: .extension,
-      attributeNode: attributeNode.detach(in: context),
+      attributeNode: attributeNode.detach(in: context, foldingWith: .standardOperators),
       declarationNode: attachedTo.detach(in: context),
       parentDeclNode: nil,
       extendedType: extendedType.detach(in: context),
@@ -1007,4 +1008,47 @@ private extension SyntaxProtocol {
 
     return self.detached
   }
+
+  /// Fold operators in this node using the given operator table, detach the
+  /// node and inform the macro expansion context, if it needs to know.
+  func detach(
+    in context: MacroExpansionContext,
+    foldingWith operatorTable: OperatorTable?
+  ) -> Syntax {
+    let folded: Syntax
+    if let operatorTable {
+      if let basicContext = context as? BasicMacroExpansionContext {
+        folded = basicContext.foldAllOperators(of: self, with: operatorTable)
+      } else {
+        folded = operatorTable.foldAll(self, errorHandler: { _ in /*ignore*/ })
+      }
+    } else {
+      folded = Syntax(self)
+    }
+    return folded.detach(in: context)
+  }
+}
+
+private extension FreestandingMacroExpansionSyntax {
+  /// Same as `SyntaxProtocol.detach(in:foldingWith:)` but returns a node of type
+  /// `Self` since we know that operator folding doesn't change the type of any
+  /// `FreestandingMacroExpansionSyntax`.
+  func detach(
+    in context: MacroExpansionContext,
+    foldingWith operatorTable: OperatorTable?
+  ) -> Self {
+    return (detach(in: context, foldingWith: operatorTable) as Syntax).cast(Self.self)
+  }
+}
+
+private extension AttributeSyntax {
+  /// Same as `SyntaxProtocol.detach(in:foldingWith:)` but returns a node of type
+  /// `Self` since we know that operator folding doesn't change the type of any
+  /// `AttributeSyntax`.
+  func detach(
+    in context: MacroExpansionContext,
+    foldingWith operatorTable: OperatorTable?
+  ) -> Self {
+    return (detach(in: context, foldingWith: operatorTable) as Syntax).cast(Self.self)
+  }
 }

Tests/SwiftSyntaxMacroExpansionTest/MacroSystemTests.swift

@@ -1894,4 +1894,146 @@ final class MacroSystemTests: XCTestCase {
       macros: ["decl": DeclMacro.self, "Peer": MyPeerMacro.self]
     )
   }
+
+  func testFoldOperatorsOfFreestandingMacro() {
+    struct ForceSubtractMacro: ExpressionMacro {
+      static func expansion(
+        of node: some FreestandingMacroExpansionSyntax,
+        in context: some MacroExpansionContext
+      ) throws -> ExprSyntax {
+        guard let argument = node.argumentList.first?.expression else {
+          fatalError("Must receive an argument")
+        }
+        guard var node = argument.as(InfixOperatorExprSyntax.self) else {
+          return argument
+        }
+        node.operator = ExprSyntax(BinaryOperatorExprSyntax(text: "-"))
+        return ExprSyntax(node)
+      }
+    }
+    assertMacroExpansion(
+      "#test(1 + 2)",
+      expandedSource: "1 - 2",
+      macros: ["test": ForceSubtractMacro.self]
+    )
+  }
+
+  func testDiagnosticFromFoldedOperatorsInFreestandingMacro() {
+    struct MyError: Error {}
+
+    struct DiagnoseFirstArgument: ExpressionMacro {
+      static func expansion(
+        of node: some FreestandingMacroExpansionSyntax,
+        in context: some MacroExpansionContext
+      ) throws -> ExprSyntax {
+        guard let argument = node.argumentList.first?.expression else {
+          fatalError("Must receive an argument")
+        }
+        context.addDiagnostics(from: MyError(), node: argument)
+        return argument
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      /// Test
+      func test() {
+        #test(1 + 2)
+      }
+      """,
+      expandedSource: """
+        /// Test
+        func test() {
+          1 + 2
+        }
+        """,
+      diagnostics: [
+        DiagnosticSpec(message: "MyError()", line: 3, column: 9, severity: .error)
+      ],
+      macros: ["test": DiagnoseFirstArgument.self]
+    )
+  }
+
+  func testFoldOperatorsInAttachedMacro() {
+    struct ForceSubtractMacro: MemberMacro {
+      static func expansion(
+        of node: AttributeSyntax,
+        providingMembersOf declaration: some DeclGroupSyntax,
+        in context: some MacroExpansionContext
+      ) throws -> [DeclSyntax] {
+        guard case .argumentList(let arguments) = node.arguments, let argument = arguments.first?.expression else {
+          fatalError("Must receive an argument")
+        }
+        guard var node = argument.as(InfixOperatorExprSyntax.self) else {
+          return []
+        }
+        node.operator = ExprSyntax(BinaryOperatorExprSyntax(text: "- "))
+        return [
+          DeclSyntax(
+            """
+            var x: Int { \(node.trimmed) }
+            """
+          )
+        ]
+      }
+    }
+    assertMacroExpansion(
+      """
+      /// Test
+      /// And another line
+      @Test(1 + 2)
+      struct Foo {
+      }
+      """,
+      expandedSource: """
+        /// Test
+        /// And another line
+        struct Foo {
+
+          var x: Int {
+            1 - 2
+          }
+        }
+        """,
+      macros: ["Test": ForceSubtractMacro.self],
+      indentationWidth: indentationWidth
+    )
+  }
+
+  func testDiagnosticFromFoldedOperatorsInAttachedMacro() {
+    struct MyError: Error {}
+
+    struct DiagnoseFirstArgument: MemberMacro {
+      static func expansion(
+        of node: AttributeSyntax,
+        providingMembersOf declaration: some DeclGroupSyntax,
+        in context: some MacroExpansionContext
+      ) throws -> [DeclSyntax] {
+        guard case .argumentList(let arguments) = node.arguments, let argument = arguments.first?.expression else {
+          fatalError("Must receive an argument")
+        }
+        context.addDiagnostics(from: MyError(), node: argument)
+        return []
+      }
+    }
+
+    assertMacroExpansion(
+      """
+      /// Test
+      /// And another line
+      @Test(1 + 2)
+      struct Foo {}
+      """,
+      expandedSource: """
+        /// Test
+        /// And another line
+        struct Foo {}
+        """,
+      diagnostics: [
+        DiagnosticSpec(message: "MyError()", line: 3, column: 7, severity: .error)
+      ],
+      macros: ["Test": DiagnoseFirstArgument.self]
+    )
+  }
+
 }