[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.

include/swift/IDE/CancellableResult.h

@@ -140,6 +140,36 @@ public:
     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 @@ public:
       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,12 +1050,12 @@ 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;
   }
@@ -1063,8 +1063,30 @@ void SwiftLangSupport::performCompletionLikeOperation(
   // 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 @@ public:
           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
   //==========================================================================//

tools/SourceKit/lib/SwiftLang/SwiftTypeContextInfo.cpp

@@ -30,100 +30,13 @@ static void translateTypeContextInfoOptions(OptionsDictionary &from,
   // TypeContextInfo doesn't receive any options at this point.
 }
 
-/// The results returned from \c swiftTypeContextInfoImpl via \c Callback.
-struct SwiftTypeContextInfoImplResult {
-  /// The actual results. If empty, no results were found.
-  ArrayRef<ide::TypeContextInfoItem> Results;
-  /// Whether an AST was reused to produce the results.
-  bool DidReuseAST;
-};
-
-static void swiftTypeContextInfoImpl(
-    SwiftLangSupport &Lang, llvm::MemoryBuffer *UnresolvedInputFile,
-    unsigned Offset, ArrayRef<const char *> Args,
-    llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
-    llvm::function_ref<void(CancellableResult<SwiftTypeContextInfoImplResult>)>
-        Callback) {
-  assert(Callback && "Must provide callback to receive results");
-
-  using ResultType = CancellableResult<SwiftTypeContextInfoImplResult>;
-
-  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}));
-    }
-  };
-
-  Lang.performCompletionLikeOperation(
-      UnresolvedInputFile, Offset, Args, FileSystem,
-      [&](CancellableResult<CompletionInstanceResult> Result) {
+void deliverResults(SourceKit::TypeContextInfoConsumer &SKConsumer,
+                    CancellableResult<TypeContextInfoResult> Result) {
   switch (Result.getKind()) {
   case CancellableResultKind::Success: {
-          ConsumerToCallbackAdapter Consumer(Result->DidReuseAST, Callback);
+    SKConsumer.setReusingASTContext(Result->DidReuseAST);
 
-          // Create a factory for code completion callbacks that will feed the
-          // Consumer.
-          std::unique_ptr<CodeCompletionCallbacksFactory> callbacksFactory(
-              ide::makeTypeContextInfoCallbacksFactory(Consumer));
-
-          if (!Result->DidFindCodeCompletionToken) {
-            Callback(
-                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 Callback gets called
-            // exactly once, call it manually with no results here.
-            Callback(
-                ResultType::success({/*Results=*/{}, Result->DidReuseAST}));
-          }
-          break;
-        }
-        case CancellableResultKind::Failure:
-          Callback(ResultType::failure(Result.getError()));
-          break;
-        case CancellableResultKind::Cancelled:
-          Callback(ResultType::cancelled());
-          break;
-        }
-      });
-}
-
-void SwiftLangSupport::getExpressionContextInfo(
-    llvm::MemoryBuffer *UnresolvedInputFile, unsigned Offset,
-    OptionsDictionary *optionsDict, ArrayRef<const char *> Args,
-    SourceKit::TypeContextInfoConsumer &SKConsumer,
-    Optional<VFSOptions> vfsOptions) {
-  std::string error;
-
-  TypeContextInfo::Options options;
-  if (optionsDict) {
-    translateTypeContextInfoOptions(*optionsDict, options);
-  }
-
-  // FIXME: the use of None as primary file is to match the fact we do not read
-  // the document contents using the editor documents infrastructure.
-  auto fileSystem = getFileSystem(vfsOptions, /*primaryFile=*/None, error);
-  if (!fileSystem)
-    return SKConsumer.failed(error);
-
-  class Consumer : public ide::TypeContextInfoConsumer {
-    SourceKit::TypeContextInfoConsumer &SKConsumer;
-
-    /// Convert an IDE result to a SK result and send it to \c SKConsumer.
-    void handleSingleResult(const ide::TypeContextInfoItem &Item) {
+    for (auto &Item : Result->Results) {
       SmallString<512> SS;
       llvm::raw_svector_ostream OS(SS);
 
@@ -201,28 +114,6 @@ void SwiftLangSupport::getExpressionContextInfo(
 
       SKConsumer.handleResult(Info);
     }
-
-  public:
-    Consumer(SourceKit::TypeContextInfoConsumer &SKConsumer)
-        : SKConsumer(SKConsumer){};
-
-    void handleResults(ArrayRef<ide::TypeContextInfoItem> Results) override {
-      for (auto &Item : Results)
-        handleSingleResult(Item);
-    }
-
-    void setReusingASTContext(bool flag) {
-      SKConsumer.setReusingASTContext(flag);
-    }
-  } Consumer(SKConsumer);
-
-  swiftTypeContextInfoImpl(
-      *this, UnresolvedInputFile, Offset, Args, fileSystem,
-      [&](CancellableResult<SwiftTypeContextInfoImplResult> Result) {
-        switch (Result.getKind()) {
-        case CancellableResultKind::Success: {
-          Consumer.handleResults(Result->Results);
-          Consumer.setReusingASTContext(Result->DidReuseAST);
     break;
   }
   case CancellableResultKind::Failure:
@@ -232,5 +123,37 @@ void SwiftLangSupport::getExpressionContextInfo(
     SKConsumer.cancelled();
     break;
   }
+}
+
+void SwiftLangSupport::getExpressionContextInfo(
+    llvm::MemoryBuffer *UnresolvedInputFile, unsigned Offset,
+    OptionsDictionary *optionsDict, ArrayRef<const char *> Args,
+    SourceKit::TypeContextInfoConsumer &SKConsumer,
+    Optional<VFSOptions> vfsOptions) {
+  std::string error;
+
+  TypeContextInfo::Options options;
+  if (optionsDict) {
+    translateTypeContextInfoOptions(*optionsDict, options);
+  }
+
+  // FIXME: the use of None as primary file is to match the fact we do not read
+  // the document contents using the editor documents infrastructure.
+  auto fileSystem = getFileSystem(vfsOptions, /*primaryFile=*/None, error);
+  if (!fileSystem) {
+    return SKConsumer.failed(error);
+  }
+
+  performWithParamsToCompletionLikeOperation(
+      UnresolvedInputFile, Offset, Args, fileSystem,
+      [&](CancellableResult<CompletionLikeOperationParams> ParamsResult) {
+        ParamsResult.mapAsync<TypeContextInfoResult>(
+            [&](auto &CIParams, auto DeliverTransformed) {
+              getCompletionInstance()->typeContextInfo(
+                  CIParams.Invocation, Args, fileSystem,
+                  CIParams.completionBuffer, Offset, CIParams.DiagC,
+                  DeliverTransformed);
+            },
+            [&](auto Result) { deliverResults(SKConsumer, Result); });
       });
 }