[SourceKit/CodeCompletion] Use callback function to run the second pass

To controls the lifetime of CompilerInstance within CompletionIntance.

- Prevent from using the same CompilerInstance from multiple completion
  requests
- Separate the stacktrace between "fast" and "normal" completions
- Further code consolidation between various completion-like requests

Author: rintaro

include/swift/IDE/CompletionInstance.h

@@ -37,47 +37,56 @@ makeCodeCompletionMemoryBuffer(const llvm::MemoryBuffer *origBuf,
 
 /// Manages \c CompilerInstance for completion like operations.
 class CompletionInstance {
-  std::unique_ptr<CompilerInstance> CachedCI = nullptr;
-  llvm::hash_code CachedArgsHash;
-  bool EnableASTCaching = false;
   unsigned MaxASTReuseCount = 100;
+  bool EnableASTCaching = false;
+
+  std::shared_ptr<CompilerInstance> CachedCI = nullptr;
+  llvm::hash_code CachedArgsHash;
   unsigned CurrentASTReuseCount = 0;
 
-  /// Returns cached \c CompilerInstance if it's usable for the specified
-  /// completion request.
-  /// Returns \c nullptr if the functionality is disabled, compiler argument has
+  /// Calls \p Callback with cached \c CompilerInstance if it's usable for the
+  /// specified completion request.
+  /// Returns \c if the callback was called. Returns \c false if the
+  /// functionality is disabled, compiler argument has
   /// changed, primary file is not the same, the \c Offset is not in function
   /// bodies, or the interface hash of the file has changed.
-  swift::CompilerInstance *
-  getCachedCompilerInstance(const swift::CompilerInvocation &Invocation,
-                            llvm::hash_code ArgsHash,
-                            llvm::MemoryBuffer *completionBuffer,
-                            unsigned int Offset, DiagnosticConsumer *DiagC);
+  bool performCachedOperaitonIfPossible(
+      const swift::CompilerInvocation &Invocation, llvm::hash_code ArgsHash,
+      llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+      DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CompilerInstance &)> Callback);
 
-  /// Returns new \c CompilerInstance for the completion request. Users still
-  /// have to call \c performParseAndResolveImportsOnly() , and perform the
-  /// second pass on it.
-  swift::CompilerInstance *renewCompilerInstance(
+  /// Calls \p Callback with new \c CompilerInstance for the completion
+  /// request. The \c CompilerInstace passed to the callback already performed
+  /// the first pass.
+  /// Returns \c false if it fails to setup the \c CompilerInstance.
+  bool performNewOperation(
       swift::CompilerInvocation &Invocation, llvm::hash_code ArgsHash,
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
-      std::string &Error, DiagnosticConsumer *DiagC);
+      std::string &Error, DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CompilerInstance &)> Callback);
 
 public:
   void setEnableASTCaching(bool Flag) { EnableASTCaching = Flag; }
 
-  /// Returns \C CompilerInstance for the completion request. Users can check if
-  /// it's cached or not by 'hasPersistentParserState()'.
+  /// Calls \p Callback with a \c CompilerInstance which is prepared for the
+  /// second pass. \p Callback is resposible to perform the second pass on it.
+  /// The \c CompilerInstance may be reused from the previous completions,
+  /// and may be cached for the next completion.
+  /// Return \c true if \p is successfully called, \c it fails. In failure
+  /// cases \p Error is populated with an error message.
   ///
   /// NOTE: \p Args is only used for checking the equaity of the invocation.
   /// Since this function assumes that it is already normalized, exact the same
   /// arguments including their order is considered as the same invocation.
-  swift::CompilerInstance *getCompilerInstance(
-      swift::CompilerInvocation &Invocation,
-      llvm::ArrayRef<const char *> Args,
-      llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
-      llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
-      std::string &Error, DiagnosticConsumer *DiagC);
+  bool
+  performOperation(swift::CompilerInvocation &Invocation,
+                   llvm::ArrayRef<const char *> Args,
+                   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+                   llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+                   std::string &Error, DiagnosticConsumer *DiagC,
+                   llvm::function_ref<void(CompilerInstance &)> Callback);
 };
 
 } // namespace ide

lib/IDE/CompletionInstance.cpp

@@ -129,38 +129,40 @@ static DeclContext *getEquivalentDeclContextFromSourceFile(DeclContext *DC,
 
 } // namespace
 
-CompilerInstance *CompletionInstance::getCachedCompilerInstance(
+bool CompletionInstance::performCachedOperaitonIfPossible(
     const swift::CompilerInvocation &Invocation, llvm::hash_code ArgsHash,
     llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
-    DiagnosticConsumer *DiagC) {
-  if (!EnableASTCaching)
-    return nullptr;
+    DiagnosticConsumer *DiagC,
+    llvm::function_ref<void(CompilerInstance &)> Callback) {
 
-  if (CurrentASTReuseCount >= MaxASTReuseCount) {
-    CurrentASTReuseCount = 0;
-    return nullptr;
-  }
+  if (!EnableASTCaching)
+    return false;
 
-  if (!CachedCI)
-    return nullptr;
+  // Temporary move the CI so other threads don't use the same instance.
+  std::shared_ptr<CompilerInstance> CI;
+  CI.swap(CachedCI);
 
+  if (!CI)
+    return false;
+  if (CurrentASTReuseCount >= MaxASTReuseCount)
+    return false;
   if (ArgsHash != CachedArgsHash)
-    return nullptr;
+    return false;
 
-  auto &oldState = CachedCI->getPersistentParserState();
+  auto &oldState = CI->getPersistentParserState();
   if (!oldState.hasCodeCompletionDelayedDeclState())
-    return nullptr;
+    return false;
 
-  auto &SM = CachedCI->getSourceMgr();
+  auto &SM = CI->getSourceMgr();
   if (SM.getIdentifierForBuffer(SM.getCodeCompletionBufferID()) !=
       completionBuffer->getBufferIdentifier())
-    return nullptr;
+    return false;
 
   auto &oldInfo = oldState.getCodeCompletionDelayedDeclState();
 
   // Currently, only completions within a function body is supported.
   if (oldInfo.Kind != CodeCompletionDelayedDeclKind::FunctionBody)
-    return nullptr;
+    return false;
 
   auto newBufferID = SM.addMemBufferCopy(completionBuffer);
   SM.setCodeCompletionPoint(newBufferID, Offset);
@@ -185,13 +187,13 @@ CompilerInstance *CompletionInstance::getCachedCompilerInstance(
   parseIntoSourceFileFull(*newSF, BufferID, &newState);
   // Couldn't find any completion token?
   if (!newState.hasCodeCompletionDelayedDeclState())
-    return nullptr;
+    return false;
 
   auto &newInfo = newState.getCodeCompletionDelayedDeclState();
 
   // The new completion must happens in function body too.
   if (newInfo.Kind != CodeCompletionDelayedDeclKind::FunctionBody)
-    return nullptr;
+    return false;
 
   auto *oldSF = oldInfo.ParentContext->getParentSourceFile();
 
@@ -201,12 +203,12 @@ CompilerInstance *CompletionInstance::getCachedCompilerInstance(
   oldSF->getInterfaceHash(oldInterfaceHash);
   newSF->getInterfaceHash(newInterfaceHash);
   if (oldInterfaceHash != newInterfaceHash)
-    return nullptr;
+    return false;
 
   DeclContext *DC =
       getEquivalentDeclContextFromSourceFile(newInfo.ParentContext, oldSF);
   if (!DC)
-    return nullptr;
+    return false;
 
   // OK, we can perform fast completion for this. Update the orignal delayed
   // decl state.
@@ -228,28 +230,33 @@ CompilerInstance *CompletionInstance::getCachedCompilerInstance(
   if (AFD->isBodySkipped())
     AFD->setBodyDelayed(AFD->getBodySourceRange());
   if (DiagC)
-    CachedCI->addDiagnosticConsumer(DiagC);
+    CI->addDiagnosticConsumer(DiagC);
 
-  CachedCI->getDiags().diagnose(
-      SM.getLocForOffset(BufferID, newInfo.StartOffset),
-      diag::completion_reusing_astcontext);
+  CI->getDiags().diagnose(SM.getLocForOffset(BufferID, newInfo.StartOffset),
+                          diag::completion_reusing_astcontext);
 
+  Callback(*CI);
+
+  if (DiagC)
+    CI->removeDiagnosticConsumer(DiagC);
+
+  CachedCI.swap(CI);
+  CachedArgsHash = ArgsHash;
   CurrentASTReuseCount += 1;
 
-  return CachedCI.get();
+  return true;
 }
 
-CompilerInstance *CompletionInstance::renewCompilerInstance(
+bool CompletionInstance::performNewOperation(
     swift::CompilerInvocation &Invocation, llvm::hash_code ArgsHash,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
     llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
-    std::string &Error, DiagnosticConsumer *DiagC) {
+    std::string &Error, DiagnosticConsumer *DiagC,
+    llvm::function_ref<void(CompilerInstance &)> Callback) {
   CachedCI.reset();
-  CachedCI = std::make_unique<CompilerInstance>();
-  CachedArgsHash = ArgsHash;
-  auto *CI = CachedCI.get();
+  auto CI = std::make_shared<CompilerInstance>();
   if (DiagC)
-    CachedCI->addDiagnosticConsumer(DiagC);
+    CI->addDiagnosticConsumer(DiagC);
 
   if (FileSystem != llvm::vfs::getRealFileSystem())
     CI->getSourceMgr().setFileSystem(FileSystem);
@@ -258,18 +265,31 @@ CompilerInstance *CompletionInstance::renewCompilerInstance(
 
   if (CI->setup(Invocation)) {
     Error = "failed to setup compiler instance";
-    return nullptr;
+    return false;
   }
   registerIDERequestFunctions(CI->getASTContext().evaluator);
 
-  return CI;
+  CI->performParseAndResolveImportsOnly();
+  Callback(*CI);
+
+  if (DiagC)
+    CI->removeDiagnosticConsumer(DiagC);
+
+  if (EnableASTCaching) {
+    CachedCI.swap(CI);
+    CachedArgsHash = ArgsHash;
+    CurrentASTReuseCount = 0;
+  }
+
+  return true;
 }
 
-CompilerInstance *swift::ide::CompletionInstance::getCompilerInstance(
+bool swift::ide::CompletionInstance::performOperation(
     swift::CompilerInvocation &Invocation, llvm::ArrayRef<const char *> Args,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
     llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
-    std::string &Error, DiagnosticConsumer *DiagC) {
+    std::string &Error, DiagnosticConsumer *DiagC,
+    llvm::function_ref<void(CompilerInstance &)> Callback) {
 
   // Always disable source location resolutions from .swiftsourceinfo file
   // because they're somewhat heavy operations and aren't needed for completion.
@@ -287,15 +307,14 @@ CompilerInstance *swift::ide::CompletionInstance::getCompilerInstance(
   for (auto arg : Args)
     ArgsHash = llvm::hash_combine(ArgsHash, StringRef(arg));
 
-  if (auto *cached = getCachedCompilerInstance(Invocation, ArgsHash,
-                                               completionBuffer, Offset, DiagC))
-    return cached;
+  if (performCachedOperaitonIfPossible(Invocation, ArgsHash, completionBuffer,
+                                       Offset, DiagC, Callback))
+    return true;
 
-  if (auto *renewed =
-          renewCompilerInstance(Invocation, ArgsHash, FileSystem,
-                                completionBuffer, Offset, Error, DiagC))
-    return renewed;
+  if (performNewOperation(Invocation, ArgsHash, FileSystem, completionBuffer,
+                          Offset, Error, DiagC, Callback))
+    return true;
 
   assert(!Error.empty());
-  return nullptr;
+  return false;
 }

test/SourceKit/CodeComplete/complete_sequence_race.swift

@@ -0,0 +1,44 @@
+class Foo {
+  var x: Int
+  var y: Int
+  func fooMethod() {}
+}
+struct Bar {
+  var a: Int
+  var b: Int
+  func barMethod() {}
+}
+func foo(arg: Foo) {
+  _ = arg.
+}
+func bar(arg: Bar) {
+  _ = arg.
+}
+
+// NOTE: Test for simultaneous completion requests don't cause compiler crashes.
+
+// ReuseASTContext disabled.
+// RUN: %sourcekitd-test \
+// RUN:   -req=complete -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=15:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=15:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=15:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=15:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -pos=15:11 %s -async -- %s
+
+// ReuseASTContext enabled.
+// RUN: %sourcekitd-test \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=15:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=15:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=15:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=15:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=12:11 %s -async -- %s == \
+// RUN:   -req=complete -req-opts=reuseastcontext=1 -pos=15:11 %s -async -- %s