[Runtime] Do a proper hash table lookup in the prespecializations library.

We were doing a linear scan of the table contents as a stopgap. Stop doing that, and compute the proper key for the lookup, matching the one used in the builder.

Author: mikeash

include/swift/Runtime/PrebuiltStringMap.h

@@ -57,7 +57,7 @@ struct PrebuiltStringMap {
   PrebuiltStringMap(uint64_t arraySize) : arraySize(arraySize) {}
 
   // Based on MurmurHash2
-  uint64_t hash(const void *data, size_t len) {
+  uint64_t hash(const void *data, size_t len) const {
     uint64_t magic = 0xc6a4a7935bd1e995ULL;
     uint64_t salt = 47;
 
@@ -148,8 +148,12 @@ struct PrebuiltStringMap {
 
   /// Look up the given string in the table. Requires that StringTy be
   /// `const char *`.
-  ArrayElement *find(const char *toFind) {
+  const ArrayElement *find(const char *toFind) const {
     size_t len = strlen(toFind);
+    return find(toFind, len);
+  }
+
+  const ArrayElement *find(const char *toFind, size_t len) const {
     uint64_t hashValue = hash(toFind, len);
 
     size_t index = hashValue % arraySize;

stdlib/public/runtime/Demangle.cpp

@@ -266,8 +266,8 @@ _buildDemanglingForMetadataPack(MetadataPackPointer pack, size_t count,
   return node;
 }
 
-/// Build an array of demangling trees for each generic argument of the given
-/// type metadata.
+/// Build an array of demangling trees for each generic argument to the given
+/// generic type context descriptor.
 ///
 /// Note:
 /// - The input array has an entry for those generic parameter descriptors which
@@ -286,15 +286,14 @@ _buildDemanglingForMetadataPack(MetadataPackPointer pack, size_t count,
 /// The output array is flat; the generic parameters of each depth are
 /// concatenated together.
 static bool _buildDemanglingForGenericArgs(
-    const Metadata *type,
     const TypeContextDescriptor *description,
+    const Metadata *const *genericArgs,
     llvm::SmallVectorImpl<NodePointer> &demangledGenerics,
     Demangle::Demangler &Dem) {
   auto generics = description->getGenericContext();
   if (!generics)
     return true;
 
-  auto genericArgs = description->getGenericArguments(type);
   auto packHeader = generics->getGenericPackShapeHeader();
   auto packDescriptors = generics->getGenericPackShapeDescriptors();
 
@@ -420,12 +419,16 @@ _buildDemanglingForNominalType(const Metadata *type, Demangle::Demangler &Dem) {
     return nullptr;
   }
 
-  // Gather the complete set of generic arguments that must be written to
-  // form this type.
   llvm::SmallVector<NodePointer, 8> demangledGenerics;
-  if (!_buildDemanglingForGenericArgs(type, description, demangledGenerics, Dem))
-    return nullptr;
-  
+  if (description->isGeneric()) {
+    auto genericArgs = description->getGenericArguments(type);
+    // Gather the complete set of generic arguments that must be written to
+    // form this type.
+    if (!_buildDemanglingForGenericArgs(description, genericArgs,
+                                        demangledGenerics, Dem))
+      return nullptr;
+  }
+
   return _buildDemanglingForContext(description, demangledGenerics, Dem);
 }
 
@@ -819,6 +822,25 @@ swift::_swift_buildDemanglingForMetadata(const Metadata *type,
   return nullptr;
 }
 
+Demangle::NodePointer
+swift::_buildDemanglingForGenericType(const TypeContextDescriptor *description,
+                                      const void *const *arguments,
+                                      Demangle::Demangler &Dem) {
+  auto kind = description->getKind();
+  if (kind != ContextDescriptorKind::Class &&
+      kind != ContextDescriptorKind::Enum &&
+      kind != ContextDescriptorKind::Struct)
+    return nullptr;
+
+  llvm::SmallVector<NodePointer, 8> demangledGenerics;
+  if (!_buildDemanglingForGenericArgs(description,
+                                      (const Metadata *const *)arguments,
+                                      demangledGenerics, Dem))
+    return nullptr;
+
+  return _buildDemanglingForContext(description, demangledGenerics, Dem);
+}
+
 // NB: This function is not used directly in the Swift codebase, but is
 // exported for Xcode support and is used by the sanitizers. Please coordinate
 // before changing.

stdlib/public/runtime/EnvironmentVariables.def

@@ -94,6 +94,9 @@ VARIABLE(SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED, bool, true,
 VARIABLE(SWIFT_DEBUG_LIB_PRESPECIALIZED_PATH, string, "",
          "A path to a prespecializations library to use at runtime.")
 
+VARIABLE(SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING, bool, false,
+         "Enable debug logging of prespecializations library use.")
+
 VARIABLE(SWIFT_ROOT, string, "",
          "Overrides the root directory of the Swift installation. "
          "This is used to locate auxiliary files relative to the runtime "

stdlib/public/runtime/LibPrespecialized.cpp

@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "swift/Runtime/LibPrespecialized.h"
+#include "Private.h"
 #include "swift/Basic/Lazy.h"
 #include "swift/Runtime/EnvironmentVariables.h"
 #include "swift/Runtime/Metadata.h"
@@ -74,36 +75,35 @@ swift::getLibPrespecializedMetadata(const TypeContextDescriptor *description,
   if (!data)
     return nullptr;
 
-  auto *metadataMap = data->getMetadataMap();
-
-  // TODO: This linear search is a stopgap until we work out the right
-  // representation for the string keys.
-  size_t count = metadataMap->arraySize;
-  auto *entries = metadataMap->array();
-  for (size_t i = 0; i < count; i++) {
-    if (entries[i].key == nullptr)
-      continue;
-
-    Metadata *candidate = entries[i].value;
-    auto *candidateDescription = candidate->getDescription();
-    if (description != candidateDescription)
-      continue;
-
-    auto *candidateArguments = candidate->getGenericArgs();
-
-    bool match = true;
-
-    const auto &genericContext = *description->getGenericContext();
-    const auto &header = genericContext.getGenericContextHeader();
-    for (unsigned param = 0; param < header.NumParams; param++)
-      if (arguments[param] != candidateArguments[param]) {
-        match = false;
-        break;
-      }
-
-    if (match)
-      return candidate;
+  Demangler dem;
+  auto mangleNode = _buildDemanglingForGenericType(description, arguments, dem);
+  if (mangleNode->getKind() != Node::Kind::Global) {
+    auto wrapper = dem.createNode(Node::Kind::Global);
+    wrapper->addChild(mangleNode, dem);
+    mangleNode = wrapper;
+  }
+  auto resolver = [](SymbolicReferenceKind kind,
+                     const void *ref) -> NodePointer {
+    swift::fatalError(0,
+                      "Unexpected symbolic reference %p in generated mangle "
+                      "tree for generic type lookup.",
+                      ref);
+  };
+  auto mangling = Demangle::mangleNode(mangleNode, resolver, dem);
+  if (!mangling.isSuccess()) {
+    swift::warning(0,
+                   "Mangling for prespecialized metadata failed with code %d",
+                   mangling.error().code);
+    return nullptr;
   }
 
-  return nullptr;
+  auto key = mangling.result();
+  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());
+  return result;
 }

stdlib/public/runtime/Private.h

@@ -280,6 +280,13 @@ class TypeInfo {
   Demangle::NodePointer _swift_buildDemanglingForMetadata(const Metadata *type,
                                                           Demangle::Demangler &Dem);
 
+  /// Build the demangling for the generic type that's created by specializing
+  /// the given type context descriptor with the given arguments.
+  Demangle::NodePointer
+  _buildDemanglingForGenericType(const TypeContextDescriptor *description,
+                                 const void *const *arguments,
+                                 Demangle::Demangler &Dem);
+
   /// Callback used to provide the substitution of a generic parameter
   /// (described by depth/index) to its metadata.
   ///