[CodeCompletion] Refactor how code completion results are returned to support cancellation

include/swift/IDE/CancellableResult.h

@@ -0,0 +1,178 @@
+//===--- CancellableResult.h ------------------------------------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_IDE_CANCELLABLE_RESULT_H
+#define SWIFT_IDE_CANCELLABLE_RESULT_H
+
+#include <string>
+
+namespace swift {
+
+namespace ide {
+
+enum class CancellableResultKind { Success, Failure, Cancelled };
+
+/// A result type that can carry one be in one of the following states:
+///  - Success and carry a value of \c ResultType
+///  - Failure and carry an error description
+///  - Cancelled in case the operation that produced the result was cancelled
+///
+/// Essentially this emulates an enum with associated values as follows
+/// \code
+/// enum CancellableResult<ResultType> {
+///   case success(ResultType)
+///   case failure(String)
+///   case cancelled
+/// }
+/// \endcode
+///
+/// The implementation is inspired by llvm::optional_detail::OptionalStorage
+template <typename ResultType>
+class CancellableResult {
+  CancellableResultKind Kind;
+  union {
+    /// If \c Kind == Success, carries the result.
+    ResultType Result;
+    /// If \c Kind == Error, carries the error description.
+    std::string Error;
+    /// If \c Kind == Cancelled, this union is not initialized.
+    char Empty;
+  };
+
+  CancellableResult(ResultType Result)
+      : Kind(CancellableResultKind::Success), Result(Result) {}
+
+  CancellableResult(std::string Error)
+      : Kind(CancellableResultKind::Failure), Error(Error) {}
+
+  explicit CancellableResult()
+      : Kind(CancellableResultKind::Cancelled), Empty() {}
+
+public:
+  CancellableResult(const CancellableResult &Other) : Kind(Other.Kind), Empty() {
+    switch (Kind) {
+    case CancellableResultKind::Success:
+      ::new ((void *)std::addressof(Result)) ResultType(Other.Result);
+      break;
+    case CancellableResultKind::Failure:
+      ::new ((void *)std::addressof(Error)) std::string(Other.Error);
+      break;
+    case CancellableResultKind::Cancelled:
+      break;
+    }
+  }
+
+  CancellableResult(CancellableResult &&Other) : Kind(Other.Kind), Empty() {
+    switch (Kind) {
+    case CancellableResultKind::Success:
+      ::new ((void *)std::addressof(Result))
+          ResultType(std::move(Other.Result));
+      break;
+    case CancellableResultKind::Failure:
+      ::new ((void *)std::addressof(Error)) std::string(std::move(Other.Error));
+      break;
+    case CancellableResultKind::Cancelled:
+      break;
+    }
+  }
+
+  ~CancellableResult() {
+    using std::string;
+    switch (Kind) {
+    case CancellableResultKind::Success:
+      Result.~ResultType();
+      break;
+    case CancellableResultKind::Failure:
+      Error.~string();
+      break;
+    case CancellableResultKind::Cancelled:
+      break;
+    }
+  }
+
+  /// Construct a \c CancellableResult that carries a successful result.
+  static CancellableResult success(ResultType Result) {
+    return std::move(CancellableResult(ResultType(Result)));
+  }
+
+  /// Construct a \c CancellableResult that carries the error message of a
+  /// failure.
+  static CancellableResult failure(std::string Error) {
+    return std::move(CancellableResult(Error));
+  }
+
+  /// Construct a \c CancellableResult representing that the producing operation
+  /// was cancelled.
+  static CancellableResult cancelled() {
+    return std::move(CancellableResult());
+  }
+
+  /// Return the result kind this \c CancellableResult represents: success,
+  /// failure or cancelled.
+  CancellableResultKind getKind() { return Kind; }
+
+  /// Assuming that the result represents success, return the underlying result
+  /// value.
+  ResultType &getResult() {
+    assert(getKind() == CancellableResultKind::Success);
+    return Result;
+  }
+
+  /// Assuming that the result represents success, retrieve members of the
+  /// underlying result value.
+  ResultType *operator->() { return &getResult(); }
+
+  /// Assuming that the result represents success, return the underlying result
+  /// value.
+  ResultType &operator*() { return getResult(); }
+
+  /// Assuming that the result represents a failure, return the error message.
+  std::string getError() {
+    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
+} // namespace swift
+
+#endif // SWIFT_IDE_CANCELLABLE_RESULT_H

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 {
@@ -1087,32 +1094,6 @@ struct SimpleCachingCodeCompletionConsumer : public CodeCompletionConsumer {
   virtual void handleResults(CodeCompletionContext &context) = 0;
 };
 
-/// A code completion result consumer that prints the results to a
-/// \c raw_ostream.
-class PrintingCodeCompletionConsumer
-    : public SimpleCachingCodeCompletionConsumer {
-  llvm::raw_ostream &OS;
-  bool IncludeKeywords;
-  bool IncludeComments;
-  bool IncludeSourceText;
-  bool PrintAnnotatedDescription;
-
-public:
- PrintingCodeCompletionConsumer(llvm::raw_ostream &OS,
-                                bool IncludeKeywords = true,
-                                bool IncludeComments = true,
-                                bool IncludeSourceText = false,
-                                bool PrintAnnotatedDescription = false)
-     : OS(OS),
-       IncludeKeywords(IncludeKeywords),
-       IncludeComments(IncludeComments),
-       IncludeSourceText(IncludeSourceText),
-       PrintAnnotatedDescription(PrintAnnotatedDescription) {}
-
-  void handleResults(CodeCompletionContext &context) override;
-  void handleResults(MutableArrayRef<CodeCompletionResult *> Results);
-};
-
 /// Create a factory for code completion callbacks.
 CodeCompletionCallbacksFactory *
 makeCodeCompletionCallbacksFactory(CodeCompletionContext &CompletionContext,

include/swift/IDE/CompletionInstance.h

@@ -14,6 +14,10 @@
 #define SWIFT_IDE_COMPLETIONINSTANCE_H
 
 #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"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringRef.h"
@@ -35,6 +39,40 @@ makeCodeCompletionMemoryBuffer(const llvm::MemoryBuffer *origBuf,
                                unsigned &Offset,
                                llvm::StringRef bufferIdentifier);
 
+/// 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;
+  /// Whether an AST was reused.
+  bool DidReuseAST;
+  /// Whether the CompletionInstance found a code completion token in the source
+  /// file. If this is \c false, the user will most likely want to return empty
+  /// results.
+  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.
+  ArrayRef<TypeContextInfoItem> Results;
+  /// Whether an AST was reused to produce the results.
+  bool DidReuseAST;
+};
+
+/// The results returned from \c CompletionInstance::conformingMethodList.
+struct ConformingMethodListResults {
+  /// The actual results. If \c nullptr, no results were found.
+  const ConformingMethodListResult *Result;
+  /// Whether an AST was reused for the completion.
+  bool DidReuseAST;
+};
+
 /// Manages \c CompilerInstance for completion like operations.
 class CompletionInstance {
   struct Options {
@@ -58,27 +96,49 @@ class CompletionInstance {
 
   /// 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 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.
+  /// Returns \c true if performing the cached operation was possible. Returns
+  /// \c false if the 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.
+  /// \p Callback will be called if and only if this function returns \c true.
   bool performCachedOperationIfPossible(
       llvm::hash_code ArgsHash,
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
       DiagnosticConsumer *DiagC,
-      llvm::function_ref<void(CompilerInstance &, bool)> Callback);
+      llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
+          Callback);
 
   /// 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(
+  void performNewOperation(
       llvm::Optional<llvm::hash_code> ArgsHash,
       swift::CompilerInvocation &Invocation,
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
-      std::string &Error, DiagnosticConsumer *DiagC,
-      llvm::function_ref<void(CompilerInstance &, bool)> Callback);
+      DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
+          Callback);
+
+  /// 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.
+  /// In case of failure or cancellation, the callback receives the
+  /// corresponding failed or cancelled result.
+  ///
+  /// 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.
+  void performOperation(
+      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<CompletionInstanceResult>)>
+          Callback);
 
 public:
   CompletionInstance() : CachedCIShouldBeInvalidated(false) {}
@@ -90,22 +150,28 @@ class CompletionInstance {
   // Update options with \c NewOpts. (Thread safe.)
   void setOptions(Options NewOpts);
 
-  /// 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.
-  bool performOperation(
+  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,
+      llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
+      DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CancellableResult<TypeContextInfoResult>)>
+          Callback);
+
+  void conformingMethodList(
       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 &, bool)> Callback);
+      DiagnosticConsumer *DiagC, ArrayRef<const char *> ExpectedTypeNames,
+      llvm::function_ref<void(CancellableResult<ConformingMethodListResults>)>
+          Callback);
 };
 
 } // namespace ide