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

This refactors a bunch of code-completion methods around `performOperation` to return their results via a callback only instead of the current mixed approach  of indicating failure via a return value, returning an error string as an inout parameter and success results via a callback. The new guarantee should be that the callback is always called exactly once on control flow graph.

Other than a support for passing the (currently unused) cancelled state through the different instance, there should be no functionality change.

Author: ahoppen

include/swift/IDE/CancellableResult.h

@@ -0,0 +1,148 @@
+//===--- 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:
+      Error = 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;
+  }
+};
+
+} // namespace ide
+} // namespace swift
+
+#endif // SWIFT_IDE_CANCELLABLE_RESULT_H

include/swift/IDE/CompletionInstance.h

@@ -14,6 +14,7 @@
 #define SWIFT_IDE_COMPLETIONINSTANCE_H
 
 #include "swift/Frontend/Frontend.h"
+#include "swift/IDE/CancellableResult.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/StringRef.h"
@@ -35,6 +36,14 @@ 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;
+};
+
 /// Manages \c CompilerInstance for completion like operations.
 class CompletionInstance {
   struct Options {
@@ -58,27 +67,31 @@ 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);
 
 public:
   CompletionInstance() : CachedCIShouldBeInvalidated(false) {}
@@ -94,18 +107,19 @@ class CompletionInstance {
   /// 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.
+  /// 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.
-  bool performOperation(
+  void 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 &, bool)> Callback);
+      DiagnosticConsumer *DiagC,
+      llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
+          Callback);
 };
 
 } // namespace ide

include/swift/IDE/ConformingMethodList.h

@@ -42,7 +42,6 @@ class ConformingMethodListConsumer {
 public:
   virtual ~ConformingMethodListConsumer() {}
   virtual void handleResult(const ConformingMethodListResult &result) = 0;
-  virtual void setReusingASTContext(bool flag) = 0;
 };
 
 /// Printing consumer
@@ -54,7 +53,6 @@ class PrintingConformingMethodListConsumer
   PrintingConformingMethodListConsumer(llvm::raw_ostream &OS) : OS(OS) {}
 
   void handleResult(const ConformingMethodListResult &result) override;
-  void setReusingASTContext(bool flag) override {}
 };
 
 /// Create a factory for code completion callbacks.

include/swift/IDE/TypeContextInfo.h

@@ -38,7 +38,6 @@ class TypeContextInfoConsumer {
 public:
   virtual ~TypeContextInfoConsumer() {}
   virtual void handleResults(ArrayRef<TypeContextInfoItem>) = 0;
-  virtual void setReusingASTContext(bool flag) = 0;
 };
 
 /// Printing consumer
@@ -49,7 +48,6 @@ class PrintingTypeContextInfoConsumer : public TypeContextInfoConsumer {
   PrintingTypeContextInfoConsumer(llvm::raw_ostream &OS) : OS(OS) {}
 
   void handleResults(ArrayRef<TypeContextInfoItem>) override;
-  void setReusingASTContext(bool flag) override {}
 };
 
 /// Create a factory for code completion callbacks.

lib/IDE/CompletionInstance.cpp

@@ -281,11 +281,12 @@ static bool areAnyDependentFilesInvalidated(
 } // namespace
 
 bool CompletionInstance::performCachedOperationIfPossible(
-   llvm::hash_code ArgsHash,
+    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) {
   llvm::PrettyStackTraceString trace(
       "While performing cached completion if possible");
 
@@ -504,7 +505,8 @@ bool CompletionInstance::performCachedOperationIfPossible(
     if (DiagC)
       CI.addDiagnosticConsumer(DiagC);
 
-    Callback(CI, /*reusingASTContext=*/true);
+    Callback(CancellableResult<CompletionInstanceResult>::success(
+        {CI, /*reusingASTContext=*/true}));
 
     if (DiagC)
       CI.removeDiagnosticConsumer(DiagC);
@@ -515,12 +517,13 @@ bool CompletionInstance::performCachedOperationIfPossible(
   return true;
 }
 
-bool CompletionInstance::performNewOperation(
+void CompletionInstance::performNewOperation(
     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) {
   llvm::PrettyStackTraceString trace("While performing new completion");
 
   auto isCachedCompletionRequested = ArgsHash.hasValue();
@@ -548,32 +551,35 @@ bool CompletionInstance::performNewOperation(
     Invocation.setCodeCompletionPoint(completionBuffer, Offset);
 
     if (CI.setup(Invocation)) {
-      Error = "failed to setup compiler instance";
-      return false;
+      Callback(CancellableResult<CompletionInstanceResult>::failure(
+          "failed to setup compiler instance"));
+      return;
     }
     registerIDERequestFunctions(CI.getASTContext().evaluator);
 
     // If we're expecting a standard library, but there either isn't one, or it
     // failed to load, let's bail early and hand back an empty completion
     // result to avoid any downstream crashes.
-    if (CI.loadStdlibIfNeeded())
-      return true;
+    if (CI.loadStdlibIfNeeded()) {
+      /// FIXME: Callback is not being called on this path.
+      return;
+    }
 
     CI.performParseAndResolveImportsOnly();
 
     // If we didn't find a code completion token, bail.
     auto *state = CI.getCodeCompletionFile()->getDelayedParserState();
     if (!state->hasCodeCompletionDelayedDeclState())
-      return true;
+      /// FIXME: Callback is not being called on this path.
+      return;
 
-    Callback(CI, /*reusingASTContext=*/false);
+    Callback(CancellableResult<CompletionInstanceResult>::success(
+        {CI, /*ReuisingASTContext=*/false}));
   }
 
   // Cache the compiler instance if fast completion is enabled.
   if (isCachedCompletionRequested)
     cacheCompilerInstance(std::move(TheInstance), *ArgsHash);
-
-  return true;
 }
 
 void CompletionInstance::cacheCompilerInstance(
@@ -608,12 +614,13 @@ void CompletionInstance::setOptions(CompletionInstance::Options NewOpts) {
   Opts = NewOpts;
 }
 
-bool swift::ide::CompletionInstance::performOperation(
+void 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,
-    llvm::function_ref<void(CompilerInstance &, bool)> Callback) {
+    DiagnosticConsumer *DiagC,
+    llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
+        Callback) {
 
   // Always disable source location resolutions from .swiftsourceinfo file
   // because they're somewhat heavy operations and aren't needed for completion.
@@ -637,14 +644,11 @@ bool swift::ide::CompletionInstance::performOperation(
 
   if (performCachedOperationIfPossible(ArgsHash, FileSystem, completionBuffer,
                                        Offset, DiagC, Callback)) {
-    return true;
+    // We were able to reuse a cached AST. Callback has already been invoked
+    // and we don't need to build a new AST. We are done.
+    return;
   }
 
-  if(performNewOperation(ArgsHash, Invocation, FileSystem, completionBuffer,
-                         Offset, Error, DiagC, Callback)) {
-    return true;
-  }
-
-  assert(!Error.empty());
-  return false;
+  performNewOperation(ArgsHash, Invocation, FileSystem, completionBuffer,
+                      Offset, DiagC, Callback);
 }