[Runtime] Don't attempt to look up prespecialized metadata involving pointers outside the shared cache.

The descriptor and arguments for prespecialized metadata will always be in the shared cache. Skip creating the mangling for any lookup involving pointers outside the shared cache, as an optimization.

Author: mikeash

stdlib/public/runtime/LibPrespecialized.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/Runtime/LibPrespecialized.h"
+#include "MetadataCache.h"
 #include "Private.h"
 #include "swift/Basic/Lazy.h"
 #include "swift/Runtime/EnvironmentVariables.h"
@@ -71,24 +72,89 @@ static const LibPrespecializedData<InProcess> *findLibPrespecialized() {
   return data;
 }
 
+struct LibPrespecializedState {
+  struct AddressRange {
+    uintptr_t start, end;
+
+    bool contains(const void *ptr) {
+      return start <= (uintptr_t)ptr && (uintptr_t)ptr < end;
+    }
+  };
+
+  bool loggingEnabled;
+  const LibPrespecializedData<InProcess> *data;
+  AddressRange sharedCacheRange{0, 0};
+  AddressRange metadataAllocatorInitialPoolRange{0, 0};
+
+  LibPrespecializedState() {
+    loggingEnabled =
+        runtime::environment::SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING();
+    data = findLibPrespecialized();
+
+#if DYLD_GET_SWIFT_PRESPECIALIZED_DATA_DEFINED
+    size_t sharedCacheLength;
+    sharedCacheRange.start =
+        (uintptr_t)_dyld_get_shared_cache_range(&sharedCacheLength);
+    sharedCacheRange.end = sharedCacheRange.start + sharedCacheLength;
+
+    auto [initialPoolStart, initialPoolLength] =
+        MetadataAllocator::InitialPoolLocation();
+    metadataAllocatorInitialPoolRange.start = (uintptr_t)initialPoolStart;
+    metadataAllocatorInitialPoolRange.end =
+        metadataAllocatorInitialPoolRange.start + initialPoolLength;
+#endif
+  }
+};
+
+static Lazy<LibPrespecializedState> LibPrespecialized;
+
+#define LOG(fmt, ...)                                                          \
+  do {                                                                         \
+    if (SWIFT_UNLIKELY(prespecialized.loggingEnabled))                         \
+      fprintf(stderr, "Prespecializations library: " fmt "\n", __VA_ARGS__);   \
+  } while (0)
+
 const LibPrespecializedData<InProcess> *swift::getLibPrespecializedData() {
   return SWIFT_LAZY_CONSTANT(findLibPrespecialized());
 }
 
 // Returns true if the type has any arguments that aren't plain types (packs or
 // unknown kinds).
-static bool hasNonTypeGenericArguments(const TypeContextDescriptor *description) {
-  auto generics = description->getGenericContext();
-  if (!generics)
-    return false;
-
+static bool
+hasNonTypeGenericArguments(const TargetGenericContext<InProcess> *generics) {
   for (auto param : generics->getGenericParams())
     if (param.getKind() != GenericParamKind::Type)
       return true;
 
   return false;
 }
 
+static bool
+isPotentialPrespecializedPointer(LibPrespecializedState &prespecialized,
+                                 const void *pointer) {
+  // Prespecialized metadata descriptors and arguments are always in the shared
+  // cache. They're either statically emitted metadata, or they're
+  // prespecialized metadata. Anything that's dynamically allocated, or
+  // statically allocated outside the shared cache, is not a possible candidate.
+
+  // If we're loading a debug libprespecialized, we can't do these checks, so
+  // just say everything is a potential argument. Performance is not so
+  // important in that case.
+  if (!prespecialized.sharedCacheRange.contains(prespecialized.data))
+    return true;
+
+  // Anything outside the shared cache isn't a potential argument.
+  if (!prespecialized.sharedCacheRange.contains(pointer))
+    return false;
+
+  // Dynamically allocated metadata could be within the shared cache, in the
+  // initial metadata allocation pool. Reject anything in that region.
+  if (prespecialized.metadataAllocatorInitialPoolRange.contains(pointer))
+    return false;
+
+  return true;
+}
+
 static bool disableForValidation = false;
 
 Metadata *
@@ -97,24 +163,40 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
   if (disableForValidation)
     return nullptr;
 
-  auto *data = getLibPrespecializedData();
+  auto &prespecialized = LibPrespecialized.get();
+
+  auto *data = prespecialized.data;
   if (!data)
     return nullptr;
 
+  auto *generics = description->getGenericContext();
+  if (!generics)
+    return nullptr;
+
   // We don't support types with pack parameters yet (and especially not types
   // with unknown parameter kinds) so don't even try to look those up.
-  if (hasNonTypeGenericArguments(description))
+  if (hasNonTypeGenericArguments(generics))
     return nullptr;
 
+  if (!isPotentialPrespecializedPointer(prespecialized, description)) {
+    LOG("Rejecting descriptor %p, not in the shared cache",
+        (const void *)description);
+    return nullptr;
+  }
+
+  auto numKeyArguments = generics->getGenericContextHeader().NumKeyArguments;
+  for (unsigned i = 0; i < numKeyArguments; i++) {
+    if (!isPotentialPrespecializedPointer(prespecialized, arguments[i])) {
+      LOG("Rejecting argument %u %p to descriptor %p, not in the shared cache",
+          i, arguments[i], (const void *)description);
+      return nullptr;
+    }
+  }
+
   Demangler dem;
   auto mangleNode = _buildDemanglingForGenericType(description, arguments, dem);
   if (!mangleNode) {
-    if (SWIFT_UNLIKELY(runtime::environment::
-                           SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING()))
-      fprintf(stderr,
-              "Prespecializations library: failed to build demangling with "
-              "descriptor %p.\n",
-              description);
+    LOG("failed to build demangling with descriptor %p.", description);
     return nullptr;
   }
 
@@ -142,10 +224,7 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
   auto *metadataMap = data->getMetadataMap();
   auto *element = metadataMap->find(key.data(), key.size());
   auto *result = element ? element->value : nullptr;
-  if (SWIFT_UNLIKELY(runtime::environment::
-                         SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING()))
-    fprintf(stderr, "Prespecializations library: found %p for key '%.*s'.\n",
-            result, (int)key.size(), key.data());
+  LOG("found %p for key '%.*s'.", result, (int)key.size(), key.data());
   return result;
 }
 

stdlib/public/runtime/Metadata.cpp

@@ -7420,6 +7420,10 @@ static swift::atomic<PoolRange>
 AllocationPool{PoolRange{InitialAllocationPool.Pool,
                          sizeof(InitialAllocationPool.Pool)}};
 
+std::tuple<const void *, size_t> MetadataAllocator::InitialPoolLocation() {
+  return {InitialAllocationPool.Pool, sizeof(InitialAllocationPool.Pool)};
+}
+
 bool swift::_swift_debug_metadataAllocationIterationEnabled = false;
 const void * const swift::_swift_debug_allocationPoolPointer = &AllocationPool;
 std::atomic<const void *> swift::_swift_debug_metadataAllocationBacktraceList;

stdlib/public/runtime/MetadataCache.h

@@ -25,6 +25,7 @@
 #include <atomic>
 #include <condition_variable>
 #include <optional>
+#include <tuple>
 
 #ifndef SWIFT_DEBUG_RUNTIME
 #define SWIFT_DEBUG_RUNTIME 0
@@ -46,6 +47,11 @@ class MetadataAllocator : public llvm::AllocatorBase<MetadataAllocator> {
 
   void Reset() {}
 
+  /// Get the location of the allocator's initial statically allocated pool.
+  /// The return values are start and size. If there is no statically allocated
+  /// pool, the return values are NULL, 0.
+  static std::tuple<const void *, size_t> InitialPoolLocation();
+
   SWIFT_RETURNS_NONNULL SWIFT_NODISCARD
   void *Allocate(size_t size, size_t alignment);
   using AllocatorBase<MetadataAllocator>::Allocate;
@@ -67,6 +73,9 @@ class MetadataAllocator : public llvm::AllocatorBase<MetadataAllocator> {
 class MetadataAllocator {
 public:
   MetadataAllocator(uint16_t tag) {}
+  static std::tuple<const void *, size_t> InitialPoolLocation() {
+    return {nullptr, 0};
+  }
   SWIFT_RETURNS_NONNULL SWIFT_NODISCARD
   void *Allocate(size_t size, size_t alignment) {
     if (alignment < sizeof(void*)) alignment = sizeof(void*);