Merge pull request #40964 from rintaro/sourcekit-compile-cancellation

[SourceKit] Add cancellation points for 'compile' requests

Author: rintaro

include/swift/Frontend/Frontend.h

@@ -676,6 +676,9 @@ class CompilerInstance {
   /// If \p fn returns true, exits early and returns true.
   bool forEachFileToTypeCheck(llvm::function_ref<bool(SourceFile &)> fn);
 
+  /// Whether the cancellation of the current operation has been requested.
+  bool isCancellationRequested() const;
+
 private:
   /// Compute the parsing options for a source file in the main module.
   SourceFile::ParsingOptions getSourceFileParsingOptions(bool forPrimary) const;

lib/Frontend/Frontend.cpp

@@ -1368,6 +1368,10 @@ bool CompilerInstance::performSILProcessing(SILModule *silModule) {
   return false;
 }
 
+bool CompilerInstance::isCancellationRequested() const {
+  auto flag = getASTContext().CancellationFlag;
+  return flag && flag->load(std::memory_order_relaxed);
+}
 
 const PrimarySpecificPaths &
 CompilerInstance::getPrimarySpecificPathsForWholeModuleOptimizationMode()

lib/FrontendTool/FrontendTool.cpp

@@ -1556,6 +1556,10 @@ static bool performCompileStepsPostSILGen(CompilerInstance &Instance,
   if (observer)
     observer->performedSILGeneration(*SM);
 
+  // Cancellation check after SILGen.
+  if (Instance.isCancellationRequested())
+    return true;
+
   auto *Stats = Instance.getASTContext().Stats;
   if (Stats)
     countStatsPostSILGen(*Stats, *SM);
@@ -1629,6 +1633,10 @@ static bool performCompileStepsPostSILGen(CompilerInstance &Instance,
   if (observer)
     observer->performedSILProcessing(*SM);
 
+  // Cancellation check after SILOptimzation.
+  if (Instance.isCancellationRequested())
+    return true;
+
   if (PSPs.haveModuleSummaryOutputPath()) {
     if (serializeModuleSummary(SM.get(), PSPs, Context)) {
       return true;
@@ -1663,6 +1671,10 @@ static bool performCompileStepsPostSILGen(CompilerInstance &Instance,
 
   runSILLoweringPasses(*SM);
 
+  // Cancellation check after SILLowering.
+  if (Instance.isCancellationRequested())
+    return true;
+
   // TODO: at this point we need to flush any the _tracing and profiling_
   // in the UnifiedStatsReporter, because the those subsystems of the USR
   // retain _pointers into_ the SILModule, and the SILModule's lifecycle is
@@ -1690,6 +1702,10 @@ static bool performCompileStepsPostSILGen(CompilerInstance &Instance,
       IRGenOpts, Invocation.getTBDGenOptions(), std::move(SM), PSPs,
       OutputFilename, MSF, HashGlobal, ParallelOutputFilenames);
 
+  // Cancellation check after IRGen.
+  if (Instance.isCancellationRequested())
+    return true;
+
   // If no IRModule is available, bail. This can either happen if IR generation
   // fails, or if parallelIRGen happened correctly (in which case it would have
   // already performed LLVM).

lib/IDE/CompileInstance.cpp

@@ -329,11 +329,18 @@ bool CompileInstance::performCompile(
     DiagnosticConsumer *DiagC,
     std::shared_ptr<std::atomic<bool>> CancellationFlag) {
 
+  // Cancellation check. This gives a chance to cancel queued up requests before
+  // processing anything.
+  if (CancellationFlag && CancellationFlag->load(std::memory_order_relaxed))
+    return true;
+
   performSema(Args, fileSystem, DiagC, CancellationFlag);
   if (CI->getDiags().hadAnyError())
     return true;
 
-  // TODO: Cancellation check.
+  // Cancellation check after Sema.
+  if (CI->isCancellationRequested())
+    return true;
 
   int ReturnValue = 0;
   return performCompileStepsPostSema(*CI, ReturnValue, /*observer=*/nullptr);