Add request to query primary files

Remove the `PrimarySourceFiles` vector from the
frontend and replace it with a request on
ModuleDecl that retrieves the primary files for
the main module.

This is in preparation for having
`CompilerInstance::getMainModule` automatically
populate the main module with files when queried.

Author: hamishknight

include/swift/AST/Module.h

@@ -548,6 +548,10 @@ class ModuleDecl : public DeclContext, public TypeDecl {
     return Bits.ModuleDecl.IsMainModule;
   }
 
+  /// For the main module, retrieves the list of primary source files being
+  /// compiled, that is, the files we're generating code for.
+  ArrayRef<SourceFile *> getPrimarySourceFiles() const;
+
   /// Retrieve the top-level module. If this module is already top-level, this
   /// returns itself. If this is a submodule such as \c Foo.Bar.Baz, this
   /// returns the module \c Foo.

include/swift/AST/SourceFile.h

@@ -158,6 +158,9 @@ class SourceFile final : public FileUnit {
   /// The parsing options for the file.
   ParsingOptions ParsingOpts;
 
+  /// Whether this is a primary source file which we'll be generating code for.
+  bool IsPrimary;
+
   /// The scope map that describes this source file.
   std::unique_ptr<ASTScope> Scope;
 
@@ -232,6 +235,10 @@ class SourceFile final : public FileUnit {
   /// Retrieve the parsing options for the file.
   ParsingOptions getParsingOptions() const { return ParsingOpts; }
 
+  /// Whether this source file is a primary file, meaning that we're generating
+  /// code for it. Note this method returns \c false in WMO.
+  bool isPrimary() const { return IsPrimary; }
+
   /// A cache of syntax nodes that can be reused when creating the syntax tree
   /// for this file.
   swift::SyntaxParsingCache *SyntaxParsingCache = nullptr;
@@ -307,7 +314,7 @@ class SourceFile final : public FileUnit {
 
   SourceFile(ModuleDecl &M, SourceFileKind K, Optional<unsigned> bufferID,
              bool KeepParsedTokens = false, bool KeepSyntaxTree = false,
-             ParsingOptions parsingOpts = {});
+             ParsingOptions parsingOpts = {}, bool isPrimary = false);
 
   ~SourceFile();
 

include/swift/AST/TypeCheckRequests.h

@@ -2459,6 +2459,26 @@ class CheckInconsistentImplementationOnlyImportsRequest
   bool isCached() const { return true; }
 };
 
+/// Retrieves the primary source files in the main module.
+// FIXME: This isn't really a type-checking request, if we ever split off a
+// zone for more basic AST requests, this should be moved there.
+class PrimarySourceFilesRequest
+    : public SimpleRequest<PrimarySourceFilesRequest,
+                           ArrayRef<SourceFile *>(ModuleDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  ArrayRef<SourceFile *> evaluate(Evaluator &evaluator, ModuleDecl *mod) const;
+
+public:
+  // Cached.
+  bool isCached() const { return true; }
+};
+
 // Allow AnyValue to compare two Type values, even though Type doesn't
 // support ==.
 template<>

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -151,6 +151,8 @@ SWIFT_REQUEST(TypeChecker, OverriddenDeclsRequest,
 SWIFT_REQUEST(TypeChecker, PatternBindingEntryRequest,
               const PatternBindingEntry *(PatternBindingDecl *, unsigned),
               SeparatelyCached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PrimarySourceFilesRequest,
+              ArrayRef<SourceFile *>(ModuleDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, PropertyWrapperBackingPropertyInfoRequest,
               PropertyWrapperBackingPropertyInfo(VarDecl *), Cached,
               NoLocationInfo)

include/swift/Frontend/Frontend.h

@@ -451,11 +451,6 @@ class CompilerInstance {
   /// considered primaries.
   llvm::SetVector<unsigned> PrimaryBufferIDs;
 
-  /// 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;
-
   /// The file that has been registered for code completion.
   NullablePtr<SourceFile> CodeCompletionFile;
 
@@ -536,7 +531,7 @@ class CompilerInstance {
   /// Gets the set of SourceFiles which are the primary inputs for this
   /// CompilerInstance.
   ArrayRef<SourceFile *> getPrimarySourceFiles() const {
-    return PrimarySourceFiles;
+    return getMainModule()->getPrimarySourceFiles();
   }
 
   /// Gets the SourceFile which is the primary input for this CompilerInstance.
@@ -546,11 +541,12 @@ class CompilerInstance {
   /// FIXME: This should be removed eventually, once there are no longer any
   /// codepaths that rely on a single primary file.
   SourceFile *getPrimarySourceFile() const {
-    if (PrimarySourceFiles.empty()) {
+    auto primaries = getPrimarySourceFiles();
+    if (primaries.empty()) {
       return nullptr;
     } else {
-      assert(PrimarySourceFiles.size() == 1);
-      return *PrimarySourceFiles.begin();
+      assert(primaries.size() == 1);
+      return *primaries.begin();
     }
   }
 

lib/AST/Module.cpp

@@ -511,6 +511,27 @@ void ModuleDecl::addFile(FileUnit &newFile) {
   clearLookupCache();
 }
 
+ArrayRef<SourceFile *>
+PrimarySourceFilesRequest::evaluate(Evaluator &evaluator,
+                                    ModuleDecl *mod) const {
+  assert(mod->isMainModule() && "Only the main module can have primaries");
+
+  SmallVector<SourceFile *, 8> primaries;
+  for (auto *file : mod->getFiles()) {
+    if (auto *SF = dyn_cast<SourceFile>(file)) {
+      if (SF->isPrimary())
+        primaries.push_back(SF);
+    }
+  }
+  return mod->getASTContext().AllocateCopy(primaries);
+}
+
+ArrayRef<SourceFile *> ModuleDecl::getPrimarySourceFiles() const {
+  auto &eval = getASTContext().evaluator;
+  auto *mutableThis = const_cast<ModuleDecl *>(this);
+  return evaluateOrDefault(eval, PrimarySourceFilesRequest{mutableThis}, {});
+}
+
 #define FORWARD(name, args) \
   for (const FileUnit *file : getFiles()) \
     file->name args;
@@ -2182,12 +2203,15 @@ ModuleDecl::computeMagicFileStringMap(bool shouldDiagnose) const {
 SourceFile::SourceFile(ModuleDecl &M, SourceFileKind K,
                        Optional<unsigned> bufferID,
                        bool KeepParsedTokens, bool BuildSyntaxTree,
-                       ParsingOptions parsingOpts)
+                       ParsingOptions parsingOpts, bool isPrimary)
     : FileUnit(FileUnitKind::Source, M), BufferID(bufferID ? *bufferID : -1),
-      ParsingOpts(parsingOpts), Kind(K),
+      ParsingOpts(parsingOpts), IsPrimary(isPrimary), Kind(K),
       SyntaxInfo(new SourceFileSyntaxInfo(BuildSyntaxTree)) {
   M.getASTContext().addDestructorCleanup(*this);
 
+  assert(!IsPrimary || M.isMainModule() &&
+         "A primary cannot appear outside the main module");
+
   if (isScriptMode()) {
     bool problem = M.registerEntryPointFile(this, SourceLoc(), None);
     assert(!problem && "multiple main files?");

lib/Frontend/Frontend.cpp

@@ -870,7 +870,7 @@ void CompilerInstance::forEachFileToTypeCheck(
       fn(*SF);
     }
   } else {
-    for (auto *SF : PrimarySourceFiles) {
+    for (auto *SF : getPrimarySourceFiles()) {
       fn(*SF);
     }
   }
@@ -900,11 +900,10 @@ SourceFile *CompilerInstance::createSourceFileForMainModule(
   SourceFile *inputFile = new (*Context)
       SourceFile(*mainModule, fileKind, bufferID,
                  Invocation.getLangOptions().CollectParsedToken,
-                 Invocation.getLangOptions().BuildSyntaxTree, opts);
+                 Invocation.getLangOptions().BuildSyntaxTree, opts, isPrimary);
   MainModule->addFile(*inputFile);
 
   if (isPrimary) {
-    PrimarySourceFiles.push_back(inputFile);
     inputFile->enableInterfaceHash();
   }
 
@@ -923,7 +922,6 @@ void CompilerInstance::freeASTContext() {
   SML = nullptr;
   MemoryBufferLoader = nullptr;
   PrimaryBufferIDs.clear();
-  PrimarySourceFiles.clear();
 }
 
 /// Perform "stable" optimizations that are invariant across compiler versions.