[IDE] Perform extension binding after AST mutation

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)

include/swift/Frontend/Frontend.h

@@ -781,6 +781,11 @@ class CompilerInstance {
   /// Parses and type-checks all input files.
   void performSema();
 
+  /// Loads any access notes for the main module.
+  ///
+  /// FIXME: This should be requestified.
+  void loadAccessNotesIfNeeded();
+
   /// Parses and performs import resolution on all input files.
   ///
   /// This is similar to a parse-only invocation, but module imports will also

lib/Frontend/Frontend.cpp

@@ -1543,33 +1543,40 @@ void CompilerInstance::setMainModule(ModuleDecl *newMod) {
   Context->MainModule = newMod;
 }
 
+void CompilerInstance::loadAccessNotesIfNeeded() {
+  if (Invocation.getFrontendOptions().AccessNotesPath.empty())
+    return;
+
+  auto *mainModule = getMainModule();
+
+  auto accessNotesPath = Invocation.getFrontendOptions().AccessNotesPath;
+
+  auto bufferOrError =
+      swift::vfs::getFileOrSTDIN(getFileSystem(), accessNotesPath);
+  if (bufferOrError) {
+    int sourceID = SourceMgr.addNewSourceBuffer(std::move(bufferOrError.get()));
+    auto buffer = SourceMgr.getLLVMSourceMgr().getMemoryBuffer(sourceID);
+
+    if (auto accessNotesFile = AccessNotesFile::load(*Context, buffer))
+      mainModule->getAccessNotes() = *accessNotesFile;
+  } else {
+    Diagnostics.diagnose(SourceLoc(), diag::access_notes_file_io_error,
+                         accessNotesPath, bufferOrError.getError().message());
+  }
+}
+
 bool CompilerInstance::performParseAndResolveImportsOnly() {
   FrontendStatsTracer tracer(getStatsReporter(), "parse-and-resolve-imports");
 
-  auto *mainModule = getMainModule();
+  // NOTE: Do not add new logic to this function, use the request evaluator to
+  // lazily evaluate instead. Once the below computations are requestified we
+  // ought to be able to remove this function.
 
   // Load access notes.
-  if (!Invocation.getFrontendOptions().AccessNotesPath.empty()) {
-    auto accessNotesPath = Invocation.getFrontendOptions().AccessNotesPath;
-
-    auto bufferOrError =
-        swift::vfs::getFileOrSTDIN(getFileSystem(), accessNotesPath);
-    if (bufferOrError) {
-      int sourceID =
-          SourceMgr.addNewSourceBuffer(std::move(bufferOrError.get()));
-      auto buffer =
-          SourceMgr.getLLVMSourceMgr().getMemoryBuffer(sourceID);
-
-      if (auto accessNotesFile = AccessNotesFile::load(*Context, buffer))
-        mainModule->getAccessNotes() = *accessNotesFile;
-    }
-    else {
-      Diagnostics.diagnose(SourceLoc(), diag::access_notes_file_io_error,
-                           accessNotesPath, bufferOrError.getError().message());
-    }
-  }
+  loadAccessNotesIfNeeded();
 
   // Resolve imports for all the source files in the module.
+  auto *mainModule = getMainModule();
   performImportResolution(mainModule);
 
   bindExtensions(*mainModule);

lib/IDE/CodeCompletion.cpp

@@ -1485,12 +1485,12 @@ void CodeCompletionCallbacksImpl::typeCheckWithLookup(
       ASTNode Call = CallExpr::create(
           CurDeclContext->getASTContext(), AttrWithCompletion->getTypeExpr(),
           AttrWithCompletion->getArgs(), /*implicit=*/true);
-      typeCheckContextAt(
+      swift::typeCheckASTNodeAtLoc(
           TypeCheckASTNodeAtLocContext::node(CurDeclContext, Call),
           CompletionLoc);
     }
   } else {
-    typeCheckContextAt(
+    swift::typeCheckASTNodeAtLoc(
         TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
         CompletionLoc);
   }
@@ -1551,8 +1551,9 @@ void CodeCompletionCallbacksImpl::postfixCompletion(SourceLoc CompletionLoc,
 
       llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
           Context.CompletionCallback, &Lookup);
-      typeCheckContextAt(TypeCheckASTNodeAtLocContext::node(CurDeclContext, AE),
-                         CompletionLoc);
+      swift::typeCheckASTNodeAtLoc(
+          TypeCheckASTNodeAtLocContext::node(CurDeclContext, AE),
+          CompletionLoc);
       Lookup.collectResults(/*IsLabeledTrailingClosure=*/true, CompletionLoc,
                             CurDeclContext, CompletionContext);
     }
@@ -1711,7 +1712,7 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
 
   if (Kind != CompletionKind::TypeSimpleWithDot) {
     // Type member completion does not need a type-checked AST.
-    typeCheckContextAt(
+    swift::typeCheckASTNodeAtLoc(
         TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
         ParsedExpr ? ParsedExpr->getLoc()
                    : CurDeclContext->getASTContext()

lib/IDE/ConformingMethodList.cpp

@@ -109,7 +109,7 @@ void ConformingMethodListCallbacks::doneParsing(SourceFile *SrcFile) {
   {
     llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
         Context.CompletionCallback, &TypeCheckCallback);
-    typeCheckContextAt(
+    swift::typeCheckASTNodeAtLoc(
         TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
         CCExpr->getLoc());
   }

lib/IDE/ExprContextAnalysis.cpp

@@ -40,45 +40,6 @@
 using namespace swift;
 using namespace ide;
 
-//===----------------------------------------------------------------------===//
-// typeCheckContextAt(DeclContext, SourceLoc)
-//===----------------------------------------------------------------------===//
-
-void swift::ide::typeCheckContextAt(TypeCheckASTNodeAtLocContext TypeCheckCtx,
-                                    SourceLoc Loc) {
-  // Make sure the extension has been bound.
-  auto DC = TypeCheckCtx.getDeclContext();
-  // Even if the extension is invalid (e.g. nested in a function or another
-  // type), we want to know the "intended nominal" of the extension so that
-  // we can know the type of 'Self'.
-  SmallVector<ExtensionDecl *, 1> extensions;
-  for (auto typeCtx = DC->getInnermostTypeContext(); typeCtx != nullptr;
-       typeCtx = typeCtx->getParent()->getInnermostTypeContext()) {
-    if (auto *ext = dyn_cast<ExtensionDecl>(typeCtx))
-      extensions.push_back(ext);
-  }
-  while (!extensions.empty()) {
-    extensions.back()->computeExtendedNominal();
-    extensions.pop_back();
-  }
-
-  // If the completion happens in the inheritance clause of the extension,
-  // 'DC' is the parent of the extension. We need to iterate the top level
-  // decls to find it. In theory, we don't need the extended nominal in the
-  // inheritance clause, but ASTScope lookup requires that. We don't care
-  // unless 'DC' is not 'SourceFile' because non-toplevel extensions are
-  // 'canNeverBeBound()' anyway.
-  if (auto *SF = dyn_cast<SourceFile>(DC)) {
-    auto &SM = DC->getASTContext().SourceMgr;
-    for (auto *decl : SF->getTopLevelDecls())
-      if (auto *ext = dyn_cast<ExtensionDecl>(decl))
-        if (SM.rangeContainsTokenLoc(ext->getSourceRange(), Loc))
-          ext->computeExtendedNominal();
-  }
-
-  swift::typeCheckASTNodeAtLoc(TypeCheckCtx, Loc);
-}
-
 //===----------------------------------------------------------------------===//
 // findParsedExpr(DeclContext, Expr)
 //===----------------------------------------------------------------------===//

lib/IDE/ExprContextAnalysis.h

@@ -28,11 +28,6 @@ class ValueDecl;
 namespace ide {
 enum class SemanticContextKind : uint8_t;
 
-/// Type check parent contexts of the given decl context, and the body of the
-/// given context until \c Loc if the context is a function body.
-void typeCheckContextAt(TypeCheckASTNodeAtLocContext TypeCheckCtx,
-                        SourceLoc Loc);
-
 /// From \p DC, find and returns the outer most expression which source range is
 /// exact the same as \p TargetRange. Returns \c nullptr if not found.
 Expr *findParsedExpr(const DeclContext *DC, SourceRange TargetRange);

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/IDE/TypeContextInfo.cpp

@@ -119,7 +119,7 @@ void ContextInfoCallbacks::doneParsing(SourceFile *SrcFile) {
   {
     llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
         Context.CompletionCallback, &TypeCheckCallback);
-    typeCheckContextAt(
+    swift::typeCheckASTNodeAtLoc(
         TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
         ParsedExpr->getLoc());
   }

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