[prespecialized metadata] Add token for caching.

Emit a once token when adding canonical prespecialized metadata records
to a nominal type descriptor and add the token itself as a trailing
object to the type descriptor.  The new token will, in subsequent
commits, enable the canonical prespecialized metadata records attached
to the type descriptor to be added to the metadata cache exactly once.

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"
@@ -3796,6 +3797,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> {
@@ -3882,7 +3888,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
@@ -4024,7 +4032,8 @@ class TargetClassDescriptor final
                               TargetObjCResilientClassStubInfo<Runtime>,
                               TargetCanonicalSpecializedMetadatasListCount<Runtime>,
                               TargetCanonicalSpecializedMetadatasListEntry<Runtime>,
-                              TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>> {
+                              TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>,
+                              TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>> {
 private:
   using TrailingGenericContextObjects =
     swift::TrailingGenericContextObjects<TargetClassDescriptor<Runtime>,
@@ -4039,7 +4048,8 @@ class TargetClassDescriptor final
                                          TargetObjCResilientClassStubInfo<Runtime>,
                                          TargetCanonicalSpecializedMetadatasListCount<Runtime>,
                                          TargetCanonicalSpecializedMetadatasListEntry<Runtime>,
-                                         TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>>;
+                                         TargetCanonicalSpecializedMetadataAccessorsListEntry<Runtime>,
+                                         TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>>;
 
   using TrailingObjects =
     typename TrailingGenericContextObjects::TrailingObjects;
@@ -4067,6 +4077,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;
@@ -4204,6 +4216,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 {
@@ -4367,6 +4383,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;
   }
@@ -4394,7 +4418,8 @@ class TargetStructDescriptor final
                             TargetForeignMetadataInitialization<Runtime>,
                             TargetSingletonMetadataInitialization<Runtime>,
                             TargetCanonicalSpecializedMetadatasListCount<Runtime>,
-                            TargetCanonicalSpecializedMetadatasListEntry<Runtime>> {
+                            TargetCanonicalSpecializedMetadatasListEntry<Runtime>,
+                            TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>> {
 public:
   using ForeignMetadataInitialization =
     TargetForeignMetadataInitialization<Runtime>;
@@ -4406,6 +4431,8 @@ class TargetStructDescriptor final
     TargetCanonicalSpecializedMetadatasListCount<Runtime>;
   using MetadataListEntry =
     TargetCanonicalSpecializedMetadatasListEntry<Runtime>;
+  using MetadataCachingOnceToken =
+      TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>;
 
 private:
   using TrailingGenericContextObjects =
@@ -4414,7 +4441,8 @@ class TargetStructDescriptor final
                                            ForeignMetadataInitialization,
                                            SingletonMetadataInitialization,
                                            MetadataListCount,
-                                           MetadataListEntry>;
+                                           MetadataListEntry,
+                                           MetadataCachingOnceToken>;
 
   using TrailingObjects =
     typename TrailingGenericContextObjects::TrailingObjects;
@@ -4444,6 +4472,10 @@ class TargetStructDescriptor final
       : 0;
   }
 
+  size_t numTrailingObjects(OverloadToken<MetadataCachingOnceToken>) const {
+    return this->hasCanonicicalMetadataPrespecializations() ? 1 : 0;
+  }
+
 public:
   using TrailingGenericContextObjects::getGenericContext;
   using TrailingGenericContextObjects::getGenericContextHeader;
@@ -4489,6 +4521,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;
   }
@@ -4505,7 +4545,8 @@ class TargetEnumDescriptor final
                             TargetForeignMetadataInitialization<Runtime>,
                             TargetSingletonMetadataInitialization<Runtime>,
                             TargetCanonicalSpecializedMetadatasListCount<Runtime>,
-                            TargetCanonicalSpecializedMetadatasListEntry<Runtime>> {
+                            TargetCanonicalSpecializedMetadatasListEntry<Runtime>,
+                            TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>> {
 public:
   using SingletonMetadataInitialization =
     TargetSingletonMetadataInitialization<Runtime>;
@@ -4517,6 +4558,8 @@ class TargetEnumDescriptor final
     TargetCanonicalSpecializedMetadatasListCount<Runtime>;
   using MetadataListEntry =
     TargetCanonicalSpecializedMetadatasListEntry<Runtime>;
+  using MetadataCachingOnceToken =
+      TargetCanonicalSpecializedMetadatasCachingOnceToken<Runtime>;
 
 private:
   using TrailingGenericContextObjects =
@@ -4525,7 +4568,8 @@ class TargetEnumDescriptor final
                                         ForeignMetadataInitialization,
                                         SingletonMetadataInitialization,
                                         MetadataListCount,
-                                        MetadataListEntry>;
+                                        MetadataListEntry, 
+                                        MetadataCachingOnceToken>;
 
   using TrailingObjects =
     typename TrailingGenericContextObjects::TrailingObjects;
@@ -4555,6 +4599,10 @@ class TargetEnumDescriptor final
       : 0;
   }
 
+  size_t numTrailingObjects(OverloadToken<MetadataCachingOnceToken>) const {
+    return this->hasCanonicicalMetadataPrespecializations() ? 1 : 0;
+  }
+
 public:
   using TrailingGenericContextObjects::getGenericContext;
   using TrailingGenericContextObjects::getGenericContextHeader;
@@ -4614,6 +4662,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;
   }
@@ -4766,6 +4822,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;

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";

lib/IRGen/GenDecl.cpp

@@ -3906,6 +3906,24 @@ IRGenModule::getAddrOfTypeMetadataLazyCacheVariable(CanType type) {
   return variable;
 }
 
+llvm::Constant *
+IRGenModule::getAddrOfCanonicalPrespecializedGenericTypeCachingOnceToken(
+    NominalTypeDecl *decl) {
+  assert(decl->isGenericContext());
+  LinkEntity entity =
+      LinkEntity::forCanonicalPrespecializedGenericTypeCachingOnceToken(decl);
+  if (auto &entry = GlobalVars[entity]) {
+    return entry;
+  }
+  auto variable = getAddrOfLLVMVariable(entity, ForDefinition, DebugTypeInfo());
+
+  // Zero-initialize if we're asking for a definition.
+  cast<llvm::GlobalVariable>(variable)->setInitializer(
+      llvm::ConstantInt::get(OnceTy, 0));
+
+  return variable;
+}
+
 llvm::Constant *
 IRGenModule::getAddrOfNoncanonicalSpecializedGenericTypeMetadataCacheVariable(CanType type) {
   assert(!type->hasArchetype() && !type->hasTypeParameter());

lib/IRGen/GenMeta.cpp

@@ -1329,6 +1329,7 @@ namespace {
     void maybeAddCanonicalMetadataPrespecializations() {
       if (Type->isGenericContext() && hasCanonicalMetadataPrespecializations()) {
         asImpl().addCanonicalMetadataPrespecializations();
+        asImpl().addCanonicalMetadataPrespecializationCachingOnceToken();
       }
     }
 
@@ -1348,6 +1349,12 @@ namespace {
       }
     }
 
+    void addCanonicalMetadataPrespecializationCachingOnceToken() {
+      auto *cachingOnceToken =
+          IGM.getAddrOfCanonicalPrespecializedGenericTypeCachingOnceToken(Type);
+      B.addRelativeAddress(cachingOnceToken);
+    }
+
     // Subclasses should provide:
     // ContextDescriptorKind getContextKind();
     // void addLayoutInfo();

lib/IRGen/IRGenMangler.h

@@ -122,6 +122,11 @@ class IRGenMangler : public Mangle::ASTMangler {
     return mangleTypeSymbol(type, "MN");
   }
 
+  std::string mangleCanonicalPrespecializedGenericTypeCachingOnceToken(
+      const NominalTypeDecl *decl) {
+    return mangleNominalTypeSymbol(decl, "Mz");
+  }
+
   std::string mangleNoncanonicalSpecializedGenericTypeMetadataCache(Type type) {
     return mangleTypeSymbol(type, "MJ");
   }