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

[5.10][Macros] Cache `SourceLocationConverter` in `ExportedSourceFile`

Author: bnbarham

lib/ASTGen/Sources/ASTGen/PluginHost.swift

@@ -379,11 +379,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

@@ -23,6 +23,11 @@ struct ExportedSourceFile {
 
   /// The syntax tree for the complete source file.
   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
 }
 
 extension Parser.ExperimentalFeatures {
@@ -53,10 +58,16 @@ 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)
+      buffer: buffer, 
+      moduleName: moduleName,
+      fileName: fileName, 
+      syntax: sourceFile,
+      sourceLocationConverter: SourceLocationConverter(fileName: fileName, tree: sourceFile)
+    )
   )
 
   return UnsafeRawPointer(exportedPtr)

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

@@ -102,15 +102,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
@@ -131,8 +122,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 adjustePresumedFile = location.presumedFile
+      if adjustePresumedFile == exportedSourceFile.fileName {
+        adjustePresumedFile = fileID
+      }
+      location = SourceLocation(
+        line: location.line,
+        column: location.column,
+        offset: location.offset,
+        file: adjustedFile,
+        presumedLine: location.presumedLine,
+        presumedFile: adjustePresumedFile
+      )
+
+    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)
   }
 }