[6.2] Improve memoization of conformance lookups during casting (#82746)

Author: Catfish-Man

benchmark/single-source/ObjectiveCBridging.swift

@@ -97,6 +97,9 @@ public let benchmarks = [
   BenchmarkInfo(name: "NSArray.bridged.repeatedBufferAccess",
                   runFunction: run_BridgedNSArrayRepeatedBufferAccess, tags: t,
                   setUpFunction: setup_bridgedArrays),
+  BenchmarkInfo(name: "NSDictionary.bridged.enumerate",
+                  runFunction: run_BridgedNSDictionaryEnumerate, tags: t,
+                  setUpFunction: setup_bridgedDictionaries),
   BenchmarkInfo(name: "NSString.bridged.byteCount.ascii.ascii",
                   runFunction: run_BridgedNSStringLengthASCII_ASCII, tags: ts,
                   setUpFunction: setup_bridgedStrings),
@@ -812,6 +815,7 @@ public func run_UnicodeStringFromCodable(_ n: Int) {
 
 #if _runtime(_ObjC)
 var bridgedArray:NSArray! = nil
+var bridgedDictionaryOfNumbersToNumbers:NSDictionary! = nil
 var bridgedArrayMutableCopy:NSMutableArray! = nil
 var nsArray:NSArray! = nil
 var nsArrayMutableCopy:NSMutableArray! = nil
@@ -824,11 +828,20 @@ public func setup_bridgedArrays() {
   var arr = Array(repeating: NSObject(), count: 100) as [AnyObject]
   bridgedArray = arr as NSArray
   bridgedArrayMutableCopy = (bridgedArray.mutableCopy() as! NSMutableArray)
+
   nsArray = NSArray(objects: &arr, count: 100)
   nsArrayMutableCopy = (nsArray.mutableCopy() as! NSMutableArray)
   #endif
 }
 
+public func setup_bridgedDictionaries() {
+  var numDict = Dictionary<Int, Int>()
+  for i in 0 ..< 100 {
+    numDict[i] = i
+  }
+  bridgedDictionaryOfNumbersToNumbers = numDict as NSDictionary
+}
+
 public func setup_bridgedStrings() {
   #if _runtime(_ObjC)
   let str = Array(repeating: "The quick brown fox jumps over the lazy dog.", count: 100).joined()
@@ -849,6 +862,23 @@ public func run_BridgedNSArrayObjectAtIndex(_ n: Int) {
   #endif
 }
 
+private func dictionaryApplier(
+  _ keyPtr: UnsafeRawPointer?,
+  _ valuePtr :UnsafeRawPointer?,
+  _ contextPtr: UnsafeMutableRawPointer?
+) -> Void {}
+
+@inline(never)
+public func run_BridgedNSDictionaryEnumerate(_ n: Int) {
+  #if _runtime(_ObjC)
+  let cf = bridgedDictionaryOfNumbersToNumbers as CFDictionary
+  for _ in 0 ..< n * 50 {
+    // Use CF to prevent Swift from providing an override, forcing going through ObjC bridging
+    CFDictionaryApplyFunction(cf, dictionaryApplier, nil)
+  }
+  #endif
+}
+
 @inline(never)
 public func run_BridgedNSArrayBufferAccess(_ n: Int) {
   #if _runtime(_ObjC)

stdlib/public/runtime/Casting.cpp

@@ -1436,10 +1436,31 @@ extern "C" const _ObjectiveCBridgeableWitnessTable BRIDGING_CONFORMANCE_SYM;
 /// Nominal type descriptor for Swift.String.
 extern "C" const StructDescriptor NOMINAL_TYPE_DESCR_SYM(SS);
 
+struct ObjCBridgeWitnessCacheEntry {
+  const Metadata *metadata;
+  const _ObjectiveCBridgeableWitnessTable *witness;
+};
+
+// String is so important that we cache it permanently, so we don't want to
+// pollute this temporary cache with the String entry
+static const _ObjectiveCBridgeableWitnessTable *
+swift_conformsToObjectiveCBridgeableNoCache(const Metadata *T) {
+  auto w = swift_conformsToProtocolCommon(
+         T, &PROTOCOL_DESCR_SYM(s21_ObjectiveCBridgeable));
+  return reinterpret_cast<const _ObjectiveCBridgeableWitnessTable *>(w);
+}
+
 static const _ObjectiveCBridgeableWitnessTable *
 swift_conformsToObjectiveCBridgeable(const Metadata *T) {
-  return reinterpret_cast<const _ObjectiveCBridgeableWitnessTable *>
-    (swift_conformsToProtocolCommon(T, &PROTOCOL_DESCR_SYM(s21_ObjectiveCBridgeable)));
+  static std::atomic<ObjCBridgeWitnessCacheEntry> _objcBridgeWitnessCache = {};
+  auto cached = _objcBridgeWitnessCache.load(SWIFT_MEMORY_ORDER_CONSUME);
+  if (cached.metadata == T) {
+    return cached.witness;
+  }
+  cached.witness = swift_conformsToObjectiveCBridgeableNoCache(T);
+  cached.metadata = T;
+  _objcBridgeWitnessCache.store(cached, std::memory_order_release);
+  return cached.witness;
 }
 
 static const _ObjectiveCBridgeableWitnessTable *
@@ -1451,7 +1472,7 @@ findBridgeWitness(const Metadata *T) {
   if (T->getKind() == MetadataKind::Struct) {
     auto structDescription = cast<StructMetadata>(T)->Description;
     if (structDescription == &NOMINAL_TYPE_DESCR_SYM(SS)) {
-      static auto *Swift_String_ObjectiveCBridgeable = swift_conformsToObjectiveCBridgeable(T);
+      static auto *Swift_String_ObjectiveCBridgeable = swift_conformsToObjectiveCBridgeableNoCache(T);
       return Swift_String_ObjectiveCBridgeable;
     }
   }

stdlib/public/runtime/DynamicCast.cpp

@@ -206,13 +206,86 @@ struct _ObjectiveCBridgeableWitnessTable : WitnessTable {
 extern "C" const ProtocolDescriptor
 PROTOCOL_DESCR_SYM(s21_ObjectiveCBridgeable);
 
+#if SWIFT_OBJC_INTEROP
+#define BRIDGING_CONFORMANCE_SYM \
+  MANGLE_SYM(s19_BridgeableMetatypeVs21_ObjectiveCBridgeablesWP)
+
+extern "C" const _ObjectiveCBridgeableWitnessTable BRIDGING_CONFORMANCE_SYM;
+#endif
+
+/// Nominal type descriptor for Swift.String.
+extern "C" const StructDescriptor NOMINAL_TYPE_DESCR_SYM(SS);
+
+struct ObjCBridgeWitnessCacheEntry {
+  const Metadata *metadata;
+  const _ObjectiveCBridgeableWitnessTable *witness;
+};
+
 static const _ObjectiveCBridgeableWitnessTable *
-findBridgeWitness(const Metadata *T) {
+swift_conformsToObjectiveCBridgeableNoCache(const Metadata *T) {
   auto w = swift_conformsToProtocolCommon(
-      T, &PROTOCOL_DESCR_SYM(s21_ObjectiveCBridgeable));
+         T, &PROTOCOL_DESCR_SYM(s21_ObjectiveCBridgeable));
   return reinterpret_cast<const _ObjectiveCBridgeableWitnessTable *>(w);
 }
 
+static const _ObjectiveCBridgeableWitnessTable *
+swift_conformsToObjectiveCBridgeable(const Metadata *T) {
+  static std::atomic<ObjCBridgeWitnessCacheEntry> _objcBridgeWitnessCache = {};
+  auto cached = _objcBridgeWitnessCache.load(SWIFT_MEMORY_ORDER_CONSUME);
+  if (cached.metadata == T) {
+    return cached.witness;
+  }
+  cached.witness = swift_conformsToObjectiveCBridgeableNoCache(T);
+  cached.metadata = T;
+  _objcBridgeWitnessCache.store(cached, std::memory_order_release);
+  return cached.witness;
+}
+
+static const _ObjectiveCBridgeableWitnessTable *
+findBridgeWitness(const Metadata *T) {
+  // Special case: Memoize the bridge witness for Swift.String.
+  // Swift.String is the most heavily used bridge because of the prevalence of
+  // string-keyed dictionaries in Obj-C.  It's worth burning a few words of static
+  // storage to avoid repeatedly looking up this conformance.
+  if (T->getKind() == MetadataKind::Struct) {
+    auto structDescription = cast<StructMetadata>(T)->Description;
+    if (structDescription == &NOMINAL_TYPE_DESCR_SYM(SS)) {
+      static auto *Swift_String_ObjectiveCBridgeable = swift_conformsToObjectiveCBridgeableNoCache(T);
+      return Swift_String_ObjectiveCBridgeable;
+    }
+  }
+
+  auto w = swift_conformsToObjectiveCBridgeable(T);
+  if (SWIFT_LIKELY(w))
+    return reinterpret_cast<const _ObjectiveCBridgeableWitnessTable *>(w);
+  // Class and ObjC existential metatypes can be bridged, but metatypes can't
+  // directly conform to protocols yet. Use a stand-in conformance for a type
+  // that looks like a metatype value if the metatype can be bridged.
+  switch (T->getKind()) {
+  case MetadataKind::Metatype: {
+#if SWIFT_OBJC_INTEROP
+    auto metaTy = static_cast<const MetatypeMetadata *>(T);
+    if (metaTy->InstanceType->isAnyClass())
+      return &BRIDGING_CONFORMANCE_SYM;
+#endif
+    break;
+  }
+  case MetadataKind::ExistentialMetatype: {
+#if SWIFT_OBJC_INTEROP
+    auto existentialMetaTy =
+      static_cast<const ExistentialMetatypeMetadata *>(T);
+    if (existentialMetaTy->isObjC())
+      return &BRIDGING_CONFORMANCE_SYM;
+#endif
+    break;
+  }
+
+  default:
+    break;
+  }
+  return nullptr;
+}
+
 /// Retrieve the bridged Objective-C type for the given type that
 /// conforms to \c _ObjectiveCBridgeable.
 MetadataResponse
@@ -734,7 +807,7 @@ struct ObjCBridgeMemo {
 #if !NDEBUG
                    memo->destType = setupData->destType;
 #endif
-                   memo->destBridgeWitness = findBridgeWitness(setupData->destType);
+                   memo->destBridgeWitness = swift_conformsToObjectiveCBridgeableNoCache(setupData->destType);
                    if (memo->destBridgeWitness == nullptr) {
                      memo->targetBridgedType = nullptr;
                      memo->targetBridgedObjCClass = nullptr;
@@ -777,6 +850,17 @@ struct ObjCBridgeMemo {
         destBridgeWitness, targetBridgedType);
     }
 };
+
+static const HashableWitnessTable* tryMemoizeNSStringHashableConformance(const Metadata *cls) {
+  auto nsString = getNSStringMetadata();
+  do {
+    if (cls == nsString) {
+      return getNSStringHashableConformance();
+    }
+    cls = _swift_class_getSuperclass(cls);
+  } while (cls != nullptr);
+  return nullptr;
+}
 #endif
 
 static DynamicCastResult
@@ -794,23 +878,16 @@ tryCastToAnyHashable(
   const HashableWitnessTable *hashableConformance = nullptr;
 
   switch (srcType->getKind()) {
+  case MetadataKind::Existential: {
+    return DynamicCastResult::Failure;
+  }
   case MetadataKind::ForeignClass: // CF -> String
   case MetadataKind::ObjCClassWrapper: { // Obj-C -> String
 #if SWIFT_OBJC_INTEROP
-    auto cls = srcType;
-    auto nsString = getNSStringMetadata();
-    do {
-      if (cls == nsString) {
-        hashableConformance = getNSStringHashableConformance();
-        break;
-      }
-      cls = _swift_class_getSuperclass(cls);
-    } while (cls != nullptr);
-    break;
-#else
+    hashableConformance = tryMemoizeNSStringHashableConformance(srcType);
+#endif
     // If no Obj-C interop, just fall through to the general case.
     break;
-#endif
   }
   case MetadataKind::Optional: {
     // FIXME: https://github.com/apple/swift/issues/51550