Merge pull request #70415 from ahoppen/ahoppen/cache-source-loc-converter

[Macros] Cache `SourceLocationConverter` in `ExportedSourceFile`

Author: bnbarham

lib/ASTGen/Sources/ASTGen/PluginHost.swift

@@ -386,11 +386,9 @@ extension PluginMessage.Syntax {
     }
 
     let source = syntax.description
-    let sourceStr = String(decoding: sourceFilePtr.pointee.buffer, as: UTF8.self)
     let fileName = sourceFilePtr.pointee.fileName
     let fileID = "\(sourceFilePtr.pointee.moduleName)/\(sourceFilePtr.pointee.fileName.basename)"
-    let converter = SourceLocationConverter(file: fileName, source: sourceStr)
-    let loc = converter.location(for: syntax.position)
+    let loc = sourceFilePtr.pointee.sourceLocationConverter.location(for: syntax.position)
 
     self.init(
       kind: kind,

lib/ASTGen/Sources/ASTGen/SourceFile.swift

@@ -32,6 +32,11 @@ public struct ExportedSourceFile {
   /// The syntax tree for the complete source file.
   public let syntax: SourceFileSyntax
 
+  /// A source location converter to convert `AbsolutePosition`s in `syntax` to line/column locations.
+  /// 
+  /// Cached so we don't need to re-build the line table every time we need to convert a position.
+  let sourceLocationConverter: SourceLocationConverter
+
   public func position(of location: BridgedSourceLoc) -> AbsolutePosition? {
     let sourceFileBaseAddress = UnsafeRawPointer(buffer.baseAddress!)
     guard let opaqueValue = location.getOpaquePointerValue() else {
@@ -74,12 +79,15 @@ public func parseSourceFile(
   let sourceFile = Parser.parse(source: buffer, experimentalFeatures: .init(from: ctx))
 
   let exportedPtr = UnsafeMutablePointer<ExportedSourceFile>.allocate(capacity: 1)
+  let moduleName = String(cString: moduleName)
+  let fileName = String(cString: filename)
   exportedPtr.initialize(
     to: .init(
       buffer: buffer,
-      moduleName: String(cString: moduleName),
-      fileName: String(cString: filename),
-      syntax: sourceFile
+      moduleName: moduleName,
+      fileName: fileName,
+      syntax: sourceFile,
+      sourceLocationConverter: SourceLocationConverter(fileName: fileName, tree: sourceFile)
     )
   )
 

lib/ASTGen/Sources/ASTGen/SourceManager+MacroExpansionContext.swift

@@ -117,15 +117,6 @@ extension SourceManager.MacroExpansionContext: MacroExpansionContext {
       return nil
     }
 
-    // Determine the filename to use in the resulting location.
-    let fileName: String
-    switch filePathMode {
-    case .fileID:
-      fileName = "\(exportedSourceFile.moduleName)/\(exportedSourceFile.fileName.basename)"
-
-    case .filePath:
-      fileName = exportedSourceFile.fileName
-    }
 
     // Find the node's offset relative to its root.
     let rawPosition: AbsolutePosition
@@ -146,8 +137,37 @@ extension SourceManager.MacroExpansionContext: MacroExpansionContext {
     let offsetWithinSyntaxNode =
       rawPosition.utf8Offset - node.position.utf8Offset
 
+    var location = exportedSourceFile.sourceLocationConverter.location(
+      for: rootPosition.advanced(by: offsetWithinSyntaxNode)
+    )
+
+    switch filePathMode {
+    case .fileID:
+      // The `SourceLocationConverter` in `exportedSourceFile` uses `filePath` as the file mode. When the `fileID` mode
+      // is requested, we need to adjust the file and presumed file to the `fileID`.
+      let fileID = "\(exportedSourceFile.moduleName)/\(exportedSourceFile.fileName.basename)"
+      var adjustedFile = location.file
+      if adjustedFile == exportedSourceFile.fileName {
+        adjustedFile = fileID
+      }
+      var adjustedPresumedFile = location.presumedFile
+      if adjustedPresumedFile == exportedSourceFile.fileName {
+        adjustedPresumedFile = fileID
+      }
+      location = SourceLocation(
+        line: location.line,
+        column: location.column,
+        offset: location.offset,
+        file: adjustedFile,
+        presumedLine: location.presumedLine,
+        presumedFile: adjustedPresumedFile
+      )
+
+    case .filePath:
+      break
+    }
+
     // Do the location lookup.
-    let converter = SourceLocationConverter(file: fileName, tree: sourceFile)
-    return AbstractSourceLocation(converter.location(for: rootPosition.advanced(by: offsetWithinSyntaxNode)))
+    return AbstractSourceLocation(location)
   }
 }