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

Author: ahoppen

include/swift/IDE/CodeCompletion.h

@@ -17,6 +17,7 @@
 #include "swift/Basic/Debug.h"
 #include "swift/Basic/LLVM.h"
 #include "swift/Basic/OptionSet.h"
+#include "swift/Frontend/Frontend.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
@@ -1062,6 +1063,12 @@ class CodeCompletionContext {
   }
 };
 
+struct SwiftCompletionInfo {
+  swift::ASTContext *swiftASTContext = nullptr;
+  const swift::CompilerInvocation *invocation = nullptr;
+  CodeCompletionContext *completionContext = nullptr;
+};
+
 /// An abstract base class for consumers of code completion results.
 /// \see \c SimpleCachingCodeCompletionConsumer.
 class CodeCompletionConsumer {

include/swift/IDE/CompletionInstance.h

@@ -15,6 +15,7 @@
 
 #include "swift/Frontend/Frontend.h"
 #include "swift/IDE/CancellableResult.h"
+#include "swift/IDE/CodeCompletion.h"
 #include "swift/IDE/ConformingMethodList.h"
 #include "swift/IDE/TypeContextInfo.h"
 #include "llvm/ADT/Hashing.h"
@@ -50,6 +51,12 @@ struct CompletionInstanceResult {
   bool DidFindCodeCompletionToken;
 };
 
+/// The results returned from \c CompletionInstance::codeComplete.
+struct CodeCompleteResult {
+  MutableArrayRef<CodeCompletionResult *> Results;
+  SwiftCompletionInfo &Info;
+};
+
 /// The results returned from \c CompletionInstance::typeContextInfo.
 struct TypeContextInfoResult {
   /// The actual results. If empty, no results were found.
@@ -143,6 +150,13 @@ class CompletionInstance {
       llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
           Callback);
 
+  void codeComplete(
+      swift::CompilerInvocation &Invocation, llvm::ArrayRef<const char *> Args,
+      llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+      llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+      DiagnosticConsumer *DiagC, ide::CodeCompletionContext &&CompletionContext,
+      llvm::function_ref<void(CancellableResult<CodeCompleteResult>)> Callback);
+
   void typeContextInfo(
       swift::CompilerInvocation &Invocation, llvm::ArrayRef<const char *> Args,
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,

lib/IDE/CompletionInstance.cpp

@@ -25,6 +25,7 @@
 #include "swift/ClangImporter/ClangModule.h"
 #include "swift/Driver/FrontendUtil.h"
 #include "swift/Frontend/Frontend.h"
+#include "swift/IDE/CodeCompletion.h"
 #include "swift/Parse/Lexer.h"
 #include "swift/Parse/PersistentParserState.h"
 #include "swift/Serialization/SerializedModuleLoader.h"
@@ -656,6 +657,81 @@ void swift::ide::CompletionInstance::performOperation(
                       Offset, DiagC, Callback);
 }
 
+void swift::ide::CompletionInstance::codeComplete(
+    swift::CompilerInvocation &Invocation, llvm::ArrayRef<const char *> Args,
+    llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
+    llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+    DiagnosticConsumer *DiagC, ide::CodeCompletionContext &&CompletionContext,
+    llvm::function_ref<void(CancellableResult<CodeCompleteResult>)> Callback) {
+  using ResultType = CancellableResult<CodeCompleteResult>;
+
+  struct ConsumerToCallbackAdapter
+      : public SimpleCachingCodeCompletionConsumer {
+    SwiftCompletionInfo SwiftContext;
+    llvm::function_ref<void(ResultType)> Callback;
+    bool HandleResultsCalled = false;
+
+    ConsumerToCallbackAdapter(llvm::function_ref<void(ResultType)> Callback)
+        : Callback(Callback) {}
+
+    void setContext(swift::ASTContext *context,
+                    const swift::CompilerInvocation *invocation,
+                    swift::ide::CodeCompletionContext *completionContext) {
+      SwiftContext.swiftASTContext = context;
+      SwiftContext.invocation = invocation;
+      SwiftContext.completionContext = completionContext;
+    }
+    void clearContext() { SwiftContext = SwiftCompletionInfo(); }
+
+    void handleResults(CodeCompletionContext &context) override {
+      HandleResultsCalled = true;
+      MutableArrayRef<CodeCompletionResult *> Results = context.takeResults();
+      assert(SwiftContext.swiftASTContext);
+      Callback(ResultType::success({Results, SwiftContext}));
+    }
+  };
+
+  performOperation(
+      Invocation, Args, FileSystem, completionBuffer, Offset, DiagC,
+      [&](CancellableResult<CompletionInstanceResult> CIResult) {
+        CIResult.mapAsync<CodeCompleteResult>(
+            [&CompletionContext](auto &Result, auto DeliverTransformed) {
+              CompletionContext.ReusingASTContext = Result.DidReuseAST;
+              CompilerInstance &CI = Result.CI;
+              ConsumerToCallbackAdapter Consumer(DeliverTransformed);
+
+              std::unique_ptr<CodeCompletionCallbacksFactory> callbacksFactory(
+                  ide::makeCodeCompletionCallbacksFactory(CompletionContext,
+                                                          Consumer));
+
+              if (!Result.DidFindCodeCompletionToken) {
+                SwiftCompletionInfo Info{&CI.getASTContext(),
+                                         &CI.getInvocation(),
+                                         &CompletionContext};
+                DeliverTransformed(ResultType::success({/*Results=*/{}, Info}));
+                return;
+              }
+
+              Consumer.setContext(&CI.getASTContext(), &CI.getInvocation(),
+                                  &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};
+                DeliverTransformed(ResultType::success({/*Results=*/{}, Info}));
+              }
+            },
+            Callback);
+      });
+}
+
 void swift::ide::CompletionInstance::typeContextInfo(
     swift::CompilerInvocation &Invocation, llvm::ArrayRef<const char *> Args,
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,

tools/SourceKit/lib/SwiftLang/CodeCompletionOrganizer.h

@@ -15,6 +15,7 @@
 
 #include "CodeCompletion.h"
 #include "SourceKit/Core/LangSupport.h"
+#include "swift/IDE/CodeCompletion.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace swift {
@@ -53,17 +54,11 @@ struct Options {
   unsigned popularityBonus = 2;
 };
 
-struct SwiftCompletionInfo {
-  swift::ASTContext *swiftASTContext = nullptr;
-  const swift::CompilerInvocation *invocation = nullptr;
-  swift::ide::CodeCompletionContext *completionContext = nullptr;
-};
-
 typedef llvm::StringMap<PopularityFactor> NameToPopularityMap;
 
 std::vector<Completion *>
 extendCompletions(ArrayRef<SwiftResult *> swiftResults, CompletionSink &sink,
-                  SwiftCompletionInfo &info,
+                  swift::ide::SwiftCompletionInfo &info,
                   const NameToPopularityMap *nameToPopularity,
                   const Options &options, Completion *prefix = nullptr,
                   bool clearFlair = false);