[CodeCompletion] Replace main module instead of file

Rather than replacing the code completion file
on the `CompilerInstance` whenever we do a cached
top-level completion, let's set a new main module
instead.

This allows us to properly update the
`LoadedModules` map, and allows the retrieval of
the code completion file to be turned into a
request.

Author: hamishknight

include/swift/AST/TypeCheckRequests.h

@@ -2479,6 +2479,25 @@ class PrimarySourceFilesRequest
   bool isCached() const { return true; }
 };
 
+/// Retrieve the file being used for code completion 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 CodeCompletionFileRequest
+    : public SimpleRequest<CodeCompletionFileRequest,
+                           SourceFile *(ModuleDecl *), RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  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

@@ -36,6 +36,8 @@ SWIFT_REQUEST(TypeChecker, CheckRedeclarationRequest,
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, ClassAncestryFlagsRequest,
               AncestryFlags(ClassDecl *), Cached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, CodeCompletionFileRequest,
+              SourceFile *(ModuleDecl *), Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, CompareDeclSpecializationRequest,
               bool (DeclContext *, ValueDecl *, ValueDecl *, bool), Cached,
               NoLocationInfo)

include/swift/Frontend/Frontend.h

@@ -451,9 +451,6 @@ class CompilerInstance {
   /// considered primaries.
   llvm::SetVector<unsigned> PrimaryBufferIDs;
 
-  /// The file that has been registered for code completion.
-  NullablePtr<SourceFile> CodeCompletionFile;
-
   /// Return whether there is an entry in PrimaryInputs for buffer \p BufID.
   bool isPrimaryInput(unsigned BufID) const {
     return PrimaryBufferIDs.count(BufID) != 0;
@@ -509,8 +506,13 @@ class CompilerInstance {
 
   UnifiedStatsReporter *getStatsReporter() const { return Stats.get(); }
 
+  /// Retrieve the main module containing the files being compiled.
   ModuleDecl *getMainModule() const;
 
+  /// Replace the current main module with a new one. This is used for top-level
+  /// cached code completion.
+  void setMainModule(ModuleDecl *newMod);
+
   MemoryBufferSerializedModuleLoader *
   getMemoryBufferSerializedModuleLoader() const {
     return MemoryBufferLoader;
@@ -557,12 +559,7 @@ class CompilerInstance {
 
   /// If a code completion buffer has been set, returns the corresponding source
   /// file.
-  NullablePtr<SourceFile> getCodeCompletionFile() { return CodeCompletionFile; }
-
-  /// Set a new file that we're performing code completion on.
-  void setCodeCompletionFile(SourceFile *file) {
-    CodeCompletionFile = file;
-  }
+  SourceFile *getCodeCompletionFile() const;
 
 private:
   /// Set up the file system by loading and validating all VFS overlay YAML

include/swift/IDE/CompletionInstance.h

@@ -43,7 +43,6 @@ class CompletionInstance {
   std::mutex mtx;
 
   std::unique_ptr<CompilerInstance> CachedCI;
-  ModuleDecl *CurrentModule = nullptr;
   llvm::hash_code CachedArgHash;
   llvm::sys::TimePoint<> DependencyCheckedTimestamp;
   llvm::StringMap<llvm::hash_code> InMemoryDependencyHash;

lib/AST/Module.cpp

@@ -532,6 +532,20 @@ ArrayRef<SourceFile *> ModuleDecl::getPrimarySourceFiles() const {
   return evaluateOrDefault(eval, PrimarySourceFilesRequest{mutableThis}, {});
 }
 
+SourceFile *CodeCompletionFileRequest::evaluate(Evaluator &evaluator,
+                                                ModuleDecl *mod) const {
+  const auto &SM = mod->getASTContext().SourceMgr;
+  assert(mod->isMainModule() && "Can only do completion in the main module");
+  assert(SM.hasCodeCompletionBuffer() && "Not performing code completion?");
+
+  for (auto *file : mod->getFiles()) {
+    auto *SF = dyn_cast<SourceFile>(file);
+    if (SF && SF->getBufferID() == SM.getCodeCompletionBufferID())
+      return SF;
+  }
+  llvm_unreachable("Couldn't find the completion file?");
+}
+
 #define FORWARD(name, args) \
   for (const FileUnit *file : getFiles()) \
     file->name args;

lib/Frontend/Frontend.cpp

@@ -489,6 +489,12 @@ Optional<unsigned> CompilerInstance::setUpCodeCompletionBuffer() {
   return codeCompletionBufferID;
 }
 
+SourceFile *CompilerInstance::getCodeCompletionFile() const {
+  auto *mod = getMainModule();
+  auto &eval = mod->getASTContext().evaluator;
+  return evaluateOrDefault(eval, CodeCompletionFileRequest{mod}, nullptr);
+}
+
 static bool shouldTreatSingleInputAsMain(InputFileKind inputKind) {
   switch (inputKind) {
   case InputFileKind::Swift:
@@ -707,6 +713,12 @@ ModuleDecl *CompilerInstance::getMainModule() const {
   return MainModule;
 }
 
+void CompilerInstance::setMainModule(ModuleDecl *newMod) {
+  assert(newMod->isMainModule());
+  MainModule = newMod;
+  Context->LoadedModules[newMod->getName()] = newMod;
+}
+
 void CompilerInstance::performParseOnly(bool EvaluateConditionals,
                                         bool CanDelayBodies) {
   const InputFileKind Kind = Invocation.getInputKind();
@@ -899,12 +911,6 @@ SourceFile *CompilerInstance::createSourceFileForMainModule(
   if (isPrimary) {
     inputFile->enableInterfaceHash();
   }
-
-  if (bufferID == SourceMgr.getCodeCompletionBufferID()) {
-    assert(!CodeCompletionFile && "Multiple code completion files?");
-    CodeCompletionFile = inputFile;
-  }
-
   return inputFile;
 }
 

lib/IDE/CompletionInstance.cpp

@@ -169,11 +169,10 @@ static DeclContext *getEquivalentDeclContextFromSourceFile(DeclContext *DC,
 /// returns \c true. Returns \c true if any callback call returns \c true, \c
 /// false otherwise.
 static bool
-forEachDependencyUntilTrue(CompilerInstance &CI, ModuleDecl *CurrentModule,
-                           unsigned excludeBufferID,
+forEachDependencyUntilTrue(CompilerInstance &CI, unsigned excludeBufferID,
                            llvm::function_ref<bool(StringRef)> callback) {
   // Check files in the current module.
-  for (FileUnit *file : CurrentModule->getFiles()) {
+  for (FileUnit *file : CI.getMainModule()->getFiles()) {
     StringRef filename;
     if (auto SF = dyn_cast<SourceFile>(file)) {
       if (SF->getBufferID() == excludeBufferID)
@@ -203,12 +202,11 @@ forEachDependencyUntilTrue(CompilerInstance &CI, ModuleDecl *CurrentModule,
 
 /// Collect hash codes of the dependencies into \c Map.
 static void cacheDependencyHashIfNeeded(CompilerInstance &CI,
-                                        ModuleDecl *CurrentModule,
                                         unsigned excludeBufferID,
                                         llvm::StringMap<llvm::hash_code> &Map) {
   auto &FS = CI.getFileSystem();
   forEachDependencyUntilTrue(
-      CI, CurrentModule, excludeBufferID, [&](StringRef filename) {
+      CI, excludeBufferID, [&](StringRef filename) {
         if (Map.count(filename))
           return false;
 
@@ -229,12 +227,12 @@ static void cacheDependencyHashIfNeeded(CompilerInstance &CI,
 
 /// Check if any dependent files are modified since \p timestamp.
 static bool areAnyDependentFilesInvalidated(
-    CompilerInstance &CI, ModuleDecl *CurrentModule, llvm::vfs::FileSystem &FS,
+    CompilerInstance &CI, llvm::vfs::FileSystem &FS,
     unsigned excludeBufferID, llvm::sys::TimePoint<> timestamp,
     llvm::StringMap<llvm::hash_code> &Map) {
 
   return forEachDependencyUntilTrue(
-      CI, CurrentModule, excludeBufferID, [&](StringRef filePath) {
+      CI, excludeBufferID, [&](StringRef filePath) {
         auto stat = FS.status(filePath);
         if (!stat)
           // Missing.
@@ -291,7 +289,7 @@ bool CompletionInstance::performCachedOperationIfPossible(
     return false;
 
   auto &CI = *CachedCI;
-  auto *oldSF = CI.getCodeCompletionFile().get();
+  auto *oldSF = CI.getCodeCompletionFile();
 
   auto *oldState = oldSF->getDelayedParserState();
   assert(oldState->hasCodeCompletionDelayedDeclState());
@@ -304,7 +302,7 @@ bool CompletionInstance::performCachedOperationIfPossible(
 
   if (shouldCheckDependencies()) {
     if (areAnyDependentFilesInvalidated(
-            CI, CurrentModule, *FileSystem, SM.getCodeCompletionBufferID(),
+            CI, *FileSystem, SM.getCodeCompletionBufferID(),
             DependencyCheckedTimestamp, InMemoryDependencyHash))
       return false;
     DependencyCheckedTimestamp = std::chrono::system_clock::now();
@@ -441,22 +439,21 @@ bool CompletionInstance::performCachedOperationIfPossible(
 
     // Create a new module and a source file using the current AST context.
     auto &Ctx = oldM->getASTContext();
-    auto *newM =
-        ModuleDecl::create(oldM->getName(), Ctx, oldM->getImplicitImportInfo());
+    auto *newM = ModuleDecl::createMainModule(Ctx, oldM->getName(),
+                                              oldM->getImplicitImportInfo());
     auto *newSF =
         new (Ctx) SourceFile(*newM, SourceFileKind::Main, newBufferID);
     newM->addFile(*newSF);
     newSF->enableInterfaceHash();
 
-    // Tell the compiler instance we've replaced the code completion file.
-    CI.setCodeCompletionFile(newSF);
+    // Tell the compiler instance we've replaced the main module.
+    CI.setMainModule(newM);
 
     // Re-process the whole file (parsing will be lazily triggered). Still
     // re-use imported modules.
     performImportResolution(*newSF);
     bindExtensions(*newM);
 
-    CurrentModule = newM;
     traceDC = newM;
 #ifndef NDEBUG
     const auto *reparsedState = newSF->getDelayedParserState();
@@ -483,7 +480,7 @@ bool CompletionInstance::performCachedOperationIfPossible(
   }
 
   CachedReuseCount += 1;
-  cacheDependencyHashIfNeeded(CI, CurrentModule, SM.getCodeCompletionBufferID(),
+  cacheDependencyHashIfNeeded(CI, SM.getCodeCompletionBufferID(),
                               InMemoryDependencyHash);
 
   return true;
@@ -536,8 +533,7 @@ bool CompletionInstance::performNewOperation(
     CI.performParseAndResolveImportsOnly();
 
     // If we didn't find a code completion token, bail.
-    auto completionFile = CI.getCodeCompletionFile();
-    auto *state = completionFile.get()->getDelayedParserState();
+    auto *state = CI.getCodeCompletionFile()->getDelayedParserState();
     if (!state->hasCodeCompletionDelayedDeclState())
       return true;
 
@@ -554,14 +550,13 @@ bool CompletionInstance::performNewOperation(
 void CompletionInstance::cacheCompilerInstance(
     std::unique_ptr<CompilerInstance> CI, llvm::hash_code ArgsHash) {
   CachedCI = std::move(CI);
-  CurrentModule = CachedCI->getMainModule();
   CachedArgHash = ArgsHash;
   auto now = std::chrono::system_clock::now();
   DependencyCheckedTimestamp = now;
   CachedReuseCount = 0;
   InMemoryDependencyHash.clear();
   cacheDependencyHashIfNeeded(
-      *CachedCI, CurrentModule,
+      *CachedCI,
       CachedCI->getASTContext().SourceMgr.getCodeCompletionBufferID(),
       InMemoryDependencyHash);
 }

tools/SourceKit/lib/SwiftLang/SwiftCompletion.cpp

@@ -142,8 +142,8 @@ static bool swiftCodeCompleteImpl(
         SwiftConsumer.setContext(&CI.getASTContext(), &CI.getInvocation(),
                                  &CompletionContext);
 
-        auto SF = CI.getCodeCompletionFile();
-        performCodeCompletionSecondPass(*SF.get(), *callbacksFactory);
+        auto *SF = CI.getCodeCompletionFile();
+        performCodeCompletionSecondPass(*SF, *callbacksFactory);
         SwiftConsumer.clearContext();
       });
 }

tools/SourceKit/lib/SwiftLang/SwiftConformingMethodList.cpp

@@ -42,8 +42,8 @@ static bool swiftConformingMethodListImpl(
             ide::makeConformingMethodListCallbacksFactory(ExpectedTypeNames,
                                                           Consumer));
 
-        auto SF = CI.getCodeCompletionFile();
-        performCodeCompletionSecondPass(*SF.get(), *callbacksFactory);
+        auto *SF = CI.getCodeCompletionFile();
+        performCodeCompletionSecondPass(*SF, *callbacksFactory);
       });
 }
 

tools/SourceKit/lib/SwiftLang/SwiftTypeContextInfo.cpp

@@ -39,8 +39,8 @@ static bool swiftTypeContextInfoImpl(
         std::unique_ptr<CodeCompletionCallbacksFactory> callbacksFactory(
             ide::makeTypeContextInfoCallbacksFactory(Consumer));
 
-        auto SF = CI.getCodeCompletionFile();
-        performCodeCompletionSecondPass(*SF.get(), *callbacksFactory);
+        auto *SF = CI.getCodeCompletionFile();
+        performCodeCompletionSecondPass(*SF, *callbacksFactory);
       });
 }