Introduce BindExtensionsForIDEInspectionRequest

This allows us to lazily bind extensions after mutating the AST,
ensuring we don't prematurely kick the building of lookup tables.

Author: hamishknight

include/swift/AST/TypeCheckRequests.h

@@ -5389,6 +5389,25 @@ class DefaultIsolationInSourceFileRequest
   bool isCached() const { return true; }
 };
 
+/// A request that allows IDE inspection to lazily kick extension binding after
+/// it has finished mutating the AST. This will eventually be subsumed when we
+/// properly requestify extension binding.
+class BindExtensionsForIDEInspectionRequest
+    : public SimpleRequest<BindExtensionsForIDEInspectionRequest,
+                           evaluator::SideEffect(ModuleDecl *),
+                           RequestFlags::Cached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  evaluator::SideEffect evaluate(Evaluator &evaluator, ModuleDecl *M) const;
+
+public:
+  bool isCached() const { return true; }
+};
+
 class ModuleHasTypeCheckerPerformanceHacksEnabledRequest
     : public SimpleRequest<ModuleHasTypeCheckerPerformanceHacksEnabledRequest,
                            bool(const ModuleDecl *),

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -641,6 +641,10 @@ SWIFT_REQUEST(TypeChecker, DefaultIsolationInSourceFileRequest,
               std::optional<DefaultIsolation>(const SourceFile *),
               Cached, NoLocationInfo)
 
+SWIFT_REQUEST(TypeChecker, BindExtensionsForIDEInspectionRequest,
+              evaluator::SideEffect(ModuleDecl *),
+              Cached, NoLocationInfo)
+
 SWIFT_REQUEST(TypeChecker, ModuleHasTypeCheckerPerformanceHacksEnabledRequest,
               bool(const ModuleDecl *),
               Cached, NoLocationInfo)

lib/IDE/REPLCodeCompletion.cpp

@@ -256,7 +256,6 @@ doCodeCompletion(SourceFile &SF, StringRef EnteredCode, unsigned *BufferID,
 
   auto &newSF = newModule->getMainSourceFile();
   performImportResolution(newSF);
-  bindExtensions(*newModule);
 
   performIDEInspectionSecondPass(newSF, *CompletionCallbacksFactory);
 

lib/IDETool/IDEInspectionInstance.cpp

@@ -429,7 +429,6 @@ bool IDEInspectionInstance::performCachedOperationIfPossible(
     // re-use imported modules.
     auto *newSF = &newM->getMainSourceFile();
     performImportResolution(*newSF);
-    bindExtensions(*newM);
 
     traceDC = newM;
 #ifndef NDEBUG
@@ -518,7 +517,8 @@ void IDEInspectionInstance::performNewOperation(
     CI->getASTContext().CancellationFlag = CancellationFlag;
     registerIDERequestFunctions(CI->getASTContext().evaluator);
 
-    CI->performParseAndResolveImportsOnly();
+    CI->loadAccessNotesIfNeeded();
+    performImportResolution(CI->getMainModule());
 
     bool DidFindIDEInspectionTarget = CI->getIDEInspectionFile()
                                         ->getDelayedParserState()

lib/Parse/Parser.cpp

@@ -162,6 +162,7 @@ void Parser::performIDEInspectionSecondPassImpl(
   // Clear any ASTScopes that were expanded.
   SF.clearScope();
 
+  // FIXME: We shouldn't be mutating the AST after-the-fact like this.
   switch (info.Kind) {
   case IDEInspectionDelayedDeclKind::TopLevelCodeDecl: {
     // Re-enter the top-level code decl context.
@@ -209,7 +210,17 @@ void Parser::performIDEInspectionSecondPassImpl(
   assert(!State->hasIDEInspectionDelayedDeclState() &&
          "Second pass should not set any code completion info");
 
-  DoneParsingCallback->doneParsing(DC->getParentSourceFile());
+  auto *SF = DC->getParentSourceFile();
+
+  // Bind extensions if needed. This needs to be done here since we may have
+  // mutated the AST above.
+  {
+    auto *M = SF->getParentModule();
+    BindExtensionsForIDEInspectionRequest req(M);
+    evaluateOrDefault(Context.evaluator, req, {});
+  }
+
+  DoneParsingCallback->doneParsing(SF);
 
   State->restoreIDEInspectionDelayedDeclState(info);
 }

lib/Sema/TypeChecker.cpp

@@ -793,3 +793,10 @@ std::pair<bool, bool> EvaluateIfConditionRequest::evaluate(
   llvm_unreachable("Must not be used in C++-only build");
 #endif
 }
+
+evaluator::SideEffect
+BindExtensionsForIDEInspectionRequest::evaluate(Evaluator &evaluator,
+                                                ModuleDecl *M) const {
+  bindExtensions(*M);
+  return {};
+}

test/IDE/pr83369.swift

@@ -0,0 +1,9 @@
+// RUN: %batch-code-completion -module-name main
+
+extension Int {}
+
+// Make sure we can resolve P here.
+protocol P<X> where X: main.#^COMPLETE^# {
+  associatedtype X
+}
+// COMPLETE: Decl[Protocol]/CurrModule: P[#P#]; name=P

validation-test/IDE/crashers/b4e591f09ce81b0.swift renamed to validation-test/IDE/crashers_fixed/b4e591f09ce81b0.swift

@@ -1,5 +1,5 @@
 // {"kind":"complete","original":"0b319c77","signature":"swift::ExtensionDecl::getExtendedNominal() const","signatureAssert":"Extension must have already been bound"}
-// RUN: not --crash %target-swift-ide-test -code-completion -batch-code-completion -skip-filecheck -code-completion-diagnostics -source-filename %s
+// RUN: %target-swift-ide-test -code-completion -batch-code-completion -skip-filecheck -code-completion-diagnostics -source-filename %s
 extension a[ {
   #^^#
 }

validation-test/IDE/crashers/f1328ad669bd7cf8.swift renamed to validation-test/IDE/crashers_fixed/f1328ad669bd7cf8.swift

@@ -1,4 +1,4 @@
 // {"kind":"complete","signature":"Extension must have already been bound","signatureAssert":"Extension must have already been bound"}
-// RUN: not --crash %target-swift-ide-test -code-completion -batch-code-completion -skip-filecheck -code-completion-diagnostics -source-filename %s
+// RUN: %target-swift-ide-test -code-completion -batch-code-completion -skip-filecheck -code-completion-diagnostics -source-filename %s
 extension a [ { func b {
 #^COMPLETE^#