Preparing CompilerInstance for batch mode.

include/swift/Frontend/Frontend.h

@@ -39,6 +39,7 @@
 #include "swift/Serialization/Validation.h"
 #include "swift/Subsystems.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/Option/ArgList.h"
 #include "llvm/Support/Host.h"
 #include "llvm/Support/MemoryBuffer.h"
@@ -336,14 +337,31 @@ class CompilerInstance {
   enum : unsigned { NO_SUCH_BUFFER = ~0U };
   unsigned MainBufferID = NO_SUCH_BUFFER;
 
-  /// PrimaryBufferID corresponds to PrimaryInput.
-  unsigned PrimaryBufferID = NO_SUCH_BUFFER;
-  bool isWholeModuleCompilation() { return PrimaryBufferID == NO_SUCH_BUFFER; }
+  /// Identifies the set of input buffers in the SourceManager that are
+  /// considered primaries.
+  llvm::SetVector<unsigned> PrimaryBufferIDs;
 
-  SourceFile *PrimarySourceFile = nullptr;
+  /// Identifies the set of SourceFiles that are considered primaries. An
+  /// invariant is that any SourceFile in this set with an associated
+  /// buffer will also have its buffer ID in PrimaryBufferIDs.
+  std::vector<SourceFile *> PrimarySourceFiles;
+
+  /// Return whether there is an entry in PrimaryInputs for buffer \p BufID.
+  bool isPrimaryInput(unsigned BufID) const {
+    return PrimaryBufferIDs.count(BufID) != 0;
+  }
+
+  /// Record in PrimaryBufferIDs the fact that \p BufID is a primary.
+  /// If \p BufID is already in the set, do nothing.
+  void recordPrimaryInputBuffer(unsigned BufID);
+
+  /// Record in PrimarySourceFiles the fact that \p SF is a primary, and
+  /// call recordPrimaryInputBuffer on \p SF's buffer (if it exists).
+  void recordPrimarySourceFile(SourceFile *SF);
+
+  bool isWholeModuleCompilation() { return PrimaryBufferIDs.empty(); }
 
   void createSILModule();
-  void setPrimarySourceFile(SourceFile *SF);
 
 public:
   SourceManager &getSourceMgr() { return SourceMgr; }
@@ -371,7 +389,7 @@ class CompilerInstance {
   }
 
   void setReferencedNameTracker(ReferencedNameTracker *tracker) {
-    assert(!PrimarySourceFile && "must be called before performSema()");
+    assert(PrimarySourceFiles.empty() && "must be called before performSema()");
     NameTracker = tracker;
   }
   ReferencedNameTracker *getReferencedNameTracker() {
@@ -413,9 +431,36 @@ class CompilerInstance {
     return Invocation.getFrontendOptions().EnableSourceImport;
   }
 
+  /// Gets the set of SourceFiles which are the primary inputs for this
+  /// CompilerInstance.
+  ArrayRef<SourceFile *> getPrimarySourceFiles() {
+    return PrimarySourceFiles;
+  }
+
+  /// Gets the Primary Source File if one exists, otherwise the main
+  /// module. If multiple Primary Source Files exist, fails with an
+  /// assertion.
+  ModuleOrSourceFile getPrimarySourceFileOrMainModule() {
+    if (PrimarySourceFiles.empty())
+      return getMainModule();
+    else
+      return getPrimarySourceFile();
+  }
+
   /// Gets the SourceFile which is the primary input for this CompilerInstance.
-  /// \returns the primary SourceFile, or nullptr if there is no primary input
-  SourceFile *getPrimarySourceFile() { return PrimarySourceFile; }
+  /// \returns the primary SourceFile, or nullptr if there is no primary input;
+  /// if there are _multiple_ primary inputs, fails with an assertion.
+  ///
+  /// FIXME: This should be removed eventually, once there are no longer any
+  /// codepaths that rely on a single primary file.
+  SourceFile *getPrimarySourceFile() {
+    if (PrimarySourceFiles.empty()) {
+      return nullptr;
+    } else {
+      assert(PrimarySourceFiles.size() == 1);
+      return *PrimarySourceFiles.begin();
+    }
+  }
 
   /// \brief Returns true if there was an error during setup.
   bool setup(const CompilerInvocation &Invocation);

lib/Frontend/Frontend.cpp

@@ -61,14 +61,16 @@ void CompilerInstance::createSILModule() {
       Invocation.getFrontendOptions().Inputs.isWholeModule());
 }
 
-void CompilerInstance::setPrimarySourceFile(SourceFile *SF) {
-  assert(SF);
+void CompilerInstance::recordPrimaryInputBuffer(unsigned BufID) {
+  PrimaryBufferIDs.insert(BufID);
+}
+
+void CompilerInstance::recordPrimarySourceFile(SourceFile *SF) {
   assert(MainModule && "main module not created yet");
-  assert(!PrimarySourceFile && "already has a primary source file");
-  assert(PrimaryBufferID == NO_SUCH_BUFFER || !SF->getBufferID().hasValue() ||
-         SF->getBufferID().getValue() == PrimaryBufferID);
-  PrimarySourceFile = SF;
-  PrimarySourceFile->setReferencedNameTracker(NameTracker);
+  PrimarySourceFiles.push_back(SF);
+  SF->setReferencedNameTracker(NameTracker);
+  if (SF->getBufferID().hasValue())
+    recordPrimaryInputBuffer(SF->getBufferID().getValue());
 }
 
 bool CompilerInstance::setup(const CompilerInvocation &Invok) {
@@ -186,9 +188,9 @@ bool CompilerInstance::setUpInputs() {
 
   // Set the primary file to the code-completion point if one exists.
   if (codeCompletionBufferID.hasValue() &&
-      *codeCompletionBufferID != PrimaryBufferID) {
-    assert(PrimaryBufferID == NO_SUCH_BUFFER && "re-setting PrimaryBufferID");
-    PrimaryBufferID = *codeCompletionBufferID;
+      !isPrimaryInput(*codeCompletionBufferID)) {
+    assert(PrimaryBufferIDs.empty() && "re-setting PrimaryBufferID");
+    recordPrimaryInputBuffer(*codeCompletionBufferID);
   }
 
   if (isInputSwift() && MainBufferID == NO_SUCH_BUFFER &&
@@ -214,8 +216,8 @@ bool CompilerInstance::setUpForInput(const InputFile &input) {
   }
 
   if (input.isPrimary()) {
-    assert(PrimaryBufferID == NO_SUCH_BUFFER && "re-setting PrimaryBufferID");
-    PrimaryBufferID = *bufferID;
+    assert(PrimaryBufferIDs.empty() && "re-setting PrimaryBufferID");
+    recordPrimaryInputBuffer(*bufferID);
   }
   return false;
 }
@@ -587,15 +589,15 @@ void CompilerInstance::parseLibraryFile(
   addAdditionalInitialImportsTo(NextInput, implicitImports);
 
   auto *DelayedCB = SecondaryDelayedCB;
-  if (BufferID == PrimaryBufferID) {
+  if (isPrimaryInput(BufferID)) {
     DelayedCB = PrimaryDelayedCB;
   }
   if (isWholeModuleCompilation())
     DelayedCB = PrimaryDelayedCB;
 
   auto &Diags = NextInput->getASTContext().Diags;
   auto DidSuppressWarnings = Diags.getSuppressWarnings();
-  auto IsPrimary = isWholeModuleCompilation() || BufferID == PrimaryBufferID;
+  auto IsPrimary = isWholeModuleCompilation() || isPrimaryInput(BufferID);
   Diags.setSuppressWarnings(DidSuppressWarnings || !IsPrimary);
 
   bool Done;
@@ -661,7 +663,7 @@ void CompilerInstance::parseAndTypeCheckMainFile(
   SharedTimer timer(
       "performSema-checkTypesWhileParsingMain-parseAndTypeCheckMainFile");
   bool mainIsPrimary =
-      (isWholeModuleCompilation() || MainBufferID == PrimaryBufferID);
+      (isWholeModuleCompilation() || isPrimaryInput(MainBufferID));
 
   SourceFile &MainFile =
       MainModule->getMainSourceFile(Invocation.getSourceFileKind());
@@ -724,7 +726,9 @@ void CompilerInstance::forEachFileToTypeCheck(
   if (isWholeModuleCompilation()) {
     forEachSourceFileIn(MainModule, [&](SourceFile &SF) { fn(SF); });
   } else {
-    fn(*PrimarySourceFile);
+    for (auto *SF : PrimarySourceFiles) {
+      fn(*SF);
+    }
   }
 }
 
@@ -746,8 +750,9 @@ SourceFile *CompilerInstance::createSourceFileForMainModule(
       SourceFile(*mainModule, fileKind, bufferID, importKind, keepSyntaxInfo);
   MainModule->addFile(*inputFile);
 
-  if (bufferID && *bufferID == PrimaryBufferID)
-    setPrimarySourceFile(inputFile);
+  if (bufferID && isPrimaryInput(*bufferID)) {
+    recordPrimarySourceFile(inputFile);
+  }
 
   return inputFile;
 }
@@ -813,6 +818,7 @@ void CompilerInstance::freeContextAndSIL() {
   TheSILModule.reset();
   MainModule = nullptr;
   SML = nullptr;
-  PrimarySourceFile = nullptr;
+  PrimaryBufferIDs.clear();
+  PrimarySourceFiles.clear();
 }