[Runtime] Support type descriptor map in LibPrespecialized.

The descriptor map is keyed by a simplified mangling that canonicalizes the differences that we accept in _contextDescriptorMatchesMangling, such as the ability to specify any kind of type with an OtherNominalType node.

This simplified mangling is not necessarily unique, but we use _contextDescriptorMatchesMangling for the final equality checking when looking up entries in the map, so occasional collisions are acceptable and get resolved when probing the table.

The table is meant to be comprehensive, so it includes all descriptors that can be looked up by name, and a negative result means the descriptor does not exist in the shared cache. We add a flag to the options that can mark it as non-definitive in case we ever need to degrade this, and fall back to a full search after a negative result.

The map encompasses the entire shared cache but we need to reject lookups for types in images that aren't loaded. The map includes an image index which allows us to cheaply query whether a given descriptor is in a loaded image or not, so we can ignore ones which are not.

TypeMetadataPrivateState now has a separate sections array for sections within the shared cache. _searchTypeMetadataRecords consults the map first. If no result is found in the map and the map is marked as comprehensive, then only the sections outside the shared cache need to be scanned.

Replace the SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED environment variable with one specifically for metadata and one for descriptor lookup so they can be controlled independently. Also add SWIFT_DEBUG_VALIDATE_LIB_PRESPECIALIZED_DESCRIPTOR_LOOKUP which consults the map and does the full scan, and ensures they produce the same result, for debugging purposes.

Remove the disablePrespecializedMetadata global and instead modify the mapConfiguration to disable prespecialized metadata when an image is loaded that overrides one in the shared cache.

rdar://113059233

Author: mikeash

include/swift/Runtime/LibPrespecialized.h

@@ -20,6 +20,7 @@
 #include "PrebuiltStringMap.h"
 #include "swift/ABI/Metadata.h"
 #include "swift/ABI/TargetLayout.h"
+#include "swift/Demangling/Demangler.h"
 
 #define LIB_PRESPECIALIZED_TOP_LEVEL_SYMBOL_NAME "_swift_prespecializationsData"
 
@@ -36,23 +37,40 @@ struct LibPrespecializedData {
 
   typename Runtime::StoredSize optionFlags;
 
+  TargetPointer<Runtime, const void> descriptorMap;
+
   // Existing fields are above, add new fields below this point.
 
+  // The major/minor version numbers for this version of the struct.
   static constexpr uint32_t currentMajorVersion = 1;
-  static constexpr uint32_t currentMinorVersion = 3;
+  static constexpr uint32_t currentMinorVersion = 4;
 
+  // Version numbers where various fields were introduced.
   static constexpr uint32_t minorVersionWithDisabledProcessesTable = 2;
   static constexpr uint32_t minorVersionWithPointerKeyedMetadataMap = 3;
   static constexpr uint32_t minorVersionWithOptionFlags = 3;
+  static constexpr uint32_t minorVersionWithDescriptorMap = 4;
 
   // 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,
+
+    // When this flag is set, descriptorMap is not comprehensive, meaning that
+    // a negative lookup result is not a definitive failure.
+    OptionFlagDescriptorMapNotComprehensive = 1ULL << 1,
   };
 
-  // Helpers for retrieving the metadata map in-process.
+  // Helpers for safely retrieving various fields. Helpers return 0 or NULL if
+  // the version number indicates that the field is not present.
+
+  typename Runtime::StoredSize getOptionFlags() const {
+    if (minorVersion < minorVersionWithOptionFlags)
+      return 0;
+    return optionFlags;
+  }
+
   static bool stringIsNull(const char *str) { return str == nullptr; }
 
   using MetadataMap = PrebuiltStringMap<const char *, Metadata *, stringIsNull>;
@@ -73,18 +91,136 @@ struct LibPrespecializedData {
     return pointerKeyedMetadataMap;
   }
 
-  typename Runtime::StoredSize getOptionFlags() const {
-    if (minorVersion < minorVersionWithOptionFlags)
-      return 0;
-    return optionFlags;
+  using DescriptorMap =
+      PrebuiltAuxDataImplicitStringMap<TargetPointer<Runtime, const void>,
+                                       uint16_t>;
+
+  const DescriptorMap *getDescriptorMap() const {
+    if (minorVersion < minorVersionWithDescriptorMap)
+      return nullptr;
+    return reinterpret_cast<const DescriptorMap *>(descriptorMap);
   }
 };
 
+enum class LibPrespecializedLookupResult {
+  // We found something.
+  Found,
+
+  // We didn't find anything, and we know it's not in the shared cache.
+  DefinitiveNotFound,
+
+  // We didn't find anything, but we couldn't rule out the shared cache. Caller
+  // must do a full search.
+  NonDefinitiveNotFound,
+};
+
 const LibPrespecializedData<InProcess> *getLibPrespecializedData();
+
 Metadata *getLibPrespecializedMetadata(const TypeContextDescriptor *description,
                                        const void *const *arguments);
 void libPrespecializedImageLoaded();
 
+std::pair<LibPrespecializedLookupResult, const TypeContextDescriptor *>
+getLibPrespecializedTypeDescriptor(Demangle::NodePointer node);
+
+/// Given the demangling referring to a particular descriptor, build the
+/// canonical simplified version of the demangling that's used for the keys in
+/// the descriptorMap. We copy across Extension and Module nodes. Type nodes are
+/// all normalized to be OtherNominalType to allow for the runtime allowing
+/// type kind mismatches on imported C types in certain cases. Other nodes are
+/// skipped.
+///
+/// The runtime always searches through duplicates in the table, and uses its
+/// own matching on all candidates, so the simplified demangling is allowed to
+/// be simplified to the point of having different descriptors sometimes produce
+/// the same demangling.
+static inline Demangle::NodePointer
+buildSimplifiedDescriptorDemangling(Demangle::NodePointer node,
+                                    Demangle::Demangler &dem) {
+  // The node that will be returned to the caller.
+  Demangle::NodePointer result = nullptr;
+
+  // The bottommost node in the result that we've generated. Additional nodes
+  // are added as children to this one.
+  Demangle::NodePointer resultBottom = nullptr;
+
+  // The current node that we're iterating over in the input node tree.
+  Demangle::NodePointer current = node;
+
+  using Kind = Demangle::Node::Kind;
+
+  // Helper to add a new node to the result. This sets `result` to the node if
+  // it hasn't already been set (indicating this is the topmost node), and adds
+  // the node as a child to `resultBottom` otherwise. `resultBottom` is updated
+  // to point to the new node.
+  auto addNode = [&](Demangle::NodePointer newNode) {
+    if (!result) {
+      result = newNode;
+    } else {
+      if (resultBottom->getKind() == Kind::Extension) {
+        resultBottom->addChild(newNode, dem);
+      } else {
+        // Shift the Identifier down, insert before it.
+        resultBottom->addChild(resultBottom->getFirstChild(), dem);
+        resultBottom->replaceChild(0, newNode);
+      }
+    }
+    resultBottom = newNode;
+  };
+
+  // Walk down the input node tree.
+  while (current) {
+    switch (current->getKind()) {
+    case Kind::Extension: {
+      // Extensions are copied across. The new extension node has the module
+      // from the original, and the second child will be added as we traverse
+      // the next node in the tree.
+      auto copy = dem.createNode(Kind::Extension);
+      auto module = current->getChild(0);
+      if (module == nullptr || module->getKind() != Kind::Module)
+        return nullptr;
+      copy->addChild(module, dem);
+      addNode(copy);
+      current = current->getChild(1);
+      break;
+    }
+    case Kind::Module: {
+      // Module contents are always in the form we want, so we can incorporate
+      // this node verbatim and terminate the walk.
+      addNode(current);
+      current = nullptr;
+      break;
+    }
+    case Kind::Protocol:
+    case Kind::OpaqueType:
+    case Kind::Class:
+    case Kind::Structure:
+    case Kind::Enum:
+    case Kind::TypeAlias:
+    case Kind::OtherNominalType: {
+      // Type nodes are copied across with the kind always set to
+      // OtherNominalType.
+      auto copy = dem.createNode(Kind::OtherNominalType);
+      auto identifier = current->getChild(1);
+      if (identifier == nullptr || identifier->getKind() != Kind::Identifier)
+        return nullptr;
+      copy->addChild(identifier, dem);
+      addNode(copy);
+      current = current->getChild(0);
+      break;
+    }
+
+    default:
+      // If we don't know about this node, continue the walk with its first
+      // child.
+      current = current->getFirstChild();
+      break;
+    }
+  }
+
+  return result;
+}
+
 } // namespace swift
 
 // Validate the prespecialized metadata map by building each entry dynamically