Introduce ResolveImportsAndBindExtensionsForIDEInspectionRequest

This allows us to lazily resolve imports and 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,26 @@ class DefaultIsolationInSourceFileRequest
   bool isCached() const { return true; }
 };
 
+/// A request that allows IDE inspection to lazily kick import resolution and
+/// extension binding after it has finished mutating the AST. This will
+/// eventually be subsumed when we requestify import resolution and extension
+/// binding.
+class ResolveImportsAndBindExtensionsForIDEInspectionRequest
+    : public SimpleRequest<
+          ResolveImportsAndBindExtensionsForIDEInspectionRequest,
+          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, ResolveImportsAndBindExtensionsForIDEInspectionRequest,
+              evaluator::SideEffect(ModuleDecl *),
+              Cached, NoLocationInfo)
+
 SWIFT_REQUEST(TypeChecker, ModuleHasTypeCheckerPerformanceHacksEnabledRequest,
               bool(const ModuleDecl *),
               Cached, NoLocationInfo)

lib/IDE/IDETypeChecking.cpp

@@ -32,6 +32,7 @@
 #include "swift/IDE/SourceEntityWalker.h"
 #include "swift/Parse/Lexer.h"
 #include "swift/Sema/IDETypeCheckingRequests.h"
+#include "swift/Subsystems.h"
 #include "llvm/ADT/SmallVector.h"
 
 using namespace swift;
@@ -1048,3 +1049,11 @@ swift::getShorthandShadows(LabeledConditionalStmt *CondStmt, DeclContext *DC) {
   }
   return Result;
 }
+
+evaluator::SideEffect
+ResolveImportsAndBindExtensionsForIDEInspectionRequest::evaluate(
+    Evaluator &evaluator, ModuleDecl *M) const {
+  performImportResolution(M);
+  bindExtensions(*M);
+  return {};
+}

lib/IDE/REPLCodeCompletion.cpp

@@ -254,11 +254,8 @@ doCodeCompletion(SourceFile &SF, StringRef EnteredCode, unsigned *BufferID,
     addFile(new (Ctx) SourceFile(*newModule, SourceFileKind::Main, *BufferID));
   });
 
-  auto &newSF = newModule->getMainSourceFile();
-  performImportResolution(newSF);
-  bindExtensions(*newModule);
-
-  performIDEInspectionSecondPass(newSF, *CompletionCallbacksFactory);
+  performIDEInspectionSecondPass(newModule->getMainSourceFile(),
+                                 *CompletionCallbacksFactory);
 
   // Reset the error state because it's only relevant to the code that we just
   // processed, which now gets thrown away.

lib/IDETool/IDEInspectionInstance.cpp

@@ -425,11 +425,9 @@ bool IDEInspectionInstance::performCachedOperationIfPossible(
     // Tell the compiler instance we've replaced the main module.
     CachedCI->setMainModule(newM);
 
-    // Re-process the whole file (parsing will be lazily triggered). Still
-    // re-use imported modules.
+    // Re-parse the SourceFile.
     auto *newSF = &newM->getMainSourceFile();
-    performImportResolution(*newSF);
-    bindExtensions(*newM);
+    (void)newSF->getTopLevelItems();
 
     traceDC = newM;
 #ifndef NDEBUG
@@ -518,7 +516,7 @@ void IDEInspectionInstance::performNewOperation(
     CI->getASTContext().CancellationFlag = CancellationFlag;
     registerIDERequestFunctions(CI->getASTContext().evaluator);
 
-    CI->performParseAndResolveImportsOnly();
+    CI->loadAccessNotesIfNeeded();
 
     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();
+
+  // Resolve imports and bind extensions if needed. This needs to be done here
+  // since we may have mutated the AST above.
+  {
+    auto *M = SF->getParentModule();
+    ResolveImportsAndBindExtensionsForIDEInspectionRequest req(M);
+    evaluateOrDefault(Context.evaluator, req, {});
+  }
+
+  DoneParsingCallback->doneParsing(SF);
 
   State->restoreIDEInspectionDelayedDeclState(info);
 }

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^#