Implement an SPI to iterate types who conform to some protocol

Author: Azoy

stdlib/public/Reflection/Sources/Reflection/CMakeLists.txt

@@ -17,6 +17,7 @@ list(APPEND SWIFT_REFLECTION_SWIFT_FLAGS
 
 add_swift_target_library(swiftReflection ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} IS_STDLIB
   Case.swift
+  Conformances.swift
   Field.swift
   GenericArguments.swift
   KeyPath.swift

stdlib/public/Reflection/Sources/Reflection/Conformances.swift

@@ -0,0 +1,66 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+//
+//===----------------------------------------------------------------------===//
+
+#if os(macOS) || os(iOS) || os(tvOS) || os(watchOS)
+
+import Swift
+import _Runtime
+
+@available(SwiftStdlib 5.9, *)
+@_silgen_name("_swift_reflection_withConformanceCache")
+func _withConformanceCache(
+  _ proto: ProtocolDescriptor,
+  _ context: UnsafeMutableRawPointer,
+  _ callback: @convention(c) (
+    /* Array of types */ UnsafePointer<UnsafeRawPointer>,
+    /* Number of types */ Int,
+    /* Context we just passed */ UnsafeMutableRawPointer
+  ) -> ()
+)
+
+@available(SwiftStdlib 5.9, *)
+@_spi(Reflection)
+public func _typesThatConform(to type: Any.Type) -> [Any.Type]? {
+  let meta = Metadata(type)
+
+  guard meta.kind == .existential else {
+    return nil
+  }
+
+  let existential = meta.existential
+
+  let protos = existential.protocols
+
+  guard protos.count == 1 else {
+    return nil
+  }
+
+  let proto = protos[0]
+
+  var result: [Any.Type] = []
+
+  withUnsafeMutablePointer(to: &result) {
+    _withConformanceCache(proto, UnsafeMutableRawPointer($0)) {
+      let buffer = UnsafeBufferPointer<Any.Type>(
+        start: UnsafePointer<Any.Type>($0._rawValue),
+        count: $1
+      )
+
+      let arrayPtr = $2.assumingMemoryBound(to: [Any.Type].self)
+
+      arrayPtr.pointee = Array(buffer)
+    }
+  }
+
+  return result
+}
+
+#endif

stdlib/public/Reflection/Sources/_Runtime/CMakeLists.txt

@@ -64,6 +64,7 @@ add_swift_target_library(swift_Runtime ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} IS_ST
   Utils/RelativePointer.swift
   Utils/TypeCache.swift
 
+  Caches.cpp
   ConformanceDescriptor.swift
   ExistentialContainer.swift
   Functions.swift

stdlib/public/Reflection/Sources/_Runtime/Caches.cpp

@@ -0,0 +1,171 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+//
+//===----------------------------------------------------------------------===//
+
+// Only as SPI for Darwin platforms for right now...
+#if defined(__MACH__)
+
+#include "swift/ABI/Metadata.h"
+#include "swift/Basic/Lazy.h"
+#include "swift/Runtime/Concurrent.h"
+#include "swift/Runtime/Config.h"
+#include "swift/Runtime/ImageInspection.h"
+#include "swift/Runtime/ImageInspectionMachO.h"
+#include <cstdint>
+#include <unordered_map>
+
+using namespace swift;
+
+//===----------------------------------------------------------------------===//
+// Protocol Conformance Cache
+//===----------------------------------------------------------------------===//
+
+namespace {
+  struct ConformanceSection {
+    const ProtocolConformanceRecord *Begin, *End;
+
+    ConformanceSection(const ProtocolConformanceRecord *begin,
+                       const ProtocolConformanceRecord *end)
+        : Begin(begin), End(end) {}
+
+    ConformanceSection(const void *ptr, uintptr_t size) {
+      auto bytes = reinterpret_cast<const char *>(ptr);
+      Begin = reinterpret_cast<const ProtocolConformanceRecord *>(ptr);
+      End = reinterpret_cast<const ProtocolConformanceRecord *>(bytes + size);
+    }
+
+    const ProtocolConformanceRecord *begin() const {
+      return Begin;
+    }
+    const ProtocolConformanceRecord *end() const {
+      return End;
+    }
+  };
+}
+
+struct ConformanceCache {
+  // All accesses to Cache and LastSectionCount must be within CacheMutex's
+  // lock scope.
+
+  std::unordered_map<
+    /* Key */   const ProtocolDescriptor *,
+    /* Value */ std::vector<const Metadata *>
+  > Cache;
+  Mutex CacheMutex;
+  ConcurrentReadableArray<ConformanceSection> Sections;
+
+  size_t LastSectionCount = 0;
+
+  ConformanceCache() {
+    initializeProtocolConformanceLookup();
+  }
+};
+
+static Lazy<ConformanceCache> Conformances;
+
+void swift::addImageProtocolConformanceBlockCallbackUnsafe(
+    const void *baseAddress,
+    const void *conformances,
+    uintptr_t conformancesSize) {
+  assert(conformancesSize % sizeof(ProtocolConformanceRecord) == 0 &&
+         "conformances section not a multiple of ProtocolConformanceRecord");
+
+  // Conformance cache should always be sufficiently initialized by this point.
+  auto &C = Conformances.unsafeGetAlreadyInitialized();
+
+  C.Sections.push_back(ConformanceSection{conformances, conformancesSize});
+}
+
+// WARNING: the callbacks are called from unsafe contexts (with the dyld and
+// ObjC runtime locks held) and must be very careful in what they do. Locking
+// must be arranged to avoid deadlocks (other code must never call out to dyld
+// or ObjC holding a lock that gets taken in one of these callbacks) and the
+// new/delete operators must not be called, in case a program supplies an
+// overload which does not cooperate with these requirements.
+
+void swift::initializeProtocolConformanceLookup() {
+  REGISTER_FUNC(
+      addImageCallback<TextSegment, ProtocolConformancesSection,
+                       addImageProtocolConformanceBlockCallbackUnsafe>);
+}
+
+static const Metadata *_getCanonicalTypeMetadata(
+                             const ProtocolConformanceDescriptor *conformance) {
+  switch (conformance->getTypeKind()) {
+  case TypeReferenceKind::DirectTypeDescriptor:
+  case TypeReferenceKind::IndirectTypeDescriptor: {
+    if (auto anyType = conformance->getTypeDescriptor()) {
+      if (auto type = dyn_cast<TypeContextDescriptor>(anyType)) {
+        if (!type->isGeneric()) {
+          if (auto accessFn = type->getAccessFunction()) {
+            return accessFn(MetadataState::Abstract).Value;
+          }
+        }
+      }
+    }
+
+    return nullptr;
+  }
+
+  case TypeReferenceKind::DirectObjCClassName:
+  case TypeReferenceKind::IndirectObjCClass:
+    return nullptr;
+  }
+}
+
+using ConformanceCacheCallback = void (*)(const Metadata **,
+                                          size_t, void *);
+
+SWIFT_RUNTIME_STDLIB_SPI SWIFT_CC(swift)
+void _swift_reflection_withConformanceCache(const ProtocolDescriptor *proto,
+                                            void *context,
+                                            ConformanceCacheCallback callback) {
+  auto &C = Conformances.get();
+
+  auto snapshot = C.Sections.snapshot();
+
+  Mutex::ScopedLock lock(C.CacheMutex);
+
+  if (C.LastSectionCount > 0 && snapshot.count() <= C.LastSectionCount) {
+    auto entry = C.Cache.find(proto);
+
+    if (entry != C.Cache.end()) {
+      callback(entry->second.data(), entry->second.size(), context);
+      return;
+    }
+  }
+
+  std::vector<const Metadata *> types = {};
+
+  for (auto &section : snapshot) {
+    for (auto &record : section) {
+      auto conformance = record.get();
+
+      if (conformance->getProtocol() != proto) {
+        continue;
+      }
+
+      if (conformance->hasConditionalRequirements()) {
+        continue;
+      }
+
+      if (auto type = _getCanonicalTypeMetadata(conformance)) {
+        types.push_back(type);
+      }
+    }
+  }
+
+  callback(types.data(), types.size(), context);
+
+  C.Cache[proto] = types;
+  C.LastSectionCount = snapshot.count();
+}
+
+#endif

stdlib/public/Reflection/Sources/_Runtime/Metadata/Metadata.swift

@@ -53,7 +53,12 @@ extension Metadata {
   public var `enum`: EnumMetadata {
     EnumMetadata(ptr)
   }
-  
+
+  @inlinable
+  public var existential: ExistentialMetadata {
+    ExistentialMetadata(ptr)
+  }
+
   @inlinable
   public var extendedExistential: ExtendedExistentialMetadata {
     ExtendedExistentialMetadata(ptr)

test/stdlib/Reflection/Conformances.swift

@@ -0,0 +1,39 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+// REQUIRES: reflection_runtime
+// REQUIRES: OS=macosx || OS=ios || OS=tvos || OS=watchos
+// UNSUPPORTED: freestanding
+
+import StdlibUnittest
+
+@_spi(Reflection)
+import Reflection
+
+let suite = TestSuite("Conformances")
+
+protocol MyCustomProtocol {}
+
+struct Conformer: MyCustomProtocol {}
+struct NotConformer {}
+struct GenericConformer<T>: MyCustomProtocol {}
+struct GenericConditionalConformer<T> {}
+
+extension GenericConditionalConformer: MyCustomProtocol where T: MyCustomProtocol {}
+
+if #available(SwiftStdlib 5.9, *) {
+  suite.test("Basic") {
+    let customTypes = _typesThatConform(to: (any MyCustomProtocol).self)
+
+    expectNotNil(customTypes)
+    expectEqual(customTypes!.count, 1)
+    expectTrue(customTypes![0] == Conformer.self)
+
+    let kpType = _typesThatConform(to: (any _AppendKeyPath).self)
+
+    expectNotNil(kpType)
+    expectEqual(kpType!.count, 1)
+    expectTrue(kpType![0] == AnyKeyPath.self)
+  }
+}
+
+runAllTests()