Merge pull request #34490 from nate-chandler/generic-metadata-prespecialization-components/special-entry

[metadata prespecialization] Add canonical records to cache.

Author: nate-chandler

docs/ABI/Mangling.rst

@@ -175,6 +175,7 @@ Globals
 
   global ::= global 'MJ'                 // noncanonical specialized generic type metadata instantiation cache associated with global
   global ::= global 'MN'                 // noncanonical specialized generic type metadata for global
+  global ::= global 'Mz'                 // canonical specialized generic type metadata caching token
 
   #if SWIFT_RUNTIME_VERSION >= 5.4
     global ::= context (decl-name '_')+ 'WZ' // global variable one-time initialization function

include/swift/ABI/Metadata.h

@@ -30,6 +30,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "swift/Strings.h"
 #include "swift/Runtime/Config.h"
+#include "swift/Runtime/Once.h"
 #include "swift/ABI/MetadataValues.h"
 #include "swift/ABI/System.h"
 #include "swift/ABI/TrailingObjects.h"
@@ -3787,6 +3788,11 @@ struct TargetCanonicalSpecializedMetadataAccessorsListEntry {
   TargetRelativeDirectPointer<Runtime, MetadataResponse(MetadataRequest), /*Nullable*/ false> accessor;
 };
 
+template <typename Runtime>
+struct TargetCanonicalSpecializedMetadatasCachingOnceToken {
+  TargetRelativeDirectPointer<Runtime, swift_once_t, /*Nullable*/ false> token;
+};
+
 template <typename Runtime>
 class TargetTypeContextDescriptor
     : public TargetContextDescriptor<Runtime> {
@@ -3873,7 +3879,9 @@ class TargetTypeContextDescriptor
   }
 
   const llvm::ArrayRef<TargetRelativeDirectPointer<Runtime, TargetMetadata<Runtime>, /*Nullable*/ false>>
-  getCanonicicalMetadataPrespecializations() const; 
+  getCanonicicalMetadataPrespecializations() const;
+
+  swift_once_t *getCanonicalMetadataPrespecializationCachingOnceToken() const;
 
   static bool classof(const TargetContextDescriptor<Runtime> *cd) {
     return cd->getKind() >= ContextDescriptorKind::Type_First
@@ -4015,7 +4023,8 @@ class TargetClassDescriptor final
                               TargetObjCResilientClassStubInfo<Runtime>,
                               TargetCanonicalSpecializedMetadatasListCount<Runtime>,
                               TargetCanonicalSpecializedMetadatasListEntry<Runtime>,
-                              TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>> {
+                              TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>,
+                              TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>> {
 private:
   using TrailingGenericContextObjects =
     swift::TrailingGenericContextObjects<TargetClassDescriptor<Runtime>,
@@ -4030,7 +4039,8 @@ class TargetClassDescriptor final
                                          TargetObjCResilientClassStubInfo<Runtime>,
                                          TargetCanonicalSpecializedMetadatasListCount<Runtime>,
                                          TargetCanonicalSpecializedMetadatasListEntry<Runtime>,
-                                         TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>>;
+                                         TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>,
+                                         TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>>;
 
   using TrailingObjects =
     typename TrailingGenericContextObjects::TrailingObjects;
@@ -4058,6 +4068,8 @@ class TargetClassDescriptor final
     TargetRelativeDirectPointer<Runtime, MetadataResponse(MetadataRequest), /*Nullable*/ false>;
   using MetadataAccessorListEntry =
       TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>;
+  using MetadataCachingOnceToken =
+      TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>;
 
   using StoredPointer = typename Runtime::StoredPointer;
   using StoredPointerDifference = typename Runtime::StoredPointerDifference;
@@ -4195,6 +4207,10 @@ class TargetClassDescriptor final
       : 0;
   }
 
+  size_t numTrailingObjects(OverloadToken<MetadataCachingOnceToken>) const {
+    return this->hasCanonicicalMetadataPrespecializations() ? 1 : 0;
+  }
+
 public:
   const TargetRelativeDirectPointer<Runtime, const void, /*nullable*/true> &
   getResilientSuperclass() const {
@@ -4358,6 +4374,14 @@ class TargetClassDescriptor final
         );
   }
 
+  swift_once_t *getCanonicalMetadataPrespecializationCachingOnceToken() const {
+    if (!this->hasCanonicicalMetadataPrespecializations()) {
+      return nullptr;
+    }
+    auto box = this->template getTrailingObjects<MetadataCachingOnceToken>();
+    return box->token.get();
+  }
+
   static bool classof(const TargetContextDescriptor<Runtime> *cd) {
     return cd->getKind() == ContextDescriptorKind::Class;
   }
@@ -4385,7 +4409,8 @@ class TargetStructDescriptor final
                             TargetForeignMetadataInitialization<Runtime>,
                             TargetSingletonMetadataInitialization<Runtime>,
                             TargetCanonicalSpecializedMetadatasListCount<Runtime>,
-                            TargetCanonicalSpecializedMetadatasListEntry<Runtime>> {
+                            TargetCanonicalSpecializedMetadatasListEntry<Runtime>,
+                            TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>> {
 public:
   using ForeignMetadataInitialization =
     TargetForeignMetadataInitialization<Runtime>;
@@ -4397,6 +4422,8 @@ class TargetStructDescriptor final
     TargetCanonicalSpecializedMetadatasListCount<Runtime>;
   using MetadataListEntry =
     TargetCanonicalSpecializedMetadatasListEntry<Runtime>;
+  using MetadataCachingOnceToken =
+      TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>;
 
 private:
   using TrailingGenericContextObjects =
@@ -4405,7 +4432,8 @@ class TargetStructDescriptor final
                                            ForeignMetadataInitialization,
                                            SingletonMetadataInitialization,
                                            MetadataListCount,
-                                           MetadataListEntry>;
+                                           MetadataListEntry,
+                                           MetadataCachingOnceToken>;
 
   using TrailingObjects =
     typename TrailingGenericContextObjects::TrailingObjects;
@@ -4435,6 +4463,10 @@ class TargetStructDescriptor final
       : 0;
   }
 
+  size_t numTrailingObjects(OverloadToken<MetadataCachingOnceToken>) const {
+    return this->hasCanonicicalMetadataPrespecializations() ? 1 : 0;
+  }
+
 public:
   using TrailingGenericContextObjects::getGenericContext;
   using TrailingGenericContextObjects::getGenericContextHeader;
@@ -4480,6 +4512,14 @@ class TargetStructDescriptor final
         );
   }
 
+  swift_once_t *getCanonicalMetadataPrespecializationCachingOnceToken() const {
+    if (!this->hasCanonicicalMetadataPrespecializations()) {
+      return nullptr;
+    }
+    auto box = this->template getTrailingObjects<MetadataCachingOnceToken>();
+    return box->token.get();
+  }
+
   static bool classof(const TargetContextDescriptor<Runtime> *cd) {
     return cd->getKind() == ContextDescriptorKind::Struct;
   }
@@ -4496,7 +4536,8 @@ class TargetEnumDescriptor final
                             TargetForeignMetadataInitialization<Runtime>,
                             TargetSingletonMetadataInitialization<Runtime>,
                             TargetCanonicalSpecializedMetadatasListCount<Runtime>,
-                            TargetCanonicalSpecializedMetadatasListEntry<Runtime>> {
+                            TargetCanonicalSpecializedMetadatasListEntry<Runtime>,
+                            TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>> {
 public:
   using SingletonMetadataInitialization =
     TargetSingletonMetadataInitialization<Runtime>;
@@ -4508,6 +4549,8 @@ class TargetEnumDescriptor final
     TargetCanonicalSpecializedMetadatasListCount<Runtime>;
   using MetadataListEntry =
     TargetCanonicalSpecializedMetadatasListEntry<Runtime>;
+  using MetadataCachingOnceToken =
+      TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>;
 
 private:
   using TrailingGenericContextObjects =
@@ -4516,7 +4559,8 @@ class TargetEnumDescriptor final
                                         ForeignMetadataInitialization,
                                         SingletonMetadataInitialization,
                                         MetadataListCount,
-                                        MetadataListEntry>;
+                                        MetadataListEntry, 
+                                        MetadataCachingOnceToken>;
 
   using TrailingObjects =
     typename TrailingGenericContextObjects::TrailingObjects;
@@ -4546,6 +4590,10 @@ class TargetEnumDescriptor final
       : 0;
   }
 
+  size_t numTrailingObjects(OverloadToken<MetadataCachingOnceToken>) const {
+    return this->hasCanonicicalMetadataPrespecializations() ? 1 : 0;
+  }
+
 public:
   using TrailingGenericContextObjects::getGenericContext;
   using TrailingGenericContextObjects::getGenericContextHeader;
@@ -4605,6 +4653,14 @@ class TargetEnumDescriptor final
         );
   }
 
+  swift_once_t *getCanonicalMetadataPrespecializationCachingOnceToken() const {
+    if (!this->hasCanonicicalMetadataPrespecializations()) {
+      return nullptr;
+    }
+    auto box = this->template getTrailingObjects<MetadataCachingOnceToken>();
+    return box->token.get();
+  }
+
   static bool classof(const TargetContextDescriptor<Runtime> *cd) {
     return cd->getKind() == ContextDescriptorKind::Enum;
   }
@@ -4757,6 +4813,24 @@ TargetTypeContextDescriptor<Runtime>::getCanonicicalMetadataPrespecializations()
   }
 }
 
+template <typename Runtime>
+inline swift_once_t *TargetTypeContextDescriptor<
+    Runtime>::getCanonicalMetadataPrespecializationCachingOnceToken() const {
+  switch (this->getKind()) {
+  case ContextDescriptorKind::Enum:
+    return llvm::cast<TargetEnumDescriptor<Runtime>>(this)
+        ->getCanonicalMetadataPrespecializationCachingOnceToken();
+  case ContextDescriptorKind::Struct:
+    return llvm::cast<TargetStructDescriptor<Runtime>>(this)
+        ->getCanonicalMetadataPrespecializationCachingOnceToken();
+  case ContextDescriptorKind::Class:
+    return llvm::cast<TargetClassDescriptor<Runtime>>(this)
+        ->getCanonicalMetadataPrespecializationCachingOnceToken();
+  default:
+    swift_unreachable("Not a type context descriptor.");
+  }
+}
+
 /// An entry in the chain of dynamic replacement functions.
 struct DynamicReplacementChainEntry {
   void *implementationFunction;

include/swift/Demangling/DemangleNodes.def

@@ -302,6 +302,7 @@ NODE(NoncanonicalSpecializedGenericTypeMetadataCache)
 NODE(GlobalVariableOnceFunction)
 NODE(GlobalVariableOnceToken)
 NODE(GlobalVariableOnceDeclList)
+NODE(CanonicalPrespecializedGenericTypeCachingOnceToken)
 
 #undef CONTEXT_NODE
 #undef NODE

include/swift/IRGen/Linking.h

@@ -286,6 +286,13 @@ class LinkEntity {
     /// The pointer is a AbstractStorageDecl*.
     DynamicallyReplaceableFunctionImpl,
 
+    /// The once token used by cacheCanonicalSpecializedMetadata, by way of
+    /// swift_getCanonicalSpecializedMetadata and
+    /// swift_getCanonicalPrespecializedGenericMetadata, to
+    /// ensure that canonical prespecialized generic records are only added to
+    /// the metadata cache once.
+    CanonicalPrespecializedGenericTypeCachingOnceToken,
+
     /// The pointer is a SILFunction*.
     DynamicallyReplaceableFunctionKey,
 
@@ -399,7 +406,7 @@ class LinkEntity {
     /// The pointer is a canonical TypeBase*.
     NoncanonicalSpecializedGenericTypeMetadata,
 
-    /// A cache variable for noncanonical specialized type metadata, to be 
+    /// A cache variable for noncanonical specialized type metadata, to be
     /// passed to swift_getCanonicalSpecializedMetadata.
     /// The pointer is a canonical TypeBase*.
     NoncanonicalSpecializedGenericTypeMetadataCacheVariable,
@@ -411,7 +418,7 @@ class LinkEntity {
   }
 
   static bool isDeclKind(Kind k) {
-    return k <= Kind::DynamicallyReplaceableFunctionImpl;
+    return k <= Kind::CanonicalPrespecializedGenericTypeCachingOnceToken;
   }
   static bool isTypeKind(Kind k) {
     return k >= Kind::ProtocolWitnessTableLazyAccessFunction;
@@ -1039,6 +1046,14 @@ class LinkEntity {
     return entity;
   }
 
+  static LinkEntity
+  forCanonicalPrespecializedGenericTypeCachingOnceToken(NominalTypeDecl *decl) {
+    LinkEntity entity;
+    entity.setForDecl(Kind::CanonicalPrespecializedGenericTypeCachingOnceToken,
+                      decl);
+    return entity;
+  }
+
   static LinkEntity
   forSpecializedGenericSwiftMetaclassStub(CanType concreteType) {
     LinkEntity entity;

include/swift/Runtime/Metadata.h

@@ -19,6 +19,7 @@
 
 #include "swift/ABI/Metadata.h"
 #include "swift/Reflection/Records.h"
+#include "swift/Runtime/Once.h"
 
 namespace swift {
 
@@ -340,6 +341,31 @@ SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) MetadataResponse
                                           const Metadata *candidate,
                                           const Metadata **cache);
 
+/// Fetch a uniqued metadata object for the generic nominal type described by
+/// the provided description and arguments, adding the canonical
+/// prespecializations attached to the type descriptor to the metadata cache on
+/// first run.
+///
+/// In contrast to swift_getGenericMetadata, this function is for use by
+/// metadata accessors for which canonical generic metadata has been specialized
+/// at compile time.
+///
+/// Runtime availability: Swift 5.4
+///
+/// \param request A specification of the metadata to be returned.
+/// \param arguments The generic arguments--metadata and witness tables--which
+///                  the returned metadata is to have been instantiated with.
+/// \param description The type descriptor for the generic type whose
+///                    generic metadata is to have been instantiated.
+/// \param token The token that ensures that prespecialized records are added to
+///              the metadata cache only once.
+/// \returns The canonical metadata for the specialized generic type described
+///          by the provided candidate metadata.
+SWIFT_RUNTIME_EXPORT SWIFT_CC(swift) MetadataResponse
+    swift_getCanonicalPrespecializedGenericMetadata(
+        MetadataRequest request, const void *const *arguments,
+        const TypeContextDescriptor *description, swift_once_t *token);
+
 /// Fetch a uniqued metadata object for a generic nominal type.
 SWIFT_RUNTIME_EXPORT SWIFT_CC(swift)
 MetadataResponse

include/swift/Runtime/RuntimeFunctions.def

@@ -653,6 +653,18 @@ FUNCTION(GetCanonicalSpecializedMetadata, swift_getCanonicalSpecializedMetadata,
          ARGS(SizeTy, TypeMetadataPtrTy, TypeMetadataPtrPtrTy),
          ATTRS(NoUnwind, ReadOnly))
 
+// MetadataResponse
+// swift_getCanonicalPrespecializedGenericMetadata(MetadataRequest request,
+//                          const void * const *arguments,
+//                          const TypeContextDescriptor *description,
+//                          swift_once_t *cachedFlag)
+FUNCTION(GetCanonicalPrespecializedGenericMetadata, 
+         swift_getCanonicalPrespecializedGenericMetadata,
+         SwiftCC, GetCanonicalPrespecializedGenericMetadataAvailability,
+         RETURNS(TypeMetadataResponseTy),
+         ARGS(SizeTy, Int8PtrTy, TypeContextDescriptorPtrTy, OnceTy->getPointerTo()),
+         ATTRS(NoUnwind, ReadOnly))
+
 // MetadataResponse swift_getOpaqueTypeMetadata(MetadataRequest request,
 //                                     const void * const *arguments,
 //                                     const OpaqueTypeDescriptor *descriptor,

lib/Demangling/Demangler.cpp

@@ -2001,6 +2001,9 @@ NodePointer Demangler::demangleMetatype() {
                              popNode(isEntity));
     case 'X':
       return demanglePrivateContextDescriptor();
+    case 'z':
+      return createWithPoppedType(
+          Node::Kind::CanonicalPrespecializedGenericTypeCachingOnceToken);
     default:
       return nullptr;
   }

lib/Demangling/NodePrinter.cpp

@@ -561,6 +561,7 @@ class NodePrinter {
     case Node::Kind::GlobalVariableOnceDeclList:
     case Node::Kind::GlobalVariableOnceFunction:
     case Node::Kind::GlobalVariableOnceToken:
+    case Node::Kind::CanonicalPrespecializedGenericTypeCachingOnceToken:
       return false;
     }
     printer_unreachable("bad node kind");
@@ -2537,6 +2538,11 @@ NodePointer NodePrinter::print(NodePointer Node, bool asPrefixContext) {
     Printer << "@objc completion handler block implementation for ";
     print(Node->getChild(0));
     return nullptr;
+  case Node::Kind::CanonicalPrespecializedGenericTypeCachingOnceToken:
+    Printer << "flag for loading of canonical specialized generic type "
+               "metadata for ";
+    print(Node->getChild(0));
+    return nullptr;
   }
   printer_unreachable("bad node kind!");
 }

lib/Demangling/OldRemangler.cpp

@@ -2191,6 +2191,12 @@ void Remangler::mangleNoncanonicalSpecializedGenericTypeMetadataCache(Node *node
   Buffer << "MJ";
 }
 
+void Remangler::mangleCanonicalPrespecializedGenericTypeCachingOnceToken(
+    Node *node) {
+  mangleSingleChildNode(node);
+  Buffer << "Mz";
+}
+
 /// The top-level interface to the remangler.
 std::string Demangle::mangleNodeOld(NodePointer node) {
   if (!node) return "";

lib/Demangling/Remangler.cpp

@@ -2633,6 +2633,12 @@ void Remangler::mangleNoncanonicalSpecializedGenericTypeMetadataCache(Node *node
   Buffer << "MJ";
 }
 
+void Remangler::mangleCanonicalPrespecializedGenericTypeCachingOnceToken(
+    Node *node) {
+  mangleSingleChildNode(node);
+  Buffer << "Mz";
+}
+
 void Remangler::mangleGlobalVariableOnceToken(Node *node) {
   mangleChildNodes(node);
   Buffer << "Wz";