[Runtime] Add pointer-based lookup of prespecialized metadata.

The existing lookup uses a hash table with the type's mangled name as the key. Building that key is costly. Add a new table that uses pointer keys, so that we can use the descriptor and arguments directly as the key.

rdar://127621414

Author: mikeash

include/swift/Runtime/LibPrespecialized.h

@@ -32,11 +32,25 @@ struct LibPrespecializedData {
 
   TargetPointer<Runtime, const void> metadataMap;
   TargetPointer<Runtime, const void> disabledProcessesTable;
+  TargetPointer<Runtime, const void> pointerKeyedMetadataMap;
+
+  typename Runtime::StoredSize optionFlags;
+
+  // Existing fields are above, add new fields below this point.
 
   static constexpr uint32_t currentMajorVersion = 1;
-  static constexpr uint32_t currentMinorVersion = 2;
+  static constexpr uint32_t currentMinorVersion = 3;
 
   static constexpr uint32_t minorVersionWithDisabledProcessesTable = 2;
+  static constexpr uint32_t minorVersionWithPointerKeyedMetadataMap = 3;
+  static constexpr uint32_t minorVersionWithOptionFlags = 3;
+
+  // Option flags values.
+  enum : typename Runtime::StoredSize {
+    // When this flag is set, the runtime should default to using the
+    // pointer-keyed table. When not set, default to using the name-keyed table.
+    OptionFlagDefaultToPointerKeyedMap = 1ULL << 0,
+  };
 
   // Helpers for retrieving the metadata map in-process.
   static bool stringIsNull(const char *str) { return str == nullptr; }
@@ -52,6 +66,18 @@ struct LibPrespecializedData {
       return nullptr;
     return reinterpret_cast<const char *const *>(disabledProcessesTable);
   }
+
+  const void *getPointerKeyedMetadataMap() const {
+    if (minorVersion < minorVersionWithPointerKeyedMetadataMap)
+      return nullptr;
+    return pointerKeyedMetadataMap;
+  }
+
+  typename Runtime::StoredSize getOptionFlags() const {
+    if (minorVersion < minorVersionWithOptionFlags)
+      return 0;
+    return optionFlags;
+  }
 };
 
 const LibPrespecializedData<InProcess> *getLibPrespecializedData();

stdlib/public/runtime/LibPrespecialized.cpp

@@ -28,6 +28,11 @@
 #include <mach-o/dyld_priv.h>
 #endif
 
+#if __has_include(<os/feature_private.h>)
+#include <os/feature_private.h> // for os_feature_enabled_simple()
+#define HAS_OS_FEATURE 1
+#endif
+
 using namespace swift;
 
 static std::atomic<bool> disablePrespecializedMetadata = false;
@@ -171,12 +176,20 @@ struct LibPrespecializedState {
   struct AddressRange {
     uintptr_t start, end;
 
-    bool contains(const void *ptr) {
+    bool contains(const void *ptr) const {
       return start <= (uintptr_t)ptr && (uintptr_t)ptr < end;
     }
   };
 
+  enum class MapConfiguration {
+    UseNameKeyedMap,
+    UsePointerKeyedMap,
+    UsePointerKeyedMapDebugMode,
+    Disabled,
+  };
+
   const LibPrespecializedData<InProcess> *data;
+  MapConfiguration mapConfiguration;
   AddressRange sharedCacheRange{0, 0};
   AddressRange metadataAllocatorInitialPoolRange{0, 0};
 
@@ -197,6 +210,60 @@ struct LibPrespecializedState {
     metadataAllocatorInitialPoolRange.end =
         metadataAllocatorInitialPoolRange.start + initialPoolLength;
 #endif
+
+    // Must do this after the shared cache range has been retrieved.
+    mapConfiguration = computeMapConfiguration(data);
+  }
+
+  MapConfiguration
+  computeMapConfiguration(const LibPrespecializedData<InProcess> *data) {
+    // If no data, we have to disable.
+    if (!data)
+      return MapConfiguration::Disabled;
+
+    auto nameKeyedMap = data->getMetadataMap();
+    auto pointerKeyedMap = data->getPointerKeyedMetadataMap();
+
+    // If we don't have either map, then disable it completely.
+    if (!nameKeyedMap && !pointerKeyedMap) {
+      LOG("No prespecializations map available from data at %p, disabling.",
+          data);
+      return MapConfiguration::Disabled;
+    }
+
+    // If we don't have the pointer-keyed map, fall back to the name-keyed map.
+    if (!pointerKeyedMap) {
+      LOG("Data at %p only contains name-keyed map.", data);
+      return MapConfiguration::UseNameKeyedMap;
+    }
+
+    // If we don't have the name-keyed map, always use the pointer-keyed map.
+    if (!nameKeyedMap) {
+      LOG("Data at %p only contains pointer-keyed map.", data);
+      return MapConfiguration::UsePointerKeyedMap;
+    }
+
+    // We have both. Consult the option flag.
+    bool usePointerKeyedMap =
+        data->getOptionFlags() &
+        LibPrespecializedData<InProcess>::OptionFlagDefaultToPointerKeyedMap;
+
+#if HAS_OS_FEATURE
+    if (os_feature_enabled_simple(Swift, useAlternatePrespecializationMap,
+                                  false))
+      usePointerKeyedMap = !usePointerKeyedMap;
+#endif
+
+    LOG("Data at %p contains both maps. Using %s keyed map.", data,
+        usePointerKeyedMap ? "pointer" : "name");
+    if (usePointerKeyedMap) {
+      // If we're using a map outside the shared cache, then we're in debug mode
+      // and need to use our own slow lookup.
+      if (!sharedCacheRange.contains(pointerKeyedMap))
+        return MapConfiguration::UsePointerKeyedMapDebugMode;
+      return MapConfiguration::UsePointerKeyedMap;
+    }
+    return MapConfiguration::UseNameKeyedMap;
   }
 };
 
@@ -218,7 +285,7 @@ hasNonTypeGenericArguments(const TargetGenericContext<InProcess> *generics) {
 }
 
 static bool
-isPotentialPrespecializedPointer(LibPrespecializedState &prespecialized,
+isPotentialPrespecializedPointer(const LibPrespecializedState &state,
                                  const void *pointer) {
   // Prespecialized metadata descriptors and arguments are always in the shared
   // cache. They're either statically emitted metadata, or they're
@@ -228,16 +295,16 @@ isPotentialPrespecializedPointer(LibPrespecializedState &prespecialized,
   // 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))
+  if (!state.sharedCacheRange.contains(state.data))
     return true;
 
   // Anything outside the shared cache isn't a potential argument.
-  if (!prespecialized.sharedCacheRange.contains(pointer))
+  if (!state.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))
+  if (state.metadataAllocatorInitialPoolRange.contains(pointer))
     return false;
 
   return true;
@@ -256,20 +323,10 @@ swift::libPrespecializedImageLoaded() {
   #endif
 }
 
-Metadata *
-swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
-                                    const void *const *arguments) {
-  if (SWIFT_UNLIKELY(
-        disableForValidation
-        || disablePrespecializedMetadata.load(std::memory_order_acquire)))
-    return nullptr;
-
-  auto &prespecialized = LibPrespecialized.get();
-
-  auto *data = prespecialized.data;
-  if (!data)
-    return nullptr;
-
+static Metadata *
+getMetadataFromNameKeyedMap(const LibPrespecializedState &state,
+                            const TypeContextDescriptor *description,
+                            const void *const *arguments) {
   auto *generics = description->getGenericContext();
   if (!generics)
     return nullptr;
@@ -279,15 +336,15 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
   if (hasNonTypeGenericArguments(generics))
     return nullptr;
 
-  if (!isPotentialPrespecializedPointer(prespecialized, description)) {
+  if (!isPotentialPrespecializedPointer(state, 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])) {
+    if (!isPotentialPrespecializedPointer(state, arguments[i])) {
       LOG("Rejecting argument %u %p to descriptor %p, not in the shared cache",
           i, arguments[i], (const void *)description);
       return nullptr;
@@ -322,13 +379,120 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
   }
 
   auto key = mangling.result();
-  auto *metadataMap = data->getMetadataMap();
+  auto *metadataMap = state.data->getMetadataMap();
   auto *element = metadataMap->find(key.data(), key.size());
   auto *result = element ? element->value : nullptr;
   LOG("found %p for key '%.*s'.", result, (int)key.size(), key.data());
   return result;
 }
 
+static Metadata *
+getMetadataFromPointerKeyedMap(const LibPrespecializedState &state,
+                               const TypeContextDescriptor *description,
+                               const void *const *arguments) {
+#if DYLD_FIND_POINTER_HASH_TABLE_ENTRY_DEFINED
+  auto *generics = description->getGenericContext();
+  if (!generics)
+    return nullptr;
+
+  auto argumentCount = generics->getGenericContextHeader().NumKeyArguments;
+
+  auto *map = state.data->getPointerKeyedMetadataMap();
+  auto result = _dyld_find_pointer_hash_table_entry(
+      map, description, argumentCount, const_cast<const void **>(arguments));
+  LOG("Looking up description %p in dyld table, found %p.", description,
+      result);
+  return reinterpret_cast<Metadata *>(const_cast<void *>(result));
+#else
+  LOG("Looking up description %p but dyld hash table call not available.",
+      description);
+  return nullptr;
+#endif
+}
+
+// When we have a pointer-keyed map from a debug library, it's not built as a
+// hash table. We just scan it linearly.
+static Metadata *getMetadataFromPointerKeyedMapDebugMode(
+    const LibPrespecializedState &state,
+    const TypeContextDescriptor *description, const void *const *arguments) {
+  auto *generics = description->getGenericContext();
+  if (!generics)
+    return nullptr;
+
+  auto argumentCount = generics->getGenericContextHeader().NumKeyArguments;
+  auto *mapPtr = state.data->getPointerKeyedMetadataMap();
+
+  struct MapKey {
+    size_t count;
+    void *pointers[];
+  };
+
+  struct MapEntry {
+    const MapKey *key;
+    Metadata *value;
+  };
+
+  struct Map {
+    size_t count;
+    MapEntry entries[];
+  };
+
+  const Map *map = reinterpret_cast<const Map *>(mapPtr);
+  for (size_t i = 0; i < map->count; i++) {
+    auto &entry = map->entries[i];
+
+    // Keys are descriptor followed by arguments, so their count is 1 plus the
+    // argument count.
+    if (entry.key->count != argumentCount + 1)
+      continue;
+
+    // Check the descriptor.
+    if (description != entry.key->pointers[0])
+      continue;
+
+    // Check the rest. The pointers array is now offset by 1 since index 0 is
+    // the descriptor.
+    bool equal = true;
+    for (size_t j = 0; j < argumentCount; j++) {
+      if (entry.key->pointers[j + 1] != arguments[j]) {
+        equal = false;
+        break;
+      }
+    }
+
+    if (equal) {
+      LOG("Looking up description %p in debug table, found %p.", description,
+          entry.value);
+      return entry.value;
+    }
+  }
+
+  LOG("Looking up description %p in debug table, no entry found.", description);
+  return nullptr;
+}
+
+Metadata *
+swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
+                                    const void *const *arguments) {
+  if (SWIFT_UNLIKELY(disableForValidation || disablePrespecializedMetadata.load(
+                                                 std::memory_order_acquire)))
+    return nullptr;
+
+  auto &state = LibPrespecialized.get();
+
+  switch (state.mapConfiguration) {
+  case LibPrespecializedState::MapConfiguration::Disabled:
+    return nullptr;
+  case LibPrespecializedState::MapConfiguration::UseNameKeyedMap:
+    return getMetadataFromNameKeyedMap(state, description, arguments);
+  case LibPrespecializedState::MapConfiguration::UsePointerKeyedMap:
+    return getMetadataFromPointerKeyedMap(state, description, arguments);
+  case LibPrespecializedState::MapConfiguration::UsePointerKeyedMapDebugMode:
+    return getMetadataFromPointerKeyedMapDebugMode(state, description,
+                                                   arguments);
+  }
+}
+
 void _swift_validatePrespecializedMetadata() {
   auto *data = getLibPrespecializedData();
   if (!data) {