[Runtime] Add prespecialized-caching getGenericMetadata variant.

Add a new entry point for getting generic metadata which adds the
canonical metadata records attached to the nominal type descriptor to
the metadata cache.

Change the implementation of the primary entry-point
swift_getGenericMetadata to stop looking through canonical
prespecialized records.

Change the implementation of swift_getCanonicalSpecializedMetadata to
use the caching token attached to the nominal type descriptor to add
canonical prespecialized metadata records to the metadata cache only
once rather than using the cache variables to limit the number of times
the attempt was made.

Author: nate-chandler

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/IRGen/IRGenModule.cpp

@@ -697,6 +697,16 @@ namespace RuntimeConstants {
     return RuntimeAvailability::AlwaysAvailable;
   }
 
+  RuntimeAvailability
+  GetCanonicalPrespecializedGenericMetadataAvailability(ASTContext &context) {
+    auto featureAvailability =
+        context.getPrespecializedGenericMetadataAvailability();
+    if (!isDeploymentAvailabilityContainedIn(context, featureAvailability)) {
+      return RuntimeAvailability::ConditionallyAvailable;
+    }
+    return RuntimeAvailability::AlwaysAvailable;
+  }
+
   RuntimeAvailability ConcurrencyAvailability(ASTContext &context) {
     auto featureAvailability = context.getConcurrencyAvailability();
     if (!isDeploymentAvailabilityContainedIn(context, featureAvailability)) {

stdlib/public/runtime/Metadata.cpp

@@ -706,24 +706,13 @@ swift::swift_allocateGenericValueMetadata(const ValueTypeDescriptor *description
   return metadata;
 }
 
-MetadataResponse swift::swift_getCanonicalSpecializedMetadata(
-    MetadataRequest request, const Metadata *candidate,
-    const Metadata **cacheMetadataPtr) {
-  assert(candidate->isStaticallySpecializedGenericMetadata() &&
-         !candidate->isCanonicalStaticallySpecializedGenericMetadata());
-  auto *description = candidate->getDescription();
-  assert(description);
-
-  using CachedMetadata = std::atomic<const Metadata *>;
-  auto cachedMetadataAddr = ((CachedMetadata *)cacheMetadataPtr);
-  auto *cachedMetadata = cachedMetadataAddr->load(SWIFT_MEMORY_ORDER_CONSUME);
-  if (SWIFT_LIKELY(cachedMetadata != nullptr)) {
-    // Cached metadata pointers are always complete.
-    return MetadataResponse{(const Metadata *)cachedMetadata,
-                            MetadataState::Complete};
-  }
-
+// Look into the canonical prespecialized metadata attached to the type
+// descriptor and add them to the metadata cache.
+static void
+_cacheCanonicalSpecializedMetadata(const TypeContextDescriptor *description) {
   auto &cache = getCache(*description);
+  auto request =
+      MetadataRequest(MetadataState::Complete, /*isNonBlocking*/ true);
   assert(description->getFullGenericContextHeader().Base.NumKeyArguments ==
          cache.NumKeyParameters + cache.NumWitnessTables);
   if (auto *classDescription = dyn_cast<ClassDescriptor>(description)) {
@@ -753,8 +742,47 @@ MetadataResponse swift::swift_getCanonicalSpecializedMetadata(
       assert(result.second.Value == canonicalMetadata);
     }
   }
+}
+
+static void
+cacheCanonicalSpecializedMetadata(const TypeContextDescriptor *description,
+                                  swift_once_t *token) {
+  swift_once(
+      token,
+      [](void *uncastDescription) {
+        auto *description = (const TypeContextDescriptor *)uncastDescription;
+        _cacheCanonicalSpecializedMetadata(description);
+      },
+      (void *)description);
+}
+
+MetadataResponse swift::swift_getCanonicalSpecializedMetadata(
+    MetadataRequest request, const Metadata *candidate,
+    const Metadata **cacheMetadataPtr) {
+  assert(candidate->isStaticallySpecializedGenericMetadata() &&
+         !candidate->isCanonicalStaticallySpecializedGenericMetadata());
+  auto *description = candidate->getDescription();
+  assert(description);
+
+  using CachedMetadata = std::atomic<const Metadata *>;
+  auto cachedMetadataAddr = ((CachedMetadata *)cacheMetadataPtr);
+  auto *cachedMetadata = cachedMetadataAddr->load(SWIFT_MEMORY_ORDER_CONSUME);
+  if (SWIFT_LIKELY(cachedMetadata != nullptr)) {
+    // Cached metadata pointers are always complete.
+    return MetadataResponse{(const Metadata *)cachedMetadata,
+                            MetadataState::Complete};
+  }
+
+  if (auto *token =
+          description
+              ->getCanonicalMetadataPrespecializationCachingOnceToken()) {
+    cacheCanonicalSpecializedMetadata(description, token);
+    // NOTE: If there is no token, then there are no canonical prespecialized
+    //       metadata records, either.
+  }
   const void *const *arguments =
       reinterpret_cast<const void *const *>(candidate->getGenericArgs());
+  auto &cache = getCache(*description);
   auto key = MetadataCacheKey(cache.NumKeyParameters, cache.NumWitnessTables,
                               arguments);
   auto result = cache.getOrInsert(key, request, candidate);
@@ -764,64 +792,38 @@ MetadataResponse swift::swift_getCanonicalSpecializedMetadata(
   return result.second;
 }
 
-// Look into the canonical prespecialized metadata attached to the type
-// descriptor and return matching records, if any.
-static Metadata *
-findCanonicalSpecializedMetadata(MetadataRequest request,
-                                 const void *const *arguments,
-                                 const TypeContextDescriptor *description) {
+SWIFT_CC(swift)
+static MetadataResponse
+_swift_getGenericMetadata(MetadataRequest request, const void *const *arguments,
+                          const TypeContextDescriptor *description) {
   auto &cache = getCache(*description);
+  assert(description->getFullGenericContextHeader().Base.NumKeyArguments ==
+         cache.NumKeyParameters + cache.NumWitnessTables);
   auto key = MetadataCacheKey(cache.NumKeyParameters, cache.NumWitnessTables,
                               arguments);
-  auto prespecializedMetadatas =
-      description->getCanonicicalMetadataPrespecializations();
-  int index = 0;
-  for (auto &prespecializedMetadataPtr : prespecializedMetadatas) {
-    Metadata *prespecializationMetadata = prespecializedMetadataPtr.get();
-    const void *const *prespecializationArguments =
-        reinterpret_cast<const void *const *>(
-            prespecializationMetadata->getGenericArgs());
-    auto prespecializationKey =
-        MetadataCacheKey(cache.NumKeyParameters, cache.NumWitnessTables,
-                         prespecializationArguments);
-    if (key == prespecializationKey) {
-      if (auto *classDescription = dyn_cast<ClassDescriptor>(description)) {
-        auto canonicalMetadataAccessors =
-            classDescription->getCanonicalMetadataPrespecializationAccessors();
-        auto &canonicalMetadataAccessorPtr = canonicalMetadataAccessors[index];
-        auto *canonicalMetadataAccessor = canonicalMetadataAccessorPtr.get();
-        auto response = canonicalMetadataAccessor(request);
-        return const_cast<Metadata *>(response.Value);
-      } else {
-        return prespecializationMetadata;
-      }
-    }
-    ++index;
-  }
-  return nullptr;
+  auto result = cache.getOrInsert(key, request, description, arguments);
+
+  return result.second;
 }
 
 /// The primary entrypoint.
 MetadataResponse
 swift::swift_getGenericMetadata(MetadataRequest request,
-                                const void * const *arguments,
+                                const void *const *arguments,
                                 const TypeContextDescriptor *description) {
-  description = swift_auth_data_non_address(description, SpecialPointerAuthDiscriminators::TypeDescriptor);
-  if (auto *prespecialization =
-          findCanonicalSpecializedMetadata(request, arguments, description)) {
-    return {prespecialization, MetadataState::Complete};
-  }
-  auto &cache = getCache(*description);
-  assert(description->getFullGenericContextHeader().Base.NumKeyArguments ==
-         cache.NumKeyParameters + cache.NumWitnessTables);
-  auto key = MetadataCacheKey(cache.NumKeyParameters, cache.NumWitnessTables,
-                              arguments);
-  auto result = cache.getOrInsert(key, request, description, arguments);
+  description = swift_auth_data_non_address(
+      description, SpecialPointerAuthDiscriminators::TypeDescriptor);
+  return _swift_getGenericMetadata(request, arguments, description);
+}
 
-  assert(
-      !result.second.Value->isCanonicalStaticallySpecializedGenericMetadata());
+MetadataResponse swift::swift_getCanonicalPrespecializedGenericMetadata(
+    MetadataRequest request, const void *const *arguments,
+    const TypeContextDescriptor *description, swift_once_t *token) {
+  description = swift_auth_data_non_address(
+      description, SpecialPointerAuthDiscriminators::TypeDescriptor);
+  cacheCanonicalSpecializedMetadata(description, token);
 
-  return result.second;
+  return _swift_getGenericMetadata(request, arguments, description);
 }
 
 /***************************************************************************/