Merge pull request #82129 from slavapestov/bind-extensions-macro

Sema: Don't expand macros when binding extensions

Author: slavapestov

include/swift/AST/Decl.h

@@ -2061,7 +2061,8 @@ class ExtensionDecl final : public GenericContext, public Decl,
   NominalTypeDecl *getExtendedNominal() const;
 
   /// Compute the nominal type declaration that is being extended.
-  NominalTypeDecl *computeExtendedNominal() const;
+  NominalTypeDecl *computeExtendedNominal(
+      bool excludeMacroExpansions=false) const;
 
   /// \c hasBeenBound means nothing if this extension can never been bound
   /// because it is not at the top level.

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 *),
+          ExtendedNominalRequest, NominalTypeDecl *(ExtensionDecl *, bool),
           RequestFlags::SeparatelyCached | RequestFlags::DependencySink> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -276,8 +276,8 @@ class ExtendedNominalRequest
   friend SimpleRequest;
 
   // Evaluation.
-  NominalTypeDecl *
-  evaluate(Evaluator &evaluator, ExtensionDecl *ext) const;
+  NominalTypeDecl * evaluate(Evaluator &evaluator, ExtensionDecl *ext,
+                             bool excludeMacroExpansions) const;
 
 public:
   // Separate caching.

lib/AST/Decl.cpp

@@ -2079,11 +2079,14 @@ NominalTypeDecl *ExtensionDecl::getExtendedNominal() const {
       "Extension must have already been bound (by bindExtensions)");
 }
 
-NominalTypeDecl *ExtensionDecl::computeExtendedNominal() const {
+NominalTypeDecl *ExtensionDecl::computeExtendedNominal(
+    bool excludeMacroExpansions) const {
   ASTContext &ctx = getASTContext();
-  return evaluateOrDefault(
-      ctx.evaluator, ExtendedNominalRequest{const_cast<ExtensionDecl *>(this)},
-      nullptr);
+  return evaluateOrDefault(ctx.evaluator,
+                           ExtendedNominalRequest{
+                               const_cast<ExtensionDecl *>(this),
+                               excludeMacroExpansions},
+                           nullptr);
 }
 
 bool ExtensionDecl::canNeverBeBound() const {

lib/AST/NameLookup.cpp

@@ -1125,6 +1125,9 @@ enum class DirectlyReferencedTypeLookupFlags {
   /// Include members that would normally be excluded because they come from
   /// modules that have not been imported directly.
   IgnoreMissingImports = 1 << 3,
+
+  /// Whenther we should exclude macro expansions.
+  ExcludeMacroExpansions = 1 << 4,
 };
 
 using DirectlyReferencedTypeLookupOptions =
@@ -3065,6 +3068,10 @@ static DirectlyReferencedTypeDecls directReferencesForUnqualifiedTypeLookup(
           DirectlyReferencedTypeLookupFlags::IgnoreMissingImports))
     options |= UnqualifiedLookupFlags::IgnoreMissingImports;
 
+  if (typeLookupOptions.contains(
+          DirectlyReferencedTypeLookupFlags::ExcludeMacroExpansions))
+    options |= UnqualifiedLookupFlags::ExcludeMacroExpansions;
+
   // Manually exclude macro expansions here since the source location
   // is overridden below.
   if (namelookup::isInMacroArgument(dc->getParentSourceFile(), loc))
@@ -3147,6 +3154,10 @@ static llvm::TinyPtrVector<TypeDecl *> directReferencesForQualifiedTypeLookup(
             DirectlyReferencedTypeLookupFlags::IgnoreMissingImports))
       options |= NL_IgnoreMissingImports;
 
+    if (typeLookupOptions.contains(
+            DirectlyReferencedTypeLookupFlags::ExcludeMacroExpansions))
+      options |= NL_ExcludeMacroExpansions;
+
     // Look through the type declarations we were given, resolving them down
     // to nominal type declarations, module declarations, and
     SmallVector<ModuleDecl *, 2> moduleDecls;
@@ -3574,7 +3585,8 @@ ProtocolRequirementsRequest::evaluate(Evaluator &evaluator,
 
 NominalTypeDecl *
 ExtendedNominalRequest::evaluate(Evaluator &evaluator,
-                                 ExtensionDecl *ext) const {
+                                 ExtensionDecl *ext,
+                                 bool excludeMacroExpansions) const {
   auto typeRepr = ext->getExtendedTypeRepr();
   if (!typeRepr) {
     // We must've seen 'extension { ... }' during parsing.
@@ -3583,9 +3595,15 @@ ExtendedNominalRequest::evaluate(Evaluator &evaluator,
 
   ASTContext &ctx = ext->getASTContext();
   auto options = defaultDirectlyReferencedTypeLookupOptions;
+
   if (ext->isInSpecializeExtensionContext()) {
     options |= DirectlyReferencedTypeLookupFlags::AllowUsableFromInline;
   }
+
+  if (excludeMacroExpansions) {
+    options |= DirectlyReferencedTypeLookupFlags::ExcludeMacroExpansions;
+  }
+
   DirectlyReferencedTypeDecls referenced = directReferencesForTypeRepr(
       evaluator, ctx, typeRepr, ext->getParent(), options);
 

lib/ClangImporter/ImportDecl.cpp

@@ -5474,8 +5474,7 @@ namespace {
                       { }, dc, nullptr, decl);
       Impl.SwiftContext.evaluator.cacheOutput(ExtendedTypeRequest{result},
                                               objcClass->getDeclaredType());
-      Impl.SwiftContext.evaluator.cacheOutput(ExtendedNominalRequest{result},
-                                              std::move(objcClass));
+      result->setExtendedNominal(objcClass);
 
       Identifier categoryName;
       if (!decl->getName().empty())
@@ -10256,8 +10255,7 @@ ClangImporter::Implementation::importDeclContextOf(
                                    getClangModuleForDecl(decl), nullptr);
   SwiftContext.evaluator.cacheOutput(ExtendedTypeRequest{ext},
                                      nominal->getDeclaredType());
-  SwiftContext.evaluator.cacheOutput(ExtendedNominalRequest{ext},
-                                     std::move(nominal));
+  ext->setExtendedNominal(nominal);
 
   // Record this extension so we can find it later. We do this early because
   // once we've set the member loader, we don't know when the compiler will use

lib/Frontend/ModuleInterfaceSupport.cpp

@@ -800,8 +800,7 @@ class InheritedProtocolCollector {
 
     ctx.evaluator.cacheOutput(ExtendedTypeRequest{extension},
                               nominal->getDeclaredType());
-    ctx.evaluator.cacheOutput(ExtendedNominalRequest{extension},
-                              const_cast<NominalTypeDecl *>(nominal));
+    extension->setExtendedNominal(const_cast<NominalTypeDecl *>(nominal));
 
     extension->print(printer, printOptions);
     printer << "\n";

lib/Sema/TypeCheckConcurrency.cpp

@@ -7158,8 +7158,7 @@ ProtocolConformance *swift::deriveImplicitSendableConformance(
 
       ctx.evaluator.cacheOutput(ExtendedTypeRequest{extension},
                                 nominal->getDeclaredType());
-      ctx.evaluator.cacheOutput(ExtendedNominalRequest{extension},
-                                std::move(nominal));
+      extension->setExtendedNominal(nominal);
       nominal->addExtension(extension);
 
       // Make it accessible to getTopLevelDecls()

lib/Sema/TypeChecker.cpp

@@ -191,12 +191,14 @@ ModuleDecl *TypeChecker::getStdlibModule(const DeclContext *dc) {
 }
 
 void swift::bindExtensions(ModuleDecl &mod) {
+  bool excludeMacroExpansions = true;
+
   // Utility function to try and resolve the extended type without diagnosing.
   // If we succeed, we go ahead and bind the extension. Otherwise, return false.
   auto tryBindExtension = [&](ExtensionDecl *ext) -> bool {
     assert(!ext->canNeverBeBound() &&
            "Only extensions that can ever be bound get here.");
-    if (auto nominal = ext->computeExtendedNominal()) {
+    if (auto nominal = ext->computeExtendedNominal(excludeMacroExpansions)) {
       nominal->addExtension(ext);
       return true;
     }
@@ -228,20 +230,28 @@ void swift::bindExtensions(ModuleDecl &mod) {
       visitTopLevelDecl(D);
   }
 
-  // Phase 2 - repeatedly go through the worklist and attempt to bind each
-  // extension there, removing it from the worklist if we succeed.
-  bool changed;
-  do {
-    changed = false;
-
-    auto last = std::remove_if(worklist.begin(), worklist.end(),
-                               tryBindExtension);
-    if (last != worklist.end()) {
-      worklist.erase(last, worklist.end());
-      changed = true;
-    }
-  } while(changed);
+  auto tryBindExtensions = [&]() {
+    // Phase 2 - repeatedly go through the worklist and attempt to bind each
+    // extension there, removing it from the worklist if we succeed.
+    bool changed;
+    do {
+      changed = false;
+
+      auto last = std::remove_if(worklist.begin(), worklist.end(),
+                                 tryBindExtension);
+      if (last != worklist.end()) {
+        worklist.erase(last, worklist.end());
+        changed = true;
+      }
+    } while(changed);
+  };
+
+  tryBindExtensions();
 
+  // If that fails, try again, but this time expand macros.
+  excludeMacroExpansions = false;
+  tryBindExtensions();
+  
   // Any remaining extensions are invalid. They will be diagnosed later by
   // typeCheckDecl().
 }

lib/Serialization/Deserialization.cpp

@@ -5465,8 +5465,7 @@ class DeclDeserializer {
     ctx.evaluator.cacheOutput(ExtendedTypeRequest{extension},
                               std::move(extendedType));
     auto nominal = dyn_cast_or_null<NominalTypeDecl>(MF.getDecl(extendedNominalID));
-    ctx.evaluator.cacheOutput(ExtendedNominalRequest{extension},
-                              std::move(nominal));
+    extension->setExtendedNominal(nominal);
 
     if (isImplicit)
       extension->setImplicit();