[CodeComplete] Store the compiler instace as a shared pointer in CompletionInstance

ahoppen · ahoppen · commit 1150e700a53e · 2022-07-21T18:29:34.000+02:00
diff --git a/include/swift/IDE/CompletionInstance.h b/include/swift/IDE/CompletionInstance.h
@@ -46,7 +46,7 @@ makeCodeCompletionMemoryBuffer(const llvm::MemoryBuffer *origBuf,
 /// The result returned via the callback from the perform*Operation methods.
 struct CompletionInstanceResult {
   /// The compiler instance that is prepared for the second pass.
-  CompilerInstance &CI;
+  std::shared_ptr<CompilerInstance> CI;
   /// Whether an AST was reused.
   bool DidReuseAST;
   /// Whether the CompletionInstance found a code completion token in the source
@@ -87,14 +87,14 @@ class CompletionInstance {
 
   std::mutex mtx;
 
-  std::unique_ptr<CompilerInstance> CachedCI;
+  std::shared_ptr<CompilerInstance> CachedCI;
   llvm::hash_code CachedArgHash;
   llvm::sys::TimePoint<> DependencyCheckedTimestamp;
   llvm::StringMap<llvm::hash_code> InMemoryDependencyHash;
   unsigned CachedReuseCount = 0;
   std::atomic<bool> CachedCIShouldBeInvalidated;
 
-  void cacheCompilerInstance(std::unique_ptr<CompilerInstance> CI,
+  void cacheCompilerInstance(std::shared_ptr<CompilerInstance> CI,
                              llvm::hash_code ArgsHash);
 
   bool shouldCheckDependencies() const;
diff --git a/include/swift/IDE/SwiftCompletionInfo.h b/include/swift/IDE/SwiftCompletionInfo.h
@@ -20,8 +20,7 @@ namespace swift {
 namespace ide {
 
 struct SwiftCompletionInfo {
-  swift::ASTContext *swiftASTContext = nullptr;
-  const swift::CompilerInvocation *invocation = nullptr;
+  std::shared_ptr<CompilerInstance> compilerInstance = nullptr;
   CodeCompletionContext *completionContext = nullptr;
 };
 
diff --git a/lib/IDE/CompletionInstance.cpp b/lib/IDE/CompletionInstance.cpp
@@ -197,15 +197,14 @@ bool CompletionInstance::performCachedOperationIfPossible(
   if (CachedArgHash != ArgsHash)
     return false;
 
-  auto &CI = *CachedCI;
-  auto *oldSF = CI.getCodeCompletionFile();
+  auto *oldSF = CachedCI->getCodeCompletionFile();
   assert(oldSF->getBufferID());
 
   auto *oldState = oldSF->getDelayedParserState();
   assert(oldState->hasCodeCompletionDelayedDeclState());
   auto &oldInfo = oldState->getCodeCompletionDelayedDeclState();
 
-  auto &SM = CI.getSourceMgr();
+  auto &SM = CachedCI->getSourceMgr();
   auto bufferName = completionBuffer->getBufferIdentifier();
   if (SM.getIdentifierForBuffer(*oldSF->getBufferID()) != bufferName)
     return false;
@@ -222,7 +221,7 @@ bool CompletionInstance::performCachedOperationIfPossible(
     }
 
     if (areAnyDependentFilesInvalidated(
-            CI, *FileSystem, *oldSF->getBufferID(),
+            *CachedCI, *FileSystem, *oldSF->getBufferID(),
             DependencyCheckedTimestamp, InMemoryDependencyHash))
       return false;
     DependencyCheckedTimestamp = std::chrono::system_clock::now();
@@ -233,10 +232,10 @@ bool CompletionInstance::performCachedOperationIfPossible(
   auto tmpBufferID = tmpSM.addMemBufferCopy(completionBuffer);
   tmpSM.setCodeCompletionPoint(tmpBufferID, Offset);
 
-  LangOptions langOpts = CI.getASTContext().LangOpts;
-  TypeCheckerOptions typeckOpts = CI.getASTContext().TypeCheckerOpts;
-  SILOptions silOpts = CI.getASTContext().SILOpts;
-  SearchPathOptions searchPathOpts = CI.getASTContext().SearchPathOpts;
+  LangOptions langOpts = CachedCI->getASTContext().LangOpts;
+  TypeCheckerOptions typeckOpts = CachedCI->getASTContext().TypeCheckerOpts;
+  SILOptions silOpts = CachedCI->getASTContext().SILOpts;
+  SearchPathOptions searchPathOpts = CachedCI->getASTContext().SearchPathOpts;
   DiagnosticEngine tmpDiags(tmpSM);
   ClangImporterOptions clangOpts;
   symbolgraphgen::SymbolGraphOptions symbolOpts;
@@ -382,7 +381,7 @@ bool CompletionInstance::performCachedOperationIfPossible(
     newM->addFile(*newSF);
 
     // Tell the compiler instance we've replaced the main module.
-    CI.setMainModule(newM);
+    CachedCI->setMainModule(newM);
 
     // Re-process the whole file (parsing will be lazily triggered). Still
     // re-use imported modules.
@@ -407,22 +406,22 @@ bool CompletionInstance::performCachedOperationIfPossible(
 
     // The diagnostic engine is keeping track of state which might modify
     // parsing and type checking behaviour. Clear the flags.
-    CI.getDiags().resetHadAnyError();
-    CI.getASTContext().CancellationFlag = CancellationFlag;
+    CachedCI->getDiags().resetHadAnyError();
+    CachedCI->getASTContext().CancellationFlag = CancellationFlag;
 
     if (DiagC)
-      CI.addDiagnosticConsumer(DiagC);
+      CachedCI->addDiagnosticConsumer(DiagC);
 
     if (CancellationFlag && CancellationFlag->load(std::memory_order_relaxed)) {
       Callback(CancellableResult<CompletionInstanceResult>::cancelled());
     } else {
       Callback(CancellableResult<CompletionInstanceResult>::success(
-          {CI, /*reusingASTContext=*/true,
+          {CachedCI, /*reusingASTContext=*/true,
            /*DidFindCodeCompletionToken=*/true}));
     }
 
     if (DiagC)
-      CI.removeDiagnosticConsumer(DiagC);
+      CachedCI->removeDiagnosticConsumer(DiagC);
   }
 
   CachedReuseCount += 1;
@@ -443,45 +442,44 @@ void CompletionInstance::performNewOperation(
   // If ArgsHash is None we shouldn't cache the compiler instance.
   bool ShouldCacheCompilerInstance = ArgsHash.hasValue();
 
-  auto TheInstance = std::make_unique<CompilerInstance>();
+  auto CI = std::make_shared<CompilerInstance>();
 
   // Track non-system dependencies in fast-completion mode to invalidate the
   // compiler instance if any dependent files are modified.
   Invocation.getFrontendOptions().IntermoduleDependencyTracking =
       IntermoduleDepTrackingMode::ExcludeSystem;
 
   {
-    auto &CI = *TheInstance;
     if (DiagC)
-      CI.addDiagnosticConsumer(DiagC);
+      CI->addDiagnosticConsumer(DiagC);
 
     SWIFT_DEFER {
       if (DiagC)
-        CI.removeDiagnosticConsumer(DiagC);
+        CI->removeDiagnosticConsumer(DiagC);
     };
 
     if (FileSystem != llvm::vfs::getRealFileSystem())
-      CI.getSourceMgr().setFileSystem(FileSystem);
+      CI->getSourceMgr().setFileSystem(FileSystem);
 
     Invocation.setCodeCompletionPoint(completionBuffer, Offset);
 
     std::string InstanceSetupError;
-    if (CI.setup(Invocation, InstanceSetupError)) {
+    if (CI->setup(Invocation, InstanceSetupError)) {
       Callback(CancellableResult<CompletionInstanceResult>::failure(
           InstanceSetupError));
       return;
     }
-    CI.getASTContext().CancellationFlag = CancellationFlag;
-    registerIDERequestFunctions(CI.getASTContext().evaluator);
+    CI->getASTContext().CancellationFlag = CancellationFlag;
+    registerIDERequestFunctions(CI->getASTContext().evaluator);
 
-    CI.performParseAndResolveImportsOnly();
+    CI->performParseAndResolveImportsOnly();
 
-    bool DidFindCodeCompletionToken = CI.getCodeCompletionFile()
+    bool DidFindCodeCompletionToken = CI->getCodeCompletionFile()
                                           ->getDelayedParserState()
                                           ->hasCodeCompletionDelayedDeclState();
     ShouldCacheCompilerInstance &= DidFindCodeCompletionToken;
 
-    auto CancellationFlag = CI.getASTContext().CancellationFlag;
+    auto CancellationFlag = CI->getASTContext().CancellationFlag;
     if (CancellationFlag && CancellationFlag->load(std::memory_order_relaxed)) {
       Callback(CancellableResult<CompletionInstanceResult>::cancelled());
       // The completion instance may be in an invalid state when it's been
@@ -502,11 +500,11 @@ void CompletionInstance::performNewOperation(
   // because performCachedOperationIfPossible wouldn't have an old code
   // completion state to compare the new one to.
   if (ShouldCacheCompilerInstance)
-    cacheCompilerInstance(std::move(TheInstance), *ArgsHash);
+    cacheCompilerInstance(std::move(CI), *ArgsHash);
 }
 
 void CompletionInstance::cacheCompilerInstance(
-    std::unique_ptr<CompilerInstance> CI, llvm::hash_code ArgsHash) {
+    std::shared_ptr<CompilerInstance> CI, llvm::hash_code ArgsHash) {
   CachedCI = std::move(CI);
   CachedArgHash = ArgsHash;
   auto now = std::chrono::system_clock::now();
@@ -602,11 +600,9 @@ void swift::ide::CompletionInstance::codeComplete(
         : ImportDep(ImportDep), CancellationFlag(CancellationFlag),
           Callback(Callback) {}
 
-    void setContext(swift::ASTContext *context,
-                    const swift::CompilerInvocation *invocation,
+    void setContext(std::shared_ptr<CompilerInstance> compilerInstance,
                     swift::ide::CodeCompletionContext *completionContext) {
-      SwiftContext.swiftASTContext = context;
-      SwiftContext.invocation = invocation;
+      SwiftContext.compilerInstance = std::move(compilerInstance);
       SwiftContext.completionContext = completionContext;
     }
     void clearContext() { SwiftContext = SwiftCompletionInfo(); }
@@ -617,7 +613,7 @@ void swift::ide::CompletionInstance::codeComplete(
           CancellationFlag->load(std::memory_order_relaxed)) {
         Callback(ResultType::cancelled());
       } else {
-        assert(SwiftContext.swiftASTContext);
+        assert(SwiftContext.compilerInstance);
         Callback(ResultType::success({context.getResultSink(), SwiftContext, ImportDep}));
       }
     }
@@ -631,9 +627,9 @@ void swift::ide::CompletionInstance::codeComplete(
             [&CompletionContext, &CancellationFlag](auto &Result,
                                                     auto DeliverTransformed) {
               CompletionContext.ReusingASTContext = Result.DidReuseAST;
-              CompilerInstance &CI = Result.CI;
-              ImportDepth ImportDep{CI.getASTContext(),
-                                    CI.getInvocation().getFrontendOptions()};
+              std::shared_ptr<CompilerInstance> CI = Result.CI;
+              ImportDepth ImportDep{CI->getASTContext(),
+                                    CI->getInvocation().getFrontendOptions()};
               ConsumerToCallbackAdapter Consumer(ImportDep, CancellationFlag,
                                                  DeliverTransformed);
 
@@ -642,27 +638,22 @@ void swift::ide::CompletionInstance::codeComplete(
                                                           Consumer));
 
               if (!Result.DidFindCodeCompletionToken) {
-                SwiftCompletionInfo Info{&CI.getASTContext(),
-                                         &CI.getInvocation(),
-                                         &CompletionContext};
+                SwiftCompletionInfo Info{CI, &CompletionContext};
                 CodeCompletionResultSink ResultSink;
                 DeliverTransformed(ResultType::success({ResultSink, Info, ImportDep}));
                 return;
               }
 
-              Consumer.setContext(&CI.getASTContext(), &CI.getInvocation(),
-                                  &CompletionContext);
-              performCodeCompletionSecondPass(*CI.getCodeCompletionFile(),
+              Consumer.setContext(CI, &CompletionContext);
+              performCodeCompletionSecondPass(*CI->getCodeCompletionFile(),
                                               *callbacksFactory);
               Consumer.clearContext();
               if (!Consumer.HandleResultsCalled) {
                 // If we didn't receive a handleResult call from the second
                 // pass, we didn't receive any results. To make sure Callback
                 // gets called exactly once, call it manually with no results
                 // here.
-                SwiftCompletionInfo Info{&CI.getASTContext(),
-                                         &CI.getInvocation(),
-                                         &CompletionContext};
+                SwiftCompletionInfo Info{CI, &CompletionContext};
                 CodeCompletionResultSink ResultSink;
                 DeliverTransformed(ResultType::success({ResultSink, Info, ImportDep}));
               }
@@ -724,7 +715,7 @@ void swift::ide::CompletionInstance::typeContextInfo(
               }
 
               performCodeCompletionSecondPass(
-                  *Result.CI.getCodeCompletionFile(), *callbacksFactory);
+                  *Result.CI->getCodeCompletionFile(), *callbacksFactory);
               if (!Consumer.HandleResultsCalled) {
                 // If we didn't receive a handleResult call from the second
                 // pass, we didn't receive any results. To make sure Callback
@@ -792,7 +783,7 @@ void swift::ide::CompletionInstance::conformingMethodList(
               }
 
               performCodeCompletionSecondPass(
-                  *Result.CI.getCodeCompletionFile(), *callbacksFactory);
+                  *Result.CI->getCodeCompletionFile(), *callbacksFactory);
               if (!Consumer.HandleResultsCalled) {
                 // If we didn't receive a handleResult call from the second
                 // pass, we didn't receive any results. To make sure Callback
diff --git a/tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp b/tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.cpp
@@ -76,10 +76,11 @@ std::vector<Completion *> SourceKit::CodeCompletion::extendCompletions(
     const Options &options, Completion *prefix, bool clearFlair) {
 
   ImportDepth depth;
-  if (info.swiftASTContext) {
+  if (info.compilerInstance) {
     // Build import depth map.
-    depth = ImportDepth(*info.swiftASTContext,
-                        info.invocation->getFrontendOptions());
+    depth = ImportDepth(
+        info.compilerInstance->getASTContext(),
+        info.compilerInstance->getInvocation().getFrontendOptions());
   }
 
   if (info.completionContext)
diff --git a/tools/swift-ide-test/swift-ide-test.cpp b/tools/swift-ide-test/swift-ide-test.cpp
@@ -1168,7 +1168,7 @@ static int printCodeCompletionResults(
         printCodeCompletionResultsImpl(
             Result.ResultSink.Results, llvm::outs(), IncludeKeywords,
             IncludeComments, IncludeSourceText, PrintAnnotatedDescription,
-            *Result.Info.swiftASTContext);
+            Result.Info.compilerInstance->getASTContext());
         return 0;
       });
 }
@@ -1547,7 +1547,8 @@ static int doBatchCodeCompletion(const CompilerInvocation &InitInvok,
             printCodeCompletionResultsImpl(
                 Result->ResultSink.Results, OS, IncludeKeywords,
                 IncludeComments, IncludeSourceText,
-                CodeCompletionAnnotateResults, *Result->Info.swiftASTContext);
+                CodeCompletionAnnotateResults,
+                Result->Info.compilerInstance->getASTContext());
             break;
           }
           case CancellableResultKind::Failure:
