[AST] RuntimeMetadata: Move getRuntimeDiscoverableAttrs to Decl

Reflection metadata attributes could be attached to a `ValueDecl`
and `ExtensionDecl`.

Author: xedin

include/swift/AST/Decl.h

@@ -859,6 +859,10 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     return Attrs;
   }
 
+  /// Retrieve runtime discoverable attributes (if any) associated
+  /// with this declaration.
+  ArrayRef<CustomAttr *> getRuntimeDiscoverableAttrs() const;
+
   /// Returns the semantic attributes attached to this declaration,
   /// including attributes that are generated as the result of member
   /// attribute macro expansion.
@@ -2851,9 +2855,6 @@ class ValueDecl : public Decl {
   GenericParameterReferenceInfo findExistentialSelfReferences(
       Type baseTy, bool treatNonResultCovariantSelfAsInvariant) const;
 
-  /// Retrieve runtime discoverable attributes (if any) associated
-  /// with this declaration.
-  ArrayRef<CustomAttr *> getRuntimeDiscoverableAttrs() const;
   /// Retrieve a nominal type declaration backing given runtime discoverable
   /// attribute.
   ///

include/swift/AST/TypeCheckRequests.h

@@ -3926,15 +3926,15 @@ class ExternalMacroDefinitionRequest
 
 class GetRuntimeDiscoverableAttributes
     : public SimpleRequest<GetRuntimeDiscoverableAttributes,
-                           ArrayRef<CustomAttr *>(ValueDecl *),
+                           ArrayRef<CustomAttr *>(Decl *),
                            RequestFlags::Cached> {
 public:
   using SimpleRequest::SimpleRequest;
 
 private:
   friend SimpleRequest;
 
-  ArrayRef<CustomAttr *> evaluate(Evaluator &evaluator, ValueDecl *decl) const;
+  ArrayRef<CustomAttr *> evaluate(Evaluator &evaluator, Decl *decl) const;
 
 public:
   bool isCached() const { return true; }

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -468,7 +468,7 @@ SWIFT_REQUEST(TypeChecker, SynthesizeRuntimeMetadataAttrGeneratorBody,
               BraceStmt *(CustomAttr *, ValueDecl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, GetRuntimeDiscoverableAttributes,
-              ArrayRef<CustomAttr *>(ValueDecl *),
+              ArrayRef<CustomAttr *>(Decl *),
               Cached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, LocalDiscriminatorsRequest,
               unsigned(DeclContext *),

lib/AST/ASTPrinter.cpp

@@ -1883,6 +1883,11 @@ bool isNonSendableExtension(const Decl *D) {
 bool ShouldPrintChecker::shouldPrint(const Decl *D,
                                      const PrintOptions &Options) {
   if (auto *ED = dyn_cast<ExtensionDecl>(D)) {
+    // Always print unavilable extensions that carry reflection
+    // metadata attributes.
+    if (!ED->getRuntimeDiscoverableAttrs().empty())
+      return true;
+
     if (Options.printExtensionContentAsMembers(ED))
       return false;
   }

lib/AST/Decl.cpp

@@ -9888,8 +9888,8 @@ ValueDecl::getRuntimeDiscoverableAttrTypeDecl(CustomAttr *attr) const {
   return nominal;
 }
 
-ArrayRef<CustomAttr *> ValueDecl::getRuntimeDiscoverableAttrs() const {
-  auto *mutableSelf = const_cast<ValueDecl *>(this);
+ArrayRef<CustomAttr *> Decl::getRuntimeDiscoverableAttrs() const {
+  auto *mutableSelf = const_cast<Decl *>(this);
   return evaluateOrDefault(getASTContext().evaluator,
                            GetRuntimeDiscoverableAttributes{mutableSelf},
                            nullptr);

lib/Sema/TypeCheckAttr.cpp

@@ -7463,7 +7463,6 @@ static void forEachCustomAttribute(
     auto *nominal = evaluateOrDefault(
         ctx.evaluator,
         CustomAttrNominalRequest{mutableAttr, decl->getDeclContext()}, nullptr);
-
     if (!nominal)
       continue;
 
@@ -7474,15 +7473,15 @@ static void forEachCustomAttribute(
 
 ArrayRef<CustomAttr *>
 GetRuntimeDiscoverableAttributes::evaluate(Evaluator &evaluator,
-                                           ValueDecl *decl) const {
+                                           Decl *decl) const {
   auto &ctx = decl->getASTContext();
 
   llvm::SmallMapVector<NominalTypeDecl *, CustomAttr *, 4> attrs;
 
   enum class GatheringMode { Direct, Inference };
 
   auto gatherRuntimeAttrsOnDecl =
-      [&](ValueDecl *decl,
+      [&](Decl *decl,
           llvm::SmallMapVector<NominalTypeDecl *, CustomAttr *, 4> &attrs,
           GatheringMode mode) {
         forEachCustomAttribute<RuntimeMetadataAttr>(
@@ -7514,6 +7513,19 @@ GetRuntimeDiscoverableAttributes::evaluate(Evaluator &evaluator,
     return copy;
   };
 
+  // Gather reflection metadata attributes only if this extension is:
+  //  - unavailable;
+  //  - unconstrained;
+  //  - declared in the same module as the extended type.
+  if (auto *ED = dyn_cast<ExtensionDecl>(decl)) {
+    if (!AvailableAttr::isUnavailable(ED))
+      return copyAttrs(attrs);
+
+    if (ED->isConstrainedExtension() ||
+        ED->getParentModule() != decl->getDeclContext()->getParentModule())
+      return copyAttrs(attrs);
+  }
+
   // First, gather all of the runtime attributes directly on the decl.
   gatherRuntimeAttrsOnDecl(decl, attrs, GatheringMode::Direct);