Re-work MacroSystem on top of the same parsing functions that the compiler invokes

This should make the MacroSystem behave more similarly to the expansions performed by the compiler.

Author: ahoppen

Package.swift

@@ -182,7 +182,7 @@ let package = Package(
 
     .target(
       name: "SwiftSyntaxMacroExpansion",
-      dependencies: ["SwiftSyntax", "SwiftSyntaxMacros", "SwiftDiagnostics"],
+      dependencies: ["SwiftSyntax", "SwiftSyntaxBuilder", "SwiftSyntaxMacros", "SwiftDiagnostics"],
       exclude: ["CMakeLists.txt"]
     ),
 

Sources/SwiftSyntaxBuilder/Indenter.swift

@@ -25,15 +25,16 @@ extension Trivia {
   }
 }
 
-class Indenter: SyntaxRewriter {
+/// Adds a given amount of indentation after every newline in a syntax tree.
+public class Indenter: SyntaxRewriter {
   let indentation: Trivia
 
   init(indentation: Trivia) {
     self.indentation = indentation
     super.init(viewMode: .sourceAccurate)
   }
 
-  /// Adds `indentation` after all newlines in the syntax tree.
+  /// Add `indentation` after all newlines in the syntax tree.
   public static func indent<SyntaxType: SyntaxProtocol>(
     _ node: SyntaxType,
     indentation: Trivia

Sources/SwiftSyntaxMacroExpansion/BasicMacroExpansionContext.swift

@@ -48,11 +48,13 @@ public class BasicMacroExpansionContext {
   /// information about that source file.
   private var sourceFiles: [SourceFileSyntax: KnownSourceFile] = [:]
 
-  /// Mapping from intentionally-disconnected syntax node roots to the
-  /// absolute offsets that have within a given source file, which is used
-  /// to establish the link between a node that been intentionally disconnected
-  /// from a source file to hide information from the macro implementation.
-  private var disconnectedNodes: [Syntax: (SourceFileSyntax, Int)] = [:]
+  /// Mapping from intentionally-disconnected syntax nodes to the corresponding
+  /// nodes in the original source file.
+  ///
+  /// This is used to establish the link between a node that been intentionally
+  /// disconnected from a source file to hide information from the macro
+  /// implementation.
+  private var detachedNodes: [Syntax: Syntax] = [:]
 
   /// The macro expansion discriminator, which is used to form unique names
   /// when requested.
@@ -69,24 +71,10 @@ public class BasicMacroExpansionContext {
 }
 
 extension BasicMacroExpansionContext {
-  /// Note that the given node that was at the given position in the provided
-  /// source file has been disconnected and is now a new root.
-  private func addDisconnected(
-    _ node: some SyntaxProtocol,
-    at offset: AbsolutePosition,
-    in sourceFile: SourceFileSyntax
-  ) {
-    disconnectedNodes[Syntax(node)] = (sourceFile, offset.utf8Offset)
-  }
-
   /// Detach the given node, and record where it came from.
   public func detach<Node: SyntaxProtocol>(_ node: Node) -> Node {
     let detached = node.detached
-
-    if let rootSourceFile = node.root.as(SourceFileSyntax.self) {
-      addDisconnected(detached, at: node.position, in: rootSourceFile)
-    }
-
+    detachedNodes[Syntax(detached)] = Syntax(node)
     return detached
   }
 }
@@ -136,28 +124,49 @@ extension BasicMacroExpansionContext: MacroExpansionContext {
     diagnostics.append(diagnostic)
   }
 
+  /// Translates a position from a detached node to the corresponding location
+  /// in the original source file.
+  ///
+  /// - Parameters:
+  ///   - position: The position to translate
+  ///   - node: The node at which the position is anchored. This node is used to
+  ///     find the offset in the original source file
+  ///   - fileName: The file name that should be used in the `SourceLocation`
+  /// - Returns: The location in the original source file
+  public func location(
+    for position: AbsolutePosition,
+    anchoredAt node: Syntax,
+    fileName: String
+  ) -> SourceLocation {
+    guard let nodeInOriginalTree = detachedNodes[node.root] else {
+      return SourceLocationConverter(file: fileName, tree: node.root).location(for: position)
+    }
+    let adjustedPosition = position + SourceLength(utf8Length: nodeInOriginalTree.position.utf8Offset)
+    return SourceLocationConverter(file: fileName, tree: nodeInOriginalTree.root).location(for: adjustedPosition)
+  }
+
   public func location(
     of node: some SyntaxProtocol,
     at position: PositionInSyntaxNode,
     filePathMode: SourceLocationFilePathMode
   ) -> AbstractSourceLocation? {
     // Dig out the root source file and figure out how we need to adjust the
     // offset of the given syntax node to adjust for it.
-    let rootSourceFile: SourceFileSyntax
-    let offsetAdjustment: Int
+    let rootSourceFile: SourceFileSyntax?
+    let offsetAdjustment: SourceLength
     if let directRootSourceFile = node.root.as(SourceFileSyntax.self) {
       // The syntax node came from the source file itself.
       rootSourceFile = directRootSourceFile
-      offsetAdjustment = 0
-    } else if let (adjustedSourceFile, offset) = disconnectedNodes[Syntax(node)] {
+      offsetAdjustment = .zero
+    } else if let nodeInOriginalTree = detachedNodes[Syntax(node)] {
       // The syntax node came from a disconnected root, so adjust for that.
-      rootSourceFile = adjustedSourceFile
-      offsetAdjustment = offset
+      rootSourceFile = nodeInOriginalTree.root.as(SourceFileSyntax.self)
+      offsetAdjustment = SourceLength(utf8Length: nodeInOriginalTree.position.utf8Offset)
     } else {
       return nil
     }
 
-    guard let knownRoot = sourceFiles[rootSourceFile] else {
+    guard let rootSourceFile, let knownRoot = sourceFiles[rootSourceFile] else {
       return nil
     }
 
@@ -189,6 +198,6 @@ extension BasicMacroExpansionContext: MacroExpansionContext {
 
     // Do the location lookup.
     let converter = SourceLocationConverter(file: fileName, tree: rootSourceFile)
-    return AbstractSourceLocation(converter.location(for: rawPosition.advanced(by: offsetAdjustment)))
+    return AbstractSourceLocation(converter.location(for: rawPosition + offsetAdjustment))
   }
 }

Sources/SwiftSyntaxMacroExpansion/CMakeLists.txt

@@ -1,10 +1,10 @@
 add_swift_host_library(SwiftSyntaxMacroExpansion
   BasicMacroExpansionContext.swift
   FunctionParameterUtils.swift
+  IndentationUtils.swift
   MacroExpansion.swift
   MacroReplacement.swift
   MacroSystem.swift
-  Syntax+MacroEvaluation.swift
 )
 
 target_link_libraries(SwiftSyntaxMacroExpansion PUBLIC

Sources/SwiftSyntaxMacroExpansion/IndentationUtils.swift

@@ -0,0 +1,95 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+
+// MARK: SyntaxProtocol.indentationOfFirstLine
+
+extension SyntaxProtocol {
+  /// The indentation of the first line in this token.
+  var indentationOfFirstLine: Trivia {
+    guard let firstToken = self.firstToken(viewMode: .sourceAccurate) else {
+      return Trivia()
+    }
+    return firstToken.indentationOfLine
+  }
+}
+
+// MARK: String.indented
+
+extension String {
+  /// Indents every new line in this string literal by `indentation`.
+  ///
+  /// - Note: The first line in the string gets indented as well.
+  func indented(by indentation: Trivia) -> String {
+    if isEmpty || indentation.isEmpty {
+      return self
+    }
+
+    var indented = ""
+    var remaining = self[...]
+    while let nextNewline = remaining.firstIndex(where: { $0.isNewline }) {
+      if nextNewline != remaining.startIndex {
+        // Don’t add indentation if the line is empty.
+        indentation.write(to: &indented)
+      }
+      indented += remaining[...nextNewline]
+      remaining = remaining[remaining.index(after: nextNewline)...]
+    }
+
+    if !remaining.isEmpty {
+      indentation.write(to: &indented)
+      indented += remaining
+    }
+
+    return indented
+  }
+}
+
+// MARK: SyntaxProtocol.stripp
+
+fileprivate class IndentationStripper: SyntaxRewriter {
+  override func visit(_ token: TokenSyntax) -> TokenSyntax {
+    if token.leadingTrivia.contains(where: \.isNewline) || token.trailingTrivia.contains(where: \.isNewline) {
+      return
+        token
+        .with(\.leadingTrivia, token.leadingTrivia.removingIndentation)
+        .with(\.trailingTrivia, token.trailingTrivia.removingIndentation)
+    } else {
+      return token
+    }
+  }
+}
+
+extension Trivia {
+  /// Remove all indentation from the trivia.
+  var removingIndentation: Trivia {
+    var resultPieces: [TriviaPiece] = []
+    var isAfterNewline = false
+    for piece in pieces {
+      if piece.isSpaceOrTab && isAfterNewline {
+        // Don’t add whitespace after a newline
+        continue
+      }
+      isAfterNewline = piece.isNewline
+      resultPieces.append(piece)
+    }
+    return Trivia(pieces: resultPieces)
+  }
+}
+
+extension SyntaxProtocol {
+  /// This syntax node with all indentation removed.
+  var withIndentationRemoved: Self {
+    return IndentationStripper().rewrite(self).cast(Self.self)
+  }
+}