[ASTGen] Handle '#sourceLocation' directives

include/swift/Basic/BasicBridging.h

@@ -450,6 +450,18 @@ enum ENUM_EXTENSIBILITY_ATTR(closed) BridgedGeneratedSourceFileKind {
   BridgedGeneratedSourceFileKindNone,
 };
 
+//===----------------------------------------------------------------------===//
+// MARK: VirtualFile
+//===----------------------------------------------------------------------===//
+
+struct BridgedVirtualFile {
+  size_t StartPosition;
+  size_t EndPosition;
+  BridgedStringRef Name;
+  ptrdiff_t LineOffset;
+  size_t NamePosition;
+};
+
 SWIFT_END_NULLABILITY_ANNOTATIONS
 
 #ifndef PURE_BRIDGING_MODE

include/swift/Bridging/ASTGen.h

@@ -94,6 +94,12 @@ intptr_t swift_ASTGen_configuredRegions(
 void swift_ASTGen_freeConfiguredRegions(
     BridgedIfConfigClauseRangeInfo *_Nullable regions, intptr_t numRegions);
 
+size_t
+swift_ASTGen_virtualFiles(void *_Nonnull sourceFile,
+                          BridgedVirtualFile *_Nullable *_Nonnull virtualFiles);
+void swift_ASTGen_freeBridgedVirtualFiles(
+    BridgedVirtualFile *_Nullable virtualFiles, size_t numFiles);
+
 #ifdef __cplusplus
 }
 #endif

lib/ASTGen/Sources/ASTGen/SourceFile.swift

@@ -15,7 +15,7 @@ import SwiftDiagnostics
 import SwiftIfConfig
 @_spi(ExperimentalLanguageFeatures) import SwiftParser
 import SwiftParserDiagnostics
-import SwiftSyntax
+@_spi(Compiler) import SwiftSyntax
 
 /// Describes a source file that has been "exported" to the C++ part of the
 /// compiler, with enough information to interface with the C++ layer.
@@ -299,3 +299,46 @@ public func findSyntaxNodeInSourceFile<Node: SyntaxProtocol>(
 
   return resultSyntax
 }
+
+@_cdecl("swift_ASTGen_virtualFiles")
+@usableFromInline
+func getVirtualFiles(
+  sourceFilePtr: UnsafeMutableRawPointer,
+  cVirtualFilesOut: UnsafeMutablePointer<UnsafeMutablePointer<BridgedVirtualFile>?>
+) -> Int {
+  let sourceFilePtr = sourceFilePtr.assumingMemoryBound(to: ExportedSourceFile.self)
+  let virtualFiles = sourceFilePtr.pointee.sourceLocationConverter.lineTable.virtualFiles
+  guard !virtualFiles.isEmpty else {
+    cVirtualFilesOut.pointee = nil
+    return 0
+  }
+
+  let cArrayBuf: UnsafeMutableBufferPointer<BridgedVirtualFile> = .allocate(capacity: virtualFiles.count)
+  _ = cArrayBuf.initialize(
+    from: virtualFiles.lazy.map({ virtualFile in
+      BridgedVirtualFile(
+        StartPosition: virtualFile.startPosition.utf8Offset,
+        EndPosition: virtualFile.endPosition.utf8Offset,
+        Name: allocateBridgedString(virtualFile.fileName),
+        LineOffset: virtualFile.lineOffset,
+        NamePosition: virtualFile.fileNamePosition.utf8Offset
+      )
+    })
+  )
+
+  cVirtualFilesOut.pointee = cArrayBuf.baseAddress
+  return cArrayBuf.count
+}
+
+@_cdecl("swift_ASTGen_freeBridgedVirtualFiles")
+func freeVirtualFiles(
+  cVirtualFiles: UnsafeMutablePointer<BridgedVirtualFile>?,
+  numFiles: Int
+) {
+  let buffer = UnsafeMutableBufferPointer<BridgedVirtualFile>(start: cVirtualFiles, count: numFiles)
+  for vFile in buffer {
+    freeBridgedString(bridged: vFile.Name)
+  }
+  buffer.deinitialize()
+  buffer.deallocate()
+}

lib/Parse/ParseRequests.cpp

@@ -252,6 +252,7 @@ void appendToVector(BridgedASTNode cNode, void *vecPtr) {
 SourceFileParsingResult parseSourceFileViaASTGen(SourceFile &SF) {
   ASTContext &Ctx = SF.getASTContext();
   DiagnosticEngine &Diags = Ctx.Diags;
+  SourceManager &SM = Ctx.SourceMgr;
   const LangOptions &langOpts = Ctx.LangOpts;
   const GeneratedSourceInfo *genInfo = SF.getGeneratedSourceFileInfo();
 
@@ -264,6 +265,24 @@ SourceFileParsingResult parseSourceFileViaASTGen(SourceFile &SF) {
   auto *exportedSourceFile = SF.getExportedSourceFile();
   assert(exportedSourceFile && "Couldn't parse via SyntaxParser");
 
+  // Collect virtual files.
+  // FIXME: Avoid side effects in the request.
+  // FIXME: Do this lazily in SourceManager::getVirtualFile().
+  BridgedVirtualFile *virtualFiles = nullptr;
+  size_t numVirtualFiles =
+      swift_ASTGen_virtualFiles(exportedSourceFile, &virtualFiles);
+  SourceLoc bufferStart = SM.getLocForBufferStart(SF.getBufferID());
+  for (size_t i = 0; i != numVirtualFiles; ++i) {
+    auto &VF = virtualFiles[i];
+    Ctx.SourceMgr.createVirtualFile(
+        bufferStart.getAdvancedLoc(VF.StartPosition), VF.Name.unbridged(),
+        VF.LineOffset, VF.EndPosition - VF.StartPosition);
+    StringRef name = Ctx.AllocateCopy(VF.Name.unbridged());
+    SF.VirtualFilePaths.emplace_back(
+        name, bufferStart.getAdvancedLoc(VF.NamePosition));
+  }
+  swift_ASTGen_freeBridgedVirtualFiles(virtualFiles, numVirtualFiles);
+
   // Emit parser diagnostics.
   (void)swift_ASTGen_emitParserDiagnostics(
       Ctx, &Diags, exportedSourceFile, /*emitOnlyErrors=*/false,

test/ASTGen/sourcelocation.swift

@@ -0,0 +1,40 @@
+func test(arg: Int) -> Int { 1 }
+
+func foo() {
+  #sourceLocation(file: "first/foo.swift", line: 100)
+  test(arg: 1)
+}
+
+func bar() {
+  #sourceLocation(file: "second/foo.swift", line: 100)
+}
+
+test(arg: 2)
+
+// RUN: %target-swift-frontend -emit-silgen -module-name MyMod %s -enable-experimental-feature ParserASTGen -diagnostic-style llvm \
+// RUN:  2>&1 >/dev/null | %FileCheck --enable-windows-compatibility --strict-whitespace %s
+
+// REQUIRES: swift_swift_parser
+// REQUIRES: swift_feature_ParserASTGen
+
+// CHECK:      {{^}}second/foo.swift:102:1: warning: result of call to 'test(arg:)' is unused
+// CHECK-NEXT: {{^}}test(arg: 2)
+// CHECK-NEXT: {{^}}^   ~~~~~~~~
+
+// CHECK:      {{^}}first/foo.swift:100:3: warning: result of call to 'test(arg:)' is unused
+// CHECK-NEXT: {{^}}  test(arg: 1)
+// CHECK-NEXT: {{^}}  ^   ~~~~~~~~
+
+// CHECK:      {{^SOURCE_DIR[\//]test[/\\]ASTGen[\//]sourcelocation\.swift}}:4:25: warning: '#sourceLocation' directive produces '#fileID' string of 'MyMod/foo.swift', which conflicts with '#fileID' strings produced by other paths in the module
+// CHECK-NEXT: {{^}}  #sourceLocation(file: "first/foo.swift", line: 100)
+// CHECK-NEXT: {{^}}                        ^
+
+// CHECK:      {{^SOURCE_DIR[\//]test[/\\]ASTGen[\//]sourcelocation\.swift}}:9:25: warning: '#sourceLocation' directive produces '#fileID' string of 'MyMod/foo.swift', which conflicts with '#fileID' strings produced by other paths in the module
+// CHECK-NEXT: {{^}}  #sourceLocation(file: "second/foo.swift", line: 100)
+// CHECK-NEXT: {{^}}                        ^
+
+// CHECK:      {{^SOURCE_DIR[\//]test[/\\]ASTGen[\//]sourcelocation\.swift}}:9:25: note: change file in '#sourceLocation' to 'first/foo.swift'
+// CHECK-NEXT: {{^}}  #sourceLocation(file: "second/foo.swift", line: 100)
+// CHECK-NEXT: {{^}}                        ^~~~~~~~~~~~~~~~~~
+// CHECK-NEXT: {{^}}                        "first/foo.swift"
+