Hold on to the SwiftSyntax source file parsed by ASTGen.

Rework the ASTGen interface to split apart parsing a source file,
turning the top-level declarations from that source file into C++ AST
nodes, and then deallocating that source file. Hold on to the source
file in the C++ SourceFile abstraction so we can query it later if we
need to.

And we will need to.

Author: DougGregor

include/swift/AST/SourceFile.h

@@ -316,6 +316,10 @@ class SourceFile final : public FileUnit {
   /// this source file.
   llvm::SmallVector<Located<StringRef>, 0> VirtualFilePaths;
 
+  /// The \c ExportedSourceFile instance produced by ASTGen, which includes
+  /// the SourceFileSyntax node corresponding to this source file.
+  void *exportedSourceFile = nullptr;
+
   /// Returns information about the file paths used for diagnostics and magic
   /// identifiers in this source file, including virtual filenames introduced by
   /// \c #sourceLocation(file:) declarations.

lib/ASTGen/Sources/ASTGen/ASTGen.swift

@@ -17,7 +17,8 @@ extension UnsafePointer {
   }
 }
 
-/// Little utility wrapper that lets us
+/// Little utility wrapper that lets us have some mutable state within
+/// immutable structs, and is therefore pretty evil.
 @propertyWrapper
 class Boxed<Value> {
   var wrappedValue: Value
@@ -29,7 +30,7 @@ class Boxed<Value> {
 
 struct ASTGenVisitor: SyntaxTransformVisitor {
   let ctx: UnsafeMutableRawPointer
-  let base: UnsafePointer<CChar>
+  let base: UnsafePointer<UInt8>
 
   @Boxed var declContext: UnsafeMutableRawPointer
 
@@ -68,15 +69,65 @@ struct ASTGenVisitor: SyntaxTransformVisitor {
   }
 }
 
-@_cdecl("parseTopLevelSwift")
-public func parseTopLevelSwift(
-    buffer: UnsafePointer<CChar>, dc: UnsafeMutableRawPointer,
-    ctx: UnsafeMutableRawPointer,
-    outputContext: UnsafeMutableRawPointer,
-    callback: @convention(c) (UnsafeMutableRawPointer, UnsafeMutableRawPointer) -> Void
+/// Describes a source file that has been "exported" to the C++ part of the
+/// compiler, with enough information to interface with the C++ layer.
+struct ExportedSourceFile {
+  /// The underlying buffer within the C++ SourceManager, which is used
+  /// for computations of source locations.
+  let buffer: UnsafeBufferPointer<UInt8>
+
+  /// The name of the enclosing module.
+  let moduleName: String
+
+  /// The name of the source file being parsed.
+  let fileName: String
+
+  /// The syntax tree for the complete source file.
+  let syntax: SourceFileSyntax
+}
+
+/// Parses the given source file and produces a pointer to a new
+/// ExportedSourceFile instance.
+@_cdecl("swift_ASTGen_parseSourceFile")
+public func parseSourceFile(
+  buffer: UnsafePointer<UInt8>, bufferLength: Int,
+  moduleName: UnsafePointer<UInt8>, filename: UnsafePointer<UInt8>
+) -> UnsafeRawPointer {
+  let buffer = UnsafeBufferPointer(start: buffer, count: bufferLength)
+  let sourceFile = Parser.parse(source: buffer)
+
+  let exportedPtr = UnsafeMutablePointer<ExportedSourceFile>.allocate(capacity: 1)
+  exportedPtr.initialize(to: .init(
+    buffer: buffer, moduleName: String(cString: moduleName),
+    fileName: String(cString: filename), syntax: sourceFile)
+  )
+
+  return UnsafeRawPointer(exportedPtr)
+}
+
+/// Deallocate a parsed source file.
+@_cdecl("swift_ASTGen_destroySourceFile")
+public func destroySourceFile(
+  sourceFilePtr: UnsafeMutableRawPointer
 ) {
-  let syntax = try! Parser.parse(source: String(cString: buffer))
-  ASTGenVisitor(ctx: ctx, base: buffer, declContext: dc)
-    .visit(syntax)
-    .forEach { callback($0, outputContext) }
+  sourceFilePtr.withMemoryRebound(to: ExportedSourceFile.self, capacity: 1) { sourceFile in
+    sourceFile.deinitialize(count: 1)
+    sourceFile.deallocate()
+  }
+}
+
+/// Generate AST nodes for all top-level entities in the given source file.
+@_cdecl("swift_ASTGen_buildTopLevelASTNodes")
+public func buildTopLevelASTNodes(
+  sourceFilePtr: UnsafeRawPointer,
+  dc: UnsafeMutableRawPointer,
+  ctx: UnsafeMutableRawPointer,
+  outputContext: UnsafeMutableRawPointer,
+  callback: @convention(c) (UnsafeMutableRawPointer, UnsafeMutableRawPointer) -> Void
+) {
+  sourceFilePtr.withMemoryRebound(to: ExportedSourceFile.self, capacity: 1) { sourceFile in
+    ASTGenVisitor(ctx: ctx, base: sourceFile.pointee.buffer.baseAddress!, declContext: dc)
+      .visit(sourceFile.pointee.syntax)
+      .forEach { callback($0, outputContext) }
+  }
 }

lib/Parse/ParseDecl.cpp

@@ -173,6 +173,22 @@ static void appendToVector(void *declPtr, void *vecPtr) {
   vec->push_back(decl);
 }
 
+/// Parse a source file.
+extern "C" void *swift_ASTGen_parseSourceFile(const char *buffer,
+                                              size_t bufferLength,
+                                              const char *moduleName,
+                                              const char *filename);
+
+/// Destroy a source file parsed with swift_ASTGen_parseSourceFile.
+extern "C" void swift_ASTGen_destroySourceFile(void *sourceFile);
+
+// Build AST nodes for the top-level entities in the syntax.
+extern "C" void swift_ASTGen_buildTopLevelASTNodes(void *sourceFile,
+                                                   void *declContext,
+                                                   void *astContext,
+                                                   void *outputContext,
+                                                   void (*)(void *, void *));
+
 /// Main entrypoint for the parser.
 ///
 /// \verbatim
@@ -182,20 +198,35 @@ static void appendToVector(void *declPtr, void *vecPtr) {
 ///     decl-sil-stage [[only in SIL mode]
 /// \endverbatim
 void Parser::parseTopLevel(SmallVectorImpl<Decl *> &decls) {
-  StringRef contents =
-      SourceMgr.extractText(SourceMgr.getRangeForBuffer(L->getBufferID()));
-
 #ifdef SWIFT_SWIFT_PARSER
   if (Context.LangOpts.hasFeature(Feature::ParserASTGen) &&
       !SourceMgr.hasCodeCompletionBuffer() &&
       SF.Kind != SourceFileKind::SIL) {
-    parseTopLevelSwift(contents.data(), CurDeclContext, &Context, &decls, appendToVector);
+    StringRef contents =
+        SourceMgr.extractText(SourceMgr.getRangeForBuffer(L->getBufferID()));
+
+    // Parse the source file.
+    auto exportedSourceFile = swift_ASTGen_parseSourceFile(
+        contents.begin(), contents.size(),
+        SF.getParentModule()->getName().str().str().c_str(),
+        SF.getFilename().str().c_str());
+    SF.exportedSourceFile = exportedSourceFile;
+    Context.addCleanup([exportedSourceFile] {
+      swift_ASTGen_destroySourceFile(exportedSourceFile);
+    });
 
-    // Spin the C++ parser to the end; we won't be using it.
-    while (!Tok.is(tok::eof)) {
-      consumeToken();
+    // If we want to do ASTGen, do so now.
+    if (Context.LangOpts.hasFeature(Feature::ParserASTGen)) {
+      swift_ASTGen_buildTopLevelASTNodes(
+          exportedSourceFile, CurDeclContext, &Context, &decls, appendToVector);
+
+      // Spin the C++ parser to the end; we won't be using it.
+      while (!Tok.is(tok::eof)) {
+        consumeToken();
+      }
+
+      return;
     }
-    return;
   }
 #endif