[SourceKit] Move invocation of code completion second pass for TypeContextInfo from SoruceKit to CompletionInstance

The invocation of the code completion second pass should be implementation detail of `CompletionInstance`. Create a method on `CompletionInstance` that correctly invokes the second pass and just reutnrs the type context info results to the caller.

Author: ahoppen

include/swift/IDE/CancellableResult.h

@@ -140,6 +140,36 @@ class CancellableResult {
     assert(getKind() == CancellableResultKind::Failure);
     return Error;
   }
+
+  /// If the result represents success, invoke \p Transform to asynchronously
+  /// transform the wrapped result type and produce a new result type that is
+  /// provided by the callback function passed to \p Transform. Afterwards call
+  /// \p Handle with either the transformed value or the failure or cancelled
+  /// result.
+  /// The \c async part of the map means that the transform might happen
+  /// asyncronously. This function does not introduce asynchronicity by itself.
+  /// \p Transform might also invoke the callback synchronously.
+  template <typename NewResultType>
+  void
+  mapAsync(llvm::function_ref<
+               void(ResultType &,
+                    llvm::function_ref<void(CancellableResult<NewResultType>)>)>
+               Transform,
+           llvm::function_ref<void(CancellableResult<NewResultType>)> Handle) {
+    switch (getKind()) {
+    case CancellableResultKind::Success:
+      Transform(getResult(), [&](CancellableResult<NewResultType> NewResult) {
+        Handle(NewResult);
+      });
+      break;
+    case CancellableResultKind::Failure:
+      Handle(CancellableResult<NewResultType>::failure(getError()));
+      break;
+    case CancellableResultKind::Cancelled:
+      Handle(CancellableResult<NewResultType>::cancelled());
+      break;
+    }
+  }
 };
 
 } // namespace ide

include/swift/IDE/CompletionInstance.h

@@ -15,6 +15,7 @@
 
 #include "swift/Frontend/Frontend.h"
 #include "swift/IDE/CancellableResult.h"
+#include "swift/IDE/TypeContextInfo.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringRef.h"
@@ -48,6 +49,14 @@ struct CompletionInstanceResult {
   bool DidFindCodeCompletionToken;
 };
 
+/// The results returned from \c CompletionInstance::typeContextInfo.
+struct TypeContextInfoResult {
+  /// The actual results. If empty, no results were found.
+  ArrayRef<ide::TypeContextInfoItem> Results;
+  /// Whether an AST was reused to produce the results.
+  bool DidReuseAST;
+};
+
 /// Manages \c CompilerInstance for completion like operations.
 class CompletionInstance {
   struct Options {
@@ -124,6 +133,14 @@ class CompletionInstance {
       DiagnosticConsumer *DiagC,
       llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
           Callback);
+
+  void typeContextInfo(
+      swift::CompilerInvocation &Invocation, llvm::ArrayRef<const char *> Args,
+      llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+      llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+      DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CancellableResult<TypeContextInfoResult>)>
+          Callback);
 };
 
 } // namespace ide

lib/IDE/CompletionInstance.cpp

@@ -655,3 +655,59 @@ void swift::ide::CompletionInstance::performOperation(
   performNewOperation(ArgsHash, Invocation, FileSystem, completionBuffer,
                       Offset, DiagC, Callback);
 }
+
+void swift::ide::CompletionInstance::typeContextInfo(
+    swift::CompilerInvocation &Invocation, llvm::ArrayRef<const char *> Args,
+    llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+    llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+    DiagnosticConsumer *DiagC,
+    llvm::function_ref<void(CancellableResult<TypeContextInfoResult>)>
+        Callback) {
+  using ResultType = CancellableResult<TypeContextInfoResult>;
+
+  struct ConsumerToCallbackAdapter : public ide::TypeContextInfoConsumer {
+    bool ReusingASTContext;
+    llvm::function_ref<void(ResultType)> Callback;
+    bool HandleResultsCalled = false;
+
+    ConsumerToCallbackAdapter(bool ReusingASTContext,
+                              llvm::function_ref<void(ResultType)> Callback)
+        : ReusingASTContext(ReusingASTContext), Callback(Callback) {}
+
+    void handleResults(ArrayRef<ide::TypeContextInfoItem> Results) override {
+      HandleResultsCalled = true;
+      Callback(ResultType::success({Results, ReusingASTContext}));
+    }
+  };
+
+  performOperation(
+      Invocation, Args, FileSystem, completionBuffer, Offset, DiagC,
+      [&](CancellableResult<CompletionInstanceResult> CIResult) {
+        CIResult.mapAsync<TypeContextInfoResult>(
+            [](auto &Result, auto DeliverTransformed) {
+              ConsumerToCallbackAdapter Consumer(Result.DidReuseAST,
+                                                 DeliverTransformed);
+              std::unique_ptr<CodeCompletionCallbacksFactory> callbacksFactory(
+                  ide::makeTypeContextInfoCallbacksFactory(Consumer));
+
+              if (!Result.DidFindCodeCompletionToken) {
+                // Deliver empty results if we didn't find a code completion
+                // token.
+                DeliverTransformed(
+                    ResultType::success({/*Results=*/{}, Result.DidReuseAST}));
+              }
+
+              performCodeCompletionSecondPass(
+                  *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
+                // DeliverTransformed gets called exactly once, call it with
+                // no results here.
+                DeliverTransformed(
+                    ResultType::success({/*Results=*/{}, Result.DidReuseAST}));
+              }
+            },
+            Callback);
+      });
+}

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.cpp

@@ -1000,12 +1000,12 @@ SwiftLangSupport::getFileSystem(const Optional<VFSOptions> &vfsOptions,
   return llvm::vfs::getRealFileSystem();
 }
 
-void SwiftLangSupport::performCompletionLikeOperation(
+void SwiftLangSupport::performWithParamsToCompletionLikeOperation(
     llvm::MemoryBuffer *UnresolvedInputFile, unsigned Offset,
     ArrayRef<const char *> Args,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
-    llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
-        Callback) {
+    llvm::function_ref<void(CancellableResult<CompletionLikeOperationParams>)>
+        PerformOperation) {
   assert(FileSystem);
 
   // Resolve symlinks for the input file; we resolve them for the input files
@@ -1050,21 +1050,43 @@ void SwiftLangSupport::performCompletionLikeOperation(
       Invocation, Args, Diags, newBuffer->getBufferIdentifier(), FileSystem,
       CompilerInvocationError);
   if (CreatingInvocationFailed) {
-    Callback(CancellableResult<CompletionInstanceResult>::failure(
+    PerformOperation(CancellableResult<CompletionLikeOperationParams>::failure(
         CompilerInvocationError));
     return;
   }
   if (!Invocation.getFrontendOptions().InputsAndOutputs.hasInputs()) {
-    Callback(CancellableResult<CompletionInstanceResult>::failure(
+    PerformOperation(CancellableResult<CompletionLikeOperationParams>::failure(
         "no input filenames specified"));
     return;
   }
 
   // Pin completion instance.
   auto CompletionInst = getCompletionInstance();
 
-  CompletionInst->performOperation(Invocation, Args, FileSystem,
-                                   newBuffer.get(), Offset, &CIDiags, Callback);
+  CompletionLikeOperationParams Params = {Invocation, newBuffer.get(),
+                                          &CIDiags};
+  PerformOperation(
+      CancellableResult<CompletionLikeOperationParams>::success(Params));
+}
+
+void SwiftLangSupport::performCompletionLikeOperation(
+    llvm::MemoryBuffer *UnresolvedInputFile, unsigned Offset,
+    ArrayRef<const char *> Args,
+    llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+    llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
+        Callback) {
+  performWithParamsToCompletionLikeOperation(
+      UnresolvedInputFile, Offset, Args, FileSystem,
+      [&](CancellableResult<CompletionLikeOperationParams> ParamsResult) {
+        ParamsResult.mapAsync<CompletionInstanceResult>(
+            [&](auto &CIParams, auto DeliverTransformed) {
+              getCompletionInstance()->performOperation(
+                  CIParams.Invocation, Args, FileSystem,
+                  CIParams.completionBuffer, Offset, CIParams.DiagC,
+                  DeliverTransformed);
+            },
+            Callback);
+      });
 }
 
 CloseClangModuleFiles::~CloseClangModuleFiles() {

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -20,6 +20,7 @@
 #include "SourceKit/Support/ThreadSafeRefCntPtr.h"
 #include "SourceKit/Support/Tracing.h"
 #include "SwiftInterfaceGenContext.h"
+#include "swift/AST/DiagnosticConsumer.h"
 #include "swift/Basic/ThreadSafeRefCounted.h"
 #include "swift/IDE/CancellableResult.h"
 #include "swift/IDE/CompletionInstance.h"
@@ -472,6 +473,23 @@ class SwiftLangSupport : public LangSupport {
           swift::ide::CancellableResult<swift::ide::CompletionInstanceResult>)>
           Callback);
 
+  /// The result returned from \c performWithParamsToCompletionLikeOperation.
+  struct CompletionLikeOperationParams {
+    swift::CompilerInvocation &Invocation;
+    llvm::MemoryBuffer *completionBuffer;
+    swift::DiagnosticConsumer *DiagC;
+  };
+
+  /// Execute \p PerformOperation sychronously with the parameters necessary to
+  /// invoke a completion-like operation on \c CompletionInstance.
+  void performWithParamsToCompletionLikeOperation(
+      llvm::MemoryBuffer *UnresolvedInputFile, unsigned Offset,
+      ArrayRef<const char *> Args,
+      llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+      llvm::function_ref<
+          void(swift::ide::CancellableResult<CompletionLikeOperationParams>)>
+          PerformOperation);
+
   //==========================================================================//
   // LangSupport Interface
   //==========================================================================//