Merge pull request #40158 from ahoppen/pr/cancel-code-completion

[SourceKit] Support cancellation of code completion like requests

Author: ahoppen

include/swift/AST/ASTContext.h

@@ -273,6 +273,15 @@ class ASTContext final {
   /// Diags - The diagnostics engine.
   DiagnosticEngine &Diags;
 
+  /// If the shared pointer is not a \c nullptr and the pointee is \c true,
+  /// all operations working on this ASTContext should be aborted at the next
+  /// possible opportunity.
+  /// This is used by SourceKit to cancel requests for which the result is no
+  /// longer of interest.
+  /// The returned result will be discarded, so the operation that acknowledges
+  /// the cancellation might return with any result.
+  std::shared_ptr<std::atomic<bool>> CancellationFlag = nullptr;
+
   TypeCheckCompletionCallback *CompletionCallback = nullptr;
 
   /// The request-evaluator that is used to process various requests.

include/swift/Basic/LangOptions.h

@@ -613,12 +613,6 @@ namespace swift {
     /// than this many seconds.
     unsigned ExpressionTimeoutThreshold = 600;
 
-    /// If the shared pointer is not a \c nullptr and the pointee is \c true,
-    /// typechecking should be aborted at the next possible opportunity.
-    /// This is used by SourceKit to cancel requests for which the result is no
-    /// longer of interest.
-    std::shared_ptr<std::atomic<bool>> CancellationFlag = nullptr;
-
     /// If non-zero, abort the switch statement exhaustiveness checker if
     /// the Space::minus function is called more than this many times.
     ///

include/swift/IDE/CompletionInstance.h

@@ -106,6 +106,7 @@ class CompletionInstance {
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
       DiagnosticConsumer *DiagC,
+      std::shared_ptr<std::atomic<bool>> CancellationFlag,
       llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
           Callback);
 
@@ -119,6 +120,7 @@ class CompletionInstance {
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
       DiagnosticConsumer *DiagC,
+      std::shared_ptr<std::atomic<bool>> CancellationFlag,
       llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
           Callback);
 
@@ -129,6 +131,14 @@ class CompletionInstance {
   /// In case of failure or cancellation, the callback receives the
   /// corresponding failed or cancelled result.
   ///
+  /// If \p CancellationFlag is not \c nullptr, code completion can be cancelled
+  /// by setting the flag to \c true.
+  /// IMPORTANT: If \p CancellationFlag is not \c nullptr, then completion might
+  /// be cancelled in the secon pass that's invoked inside \p Callback.
+  /// Therefore, \p Callback MUST check whether completion was cancelled before
+  /// interpreting the results, since invalid results may be returned in case
+  /// of cancellation.
+  ///
   /// 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.
@@ -137,6 +147,7 @@ class CompletionInstance {
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
       DiagnosticConsumer *DiagC,
+      std::shared_ptr<std::atomic<bool>> CancellationFlag,
       llvm::function_ref<void(CancellableResult<CompletionInstanceResult>)>
           Callback);
 
@@ -155,13 +166,15 @@ class CompletionInstance {
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
       DiagnosticConsumer *DiagC, ide::CodeCompletionContext &CompletionContext,
+      std::shared_ptr<std::atomic<bool>> CancellationFlag,
       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,
+      std::shared_ptr<std::atomic<bool>> CancellationFlag,
       llvm::function_ref<void(CancellableResult<TypeContextInfoResult>)>
           Callback);
 
@@ -170,6 +183,7 @@ class CompletionInstance {
       llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FileSystem,
       llvm::MemoryBuffer *completionBuffer, unsigned int Offset,
       DiagnosticConsumer *DiagC, ArrayRef<const char *> ExpectedTypeNames,
+      std::shared_ptr<std::atomic<bool>> CancellationFlag,
       llvm::function_ref<void(CancellableResult<ConformingMethodListResults>)>
           Callback);
 };

include/swift/Sema/ConstraintSystem.h

@@ -5179,7 +5179,7 @@ class ConstraintSystem {
     if (isExpressionAlreadyTooComplex)
       return true;
 
-    auto CancellationFlag = getASTContext().TypeCheckerOpts.CancellationFlag;
+    auto CancellationFlag = getASTContext().CancellationFlag;
     if (CancellationFlag && CancellationFlag->load(std::memory_order_relaxed))
       return true;