Merge pull request #83777 from hamishknight/ext-cleanup

[Sema] Clean up extension binding a little

Author: hamishknight

include/swift/AST/Decl.h

@@ -2053,6 +2053,7 @@ class ExtensionDecl final : public GenericContext, public Decl,
   std::pair<LazyMemberLoader *, uint64_t> takeConformanceLoaderSlow();
 
   friend class ExtendedNominalRequest;
+  friend class BindExtensionsRequest;
   friend class Decl;
 public:
   using Decl::getASTContext;
@@ -2090,13 +2091,14 @@ class ExtensionDecl final : public GenericContext, public Decl,
   Type getExtendedType() const;
 
   /// Retrieve the nominal type declaration that is being extended.
-  /// Will  trip an assertion if the declaration has not already been computed.
+  /// Will trip an assertion if the declaration has not already been computed.
   /// In order to fail fast when type checking work is attempted
   /// before extension binding has taken place.
-
   NominalTypeDecl *getExtendedNominal() const;
 
-  /// Compute the nominal type declaration that is being extended.
+  /// Compute the nominal type declaration that is being extended. The result
+  /// is not cached, this should only be invoked by extension binding itself.
+  /// FIXME: Make this private once lldb has been migrated off it.
   NominalTypeDecl *computeExtendedNominal(
       bool excludeMacroExpansions=false) const;
 

include/swift/AST/NameLookupRequests.h

@@ -267,7 +267,7 @@ class HasMissingDesignatedInitializersRequest :
 /// Request the nominal declaration extended by a given extension declaration.
 class ExtendedNominalRequest
     : public SimpleRequest<
-          ExtendedNominalRequest, NominalTypeDecl *(ExtensionDecl *, bool),
+          ExtendedNominalRequest, NominalTypeDecl *(const ExtensionDecl *),
           RequestFlags::SeparatelyCached | RequestFlags::DependencySink> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -276,8 +276,7 @@ class ExtendedNominalRequest
   friend SimpleRequest;
 
   // Evaluation.
-  NominalTypeDecl * evaluate(Evaluator &evaluator, ExtensionDecl *ext,
-                             bool excludeMacroExpansions) const;
+  NominalTypeDecl * evaluate(Evaluator &evaluator, const ExtensionDecl *ext) const;
 
 public:
   // Separate caching.

include/swift/AST/NameLookupTypeIDZone.def

@@ -32,7 +32,7 @@ SWIFT_REQUEST(NameLookup, DirectPrecedenceGroupLookupRequest,
               TinyPtrVector<PrecedenceGroupDecl *>(OperatorLookupDescriptor),
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(NameLookup, ExtendedNominalRequest,
-              NominalTypeDecl *(ExtensionDecl *), SeparatelyCached,
+              NominalTypeDecl *(const ExtensionDecl *), SeparatelyCached,
               NoLocationInfo)
 SWIFT_REQUEST(NameLookup, GenericParamListRequest,
               GenericParamList *(GenericContext *), SeparatelyCached,

include/swift/AST/TypeCheckRequests.h

@@ -5437,11 +5437,9 @@ 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,
+/// Performs extension binding for all of the extensions in a module.
+class BindExtensionsRequest
+    : public SimpleRequest<BindExtensionsRequest,
                            evaluator::SideEffect(ModuleDecl *),
                            RequestFlags::Cached> {
 public:

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -647,7 +647,7 @@ SWIFT_REQUEST(TypeChecker, DefaultIsolationInSourceFileRequest,
               std::optional<DefaultIsolation>(const SourceFile *),
               Cached, NoLocationInfo)
 
-SWIFT_REQUEST(TypeChecker, BindExtensionsForIDEInspectionRequest,
+SWIFT_REQUEST(TypeChecker, BindExtensionsRequest,
               evaluator::SideEffect(ModuleDecl *),
               Cached, NoLocationInfo)
 

include/swift/Sema/SourceLoader.h

@@ -24,6 +24,7 @@ class ModuleDecl;
 class SourceLoader : public ModuleLoader {
 private:
   ASTContext &Ctx;
+  std::vector<ModuleDecl *> ModulesToBindExtensions;
   bool EnableLibraryEvolution;
 
   explicit SourceLoader(ASTContext &ctx,

lib/AST/Decl.cpp

@@ -2097,24 +2097,8 @@ ExtensionDecl::takeConformanceLoaderSlow() {
 }
 
 NominalTypeDecl *ExtensionDecl::getExtendedNominal() const {
-  if (hasBeenBound()) {
-    return ExtendedNominal.getPointer();
-  } else if (canNeverBeBound()) {
-    return computeExtendedNominal();
-  }
-
-  llvm_unreachable(
-      "Extension must have already been bound (by bindExtensions)");
-}
-
-NominalTypeDecl *ExtensionDecl::computeExtendedNominal(
-    bool excludeMacroExpansions) const {
-  ASTContext &ctx = getASTContext();
-  return evaluateOrDefault(ctx.evaluator,
-                           ExtendedNominalRequest{
-                               const_cast<ExtensionDecl *>(this),
-                               excludeMacroExpansions},
-                           nullptr);
+  auto &eval = getASTContext().evaluator;
+  return evaluateOrDefault(eval, ExtendedNominalRequest{this}, nullptr);
 }
 
 bool ExtensionDecl::canNeverBeBound() const {

lib/AST/NameLookup.cpp

@@ -3583,20 +3583,19 @@ ProtocolRequirementsRequest::evaluate(Evaluator &evaluator,
   return PD->getASTContext().AllocateCopy(requirements);
 }
 
-NominalTypeDecl *
-ExtendedNominalRequest::evaluate(Evaluator &evaluator,
-                                 ExtensionDecl *ext,
-                                 bool excludeMacroExpansions) const {
-  auto typeRepr = ext->getExtendedTypeRepr();
+NominalTypeDecl *ExtensionDecl::computeExtendedNominal(
+    bool excludeMacroExpansions) const {
+  auto typeRepr = getExtendedTypeRepr();
   if (!typeRepr) {
     // We must've seen 'extension { ... }' during parsing.
     return nullptr;
   }
 
-  ASTContext &ctx = ext->getASTContext();
+  ASTContext &ctx = getASTContext();
+  auto &evaluator = ctx.evaluator;
   auto options = defaultDirectlyReferencedTypeLookupOptions;
 
-  if (ext->isInSpecializeExtensionContext()) {
+  if (isInSpecializeExtensionContext()) {
     options |= DirectlyReferencedTypeLookupFlags::AllowUsableFromInline;
   }
 
@@ -3605,7 +3604,7 @@ ExtendedNominalRequest::evaluate(Evaluator &evaluator,
   }
 
   DirectlyReferencedTypeDecls referenced = directReferencesForTypeRepr(
-      evaluator, ctx, typeRepr, ext->getParent(), options);
+      evaluator, ctx, typeRepr, getParent(), options);
 
   // If there were no results, expand the lookup to include members that are
   // inaccessible due to missing imports. The missing imports will be diagnosed
@@ -3615,7 +3614,7 @@ ExtendedNominalRequest::evaluate(Evaluator &evaluator,
                               /*allowMigration=*/true)) {
     options |= DirectlyReferencedTypeLookupFlags::IgnoreMissingImports;
     referenced = directReferencesForTypeRepr(evaluator, ctx, typeRepr,
-                                             ext->getParent(), options);
+                                             getParent(), options);
   }
 
   // Resolve those type declarations to nominal type declarations.
@@ -3634,6 +3633,13 @@ ExtendedNominalRequest::evaluate(Evaluator &evaluator,
   return nominalTypes[0];
 }
 
+NominalTypeDecl *
+ExtendedNominalRequest::evaluate(Evaluator &evaluator,
+                                 const ExtensionDecl *ext) const {
+  ASSERT(ext->canNeverBeBound() && "Should have been bound by bindExtensions");
+  return ext->computeExtendedNominal();
+}
+
 /// Whether there are only associated types in the set of declarations.
 static bool declsAreAssociatedTypes(ArrayRef<TypeDecl *> decls) {
   if (decls.empty())

lib/AST/NameLookupRequests.cpp

@@ -205,22 +205,15 @@ HasMissingDesignatedInitializersRequest::evaluate(Evaluator &evaluator,
 
 std::optional<NominalTypeDecl *>
 ExtendedNominalRequest::getCachedResult() const {
-  // Note: if we fail to compute any nominal declaration, it's considered
-  // a cache miss. This allows us to recompute the extended nominal types
-  // during extension binding.
-  // This recomputation is also what allows you to extend types defined inside
-  // other extensions, regardless of source file order. See \c bindExtensions(),
-  // which uses a worklist algorithm that attempts to bind everything until
-  // fixed point.
   auto ext = std::get<0>(getStorage());
-  if (!ext->hasBeenBound() || !ext->getExtendedNominal())
+  if (!ext->hasBeenBound())
     return std::nullopt;
-  return ext->getExtendedNominal();
+  return ext->ExtendedNominal.getPointer();
 }
 
 void ExtendedNominalRequest::cacheResult(NominalTypeDecl *value) const {
   auto ext = std::get<0>(getStorage());
-  ext->setExtendedNominal(value);
+  const_cast<ExtensionDecl *>(ext)->setExtendedNominal(value);
 }
 
 void ExtendedNominalRequest::writeDependencySink(

lib/IDE/CodeCompletion.cpp

@@ -14,6 +14,7 @@
 #include "CodeCompletionDiagnostics.h"
 #include "CodeCompletionResultBuilder.h"
 #include "ExprContextAnalysis.h"
+#include "ReadyForTypeCheckingCallback.h"
 #include "swift/AST/ASTPrinter.h"
 #include "swift/AST/ASTWalker.h"
 #include "swift/AST/Comment.h"
@@ -105,7 +106,7 @@ std::string swift::ide::removeCodeCompletionTokens(
 namespace {
 
 class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
-                                    public DoneParsingCallback {
+                                    public ReadyForTypeCheckingCallback {
   CodeCompletionContext &CompletionContext;
   CodeCompletionConsumer &Consumer;
   CodeCompletionExpr *CodeCompleteTokenExpr = nullptr;
@@ -175,25 +176,6 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
 
   /// \returns true on success, false on failure.
   bool typecheckParsedType() {
-    // If the type appeared inside an extension, make sure that extension has
-    // been bound.
-    auto SF = CurDeclContext->getParentSourceFile();
-    auto visitTopLevelDecl = [&](Decl *D) {
-      if (auto ED = dyn_cast<ExtensionDecl>(D)) {
-        if (ED->getSourceRange().contains(ParsedTypeLoc.getLoc())) {
-          ED->computeExtendedNominal();
-        }
-      }
-    };
-    for (auto item : SF->getTopLevelItems()) {
-      if (auto D = item.dyn_cast<Decl *>()) {
-        visitTopLevelDecl(D);
-      }
-    }
-    for (auto *D : SF->getHoistedDecls()) {
-      visitTopLevelDecl(D);
-    }
-
     assert(ParsedTypeLoc.getTypeRepr() && "should have a TypeRepr");
     if (ParsedTypeLoc.wasValidated() && !ParsedTypeLoc.isError()) {
       return true;
@@ -233,8 +215,8 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
   CodeCompletionCallbacksImpl(Parser &P,
                               CodeCompletionContext &CompletionContext,
                               CodeCompletionConsumer &Consumer)
-      : CodeCompletionCallbacks(P), DoneParsingCallback(),
-        CompletionContext(CompletionContext), Consumer(Consumer) {}
+      : CodeCompletionCallbacks(P), CompletionContext(CompletionContext),
+        Consumer(Consumer) {}
 
   void setAttrTargetDeclKind(std::optional<DeclKind> DK) override {
     if (DK == DeclKind::PatternBinding)
@@ -309,7 +291,7 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks,
   void completeTypeAttrInheritanceBeginning() override;
   void completeOptionalBinding() override;
 
-  void doneParsing(SourceFile *SrcFile) override;
+  void readyForTypeChecking(SourceFile *SrcFile) override;
 
 private:
   void addKeywords(CodeCompletionResultSink &Sink, bool MaybeFuncBody);
@@ -1641,7 +1623,7 @@ void CodeCompletionCallbacksImpl::afterPoundCompletion(SourceLoc CompletionLoc,
   Consumer.handleResults(CompletionContext);
 }
 
-void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
+void CodeCompletionCallbacksImpl::readyForTypeChecking(SourceFile *SrcFile) {
   CompletionContext.CodeCompletionKind = Kind;
 
   if (Kind == CompletionKind::None) {