Add a runtime-internal function, _enumerateTypes(fromImageAt:conformingTo:_:), that will walk the protocol conformance tables looking for types that conform to a given protocol.

Author: grynspan

stdlib/public/core/CMakeLists.txt

@@ -173,6 +173,7 @@ set(SWIFTLIB_ESSENTIAL
   SwiftNativeNSArray.swift
   TemporaryAllocation.swift
   ThreadLocalStorage.swift
+  TypeDiscovery.swift
   UIntBuffer.swift
   UnavailableStringAPIs.swift
   UnicodeData.swift

stdlib/public/core/GroupInfo.json

@@ -155,7 +155,8 @@
     "Mirror.swift",
     "Mirrors.swift",
     "ReflectionMirror.swift",
-    "ObjectIdentifier.swift"
+    "ObjectIdentifier.swift",
+    "TypeDiscovery.swift"
   ],
   "Math": [
     "SetAlgebra.swift",

stdlib/public/core/TypeDiscovery.swift

@@ -0,0 +1,105 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2022 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
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+/// A function that instantiates the type described by a type descriptor.
+///
+/// - Parameters:
+///   - descriptor: A type descriptor as passed to `_swift_enumerateTypes()`'s
+///     callback.
+///
+/// - Returns: The initialized type corresponding to `descriptor`, or `nil` if
+///   the type could not be initialized. This value can subsequently be cast to
+///   `Any.Type` using `unsafeBitCast(_:to:)`.
+private typealias _TypeGetter = @convention(c) (
+  _ descriptor: UnsafeRawPointer
+) -> UnsafeRawPointer?
+
+/// A type describing another type available in the current process.
+@available(SwiftStdlib 9999, *)
+public struct _TypeRecord {
+  /// Initialize an instance of this type.
+  ///
+  /// - Parameters:
+  ///   - name: The name of this type.
+  ///   - typeInstantiator: A function to call that calls back into the runtime
+  ///     to instantiate the type.
+  fileprivate init(
+    _name name: String,
+    instantiatingTypeUsing typeInstantiator: @escaping () -> Any.Type?
+  ) {
+    self.name = name
+    _typeInstantiator = typeInstantiator
+  }
+
+  /// The name of this type.
+  public let name: String
+
+  /// Storage for `type`.
+  ///
+  /// This function calls back into the runtime to instantiate the type.
+  private let _typeInstantiator: () -> Any.Type?
+
+  /// The Swift type described by this type record.
+  ///
+  /// On first use, this property will "instantiate" the type, meaning that the
+  /// type will become fully realized and will be available for use in Swift. If
+  /// that operation fails, the value of this property is `nil`.
+  public var type: Any.Type? {
+    return _typeInstantiator()
+  }
+}
+
+@_silgen_name("swift_enumerateAllTypesFromImage")
+private func _swift_enumerateTypes(
+  fromImageAt imageAddress: UnsafeRawPointer?,
+  _ body: (
+    _ name: UnsafePointer<CChar>,
+    _ descriptor: UnsafeRawPointer,
+    _ typeGetter: @escaping _TypeGetter,
+    _ stop: inout Bool
+  ) throws -> Void
+) rethrows
+
+/// Enumerate all types in a given image.
+///
+/// - Parameters:
+///   - imageAddress: A platform-specific pointer to the image of interest. The
+///     image must have been loaded into the current process. For the binary of
+///     the calling function, you can pass `#dsohandle`. For all binaries, pass
+///     `nil` (the default.)
+///   - body: A closure to invoke once per conforming type.
+///
+/// - Throws: Whatever is thrown by `body`.
+///
+/// This function walks all known types in the given image and passes them to
+/// `body` for evaluation.
+///
+/// Generic types are not enumerated.
+///
+/// - Bug: Objective-C class lookups are not supported yet.
+@available(SwiftStdlib 9999, *)
+public func _enumerateTypes(
+  fromImageAt imageAddress: UnsafeRawPointer? = nil,
+  _ body: (_ typeRecord: _TypeRecord, _ stop: inout Bool) throws -> Void
+) rethrows {
+  try _swift_enumerateTypes(
+    fromImageAt: imageAddress
+  ) { name, descriptor, getType, stop in
+    let typeRecord = _TypeRecord(
+      _name: String(cString: name),
+      instantiatingTypeUsing: {
+        getType(descriptor).map { unsafeBitCast($0, to: Any.Type.self) }
+      }
+    )
+    try body(typeRecord, &stop)
+  }
+}

stdlib/public/runtime/ImageInspection.h

@@ -74,17 +74,27 @@ void initializeTypeMetadataRecordLookup();
 void initializeDynamicReplacementLookup();
 
 // Callbacks to register metadata from an image to the runtime.
-void addImageProtocolsBlockCallback(const void *start, uintptr_t size);
-void addImageProtocolsBlockCallbackUnsafe(const void *start, uintptr_t size);
-void addImageProtocolConformanceBlockCallback(const void *start,
+void addImageProtocolsBlockCallback(const void *image,
+                                    const void *start,
+                                    uintptr_t size);
+void addImageProtocolsBlockCallbackUnsafe(const void *image,
+                                          const void *start,
+                                          uintptr_t size);
+void addImageProtocolConformanceBlockCallback(const void *image,
+                                              const void *start,
                                               uintptr_t size);
-void addImageProtocolConformanceBlockCallbackUnsafe(const void *start,
+void addImageProtocolConformanceBlockCallbackUnsafe(const void *image,
+                                                    const void *start,
                                                     uintptr_t size);
-void addImageTypeMetadataRecordBlockCallback(const void *start,
+void addImageTypeMetadataRecordBlockCallback(const void *image,
+                                             const void *start,
                                              uintptr_t size);
-void addImageTypeMetadataRecordBlockCallbackUnsafe(const void *start,
+void addImageTypeMetadataRecordBlockCallbackUnsafe(const void *image,
+                                                   const void *start,
                                                    uintptr_t size);
-void addImageDynamicReplacementBlockCallback(const void *start, uintptr_t size,
+void addImageDynamicReplacementBlockCallback(const void *image,
+                                             const void *start,
+                                             uintptr_t size,
                                              const void *start2,
                                              uintptr_t size2);
 

stdlib/public/runtime/ImageInspectionCommon.cpp

@@ -44,7 +44,7 @@ void record(swift::MetadataSections *sections) {
 }
 
 SWIFT_RUNTIME_EXPORT
-void swift_addNewDSOImage(const void *addr) {
+void swift_addNewDSOImage(const void *imageAddress, const void *addr) {
   // We cast off the const in order to update the linked list
   // data structure. This is safe to do since we don't touch 
   // any other fields.
@@ -56,19 +56,23 @@ void swift_addNewDSOImage(const void *addr) {
   const auto &protocols_section = sections->swift5_protocols;
   const void *protocols = reinterpret_cast<void *>(protocols_section.start);
   if (protocols_section.length)
-    swift::addImageProtocolsBlockCallback(protocols, protocols_section.length);
+    swift::addImageProtocolsBlockCallback(image,
+                                          protocols,
+                                          protocols_section.length);
 
   const auto &protocol_conformances = sections->swift5_protocol_conformances;
   const void *conformances =
       reinterpret_cast<void *>(protocol_conformances.start);
   if (protocol_conformances.length)
-    swift::addImageProtocolConformanceBlockCallback(conformances,
+    swift::addImageProtocolConformanceBlockCallback(image, conformances,
                                              protocol_conformances.length);
 
   const auto &type_metadata = sections->swift5_type_metadata;
   const void *metadata = reinterpret_cast<void *>(type_metadata.start);
   if (type_metadata.length)
-    swift::addImageTypeMetadataRecordBlockCallback(metadata, type_metadata.length);
+    swift::addImageTypeMetadataRecordBlockCallback(image,
+                                                   metadata,
+                                                   type_metadata.length);
 
   const auto &dynamic_replacements = sections->swift5_replace;
   const auto *replacements =
@@ -77,7 +81,7 @@ void swift_addNewDSOImage(const void *addr) {
     const auto &dynamic_replacements_some = sections->swift5_replac2;
     const auto *replacements_some =
       reinterpret_cast<void *>(dynamic_replacements_some.start);
-    swift::addImageDynamicReplacementBlockCallback(
+    swift::addImageDynamicReplacementBlockCallback(image,
         replacements, dynamic_replacements.length, replacements_some,
         dynamic_replacements_some.length);
   }
@@ -89,7 +93,7 @@ void swift::initializeProtocolLookup() {
     const swift::MetadataSectionRange &protocols =
       sections->swift5_protocols;
     if (protocols.length)
-      addImageProtocolsBlockCallbackUnsafe(
+      addImageProtocolsBlockCallbackUnsafe(image,
           reinterpret_cast<void *>(protocols.start), protocols.length);
 
     if (sections->next == registered)
@@ -104,7 +108,7 @@ void swift::initializeProtocolConformanceLookup() {
     const swift::MetadataSectionRange &conformances =
         sections->swift5_protocol_conformances;
     if (conformances.length)
-      addImageProtocolConformanceBlockCallbackUnsafe(
+      addImageProtocolConformanceBlockCallbackUnsafe(image,
           reinterpret_cast<void *>(conformances.start), conformances.length);
 
     if (sections->next == registered)
@@ -119,7 +123,7 @@ void swift::initializeTypeMetadataRecordLookup() {
     const swift::MetadataSectionRange &type_metadata =
         sections->swift5_type_metadata;
     if (type_metadata.length)
-      addImageTypeMetadataRecordBlockCallbackUnsafe(
+      addImageTypeMetadataRecordBlockCallbackUnsafe(image,
           reinterpret_cast<void *>(type_metadata.start), type_metadata.length);
 
     if (sections->next == registered)
@@ -177,4 +181,4 @@ size_t swift_getMetadataSectionCount() {
 
 #endif // !defined(__MACH__)
 
-#endif // SWIFT_RUNTIME_IMAGEINSPECTIONCOMMON_H
+#endif // SWIFT_RUNTIME_IMAGEINSPECTIONCOMMON_H

stdlib/public/runtime/ImageInspectionCommon.h

@@ -59,7 +59,7 @@ struct SectionInfo {
 
 // Called by injected constructors when a dynamic library is loaded.
 SWIFT_RUNTIME_EXPORT
-void swift_addNewDSOImage(const void *addr);
+void swift_addNewDSOImage(const void *imageAddress, const void *addr);
 
 #ifndef NDEBUG
 

stdlib/public/runtime/ImageInspectionMachO.cpp

@@ -52,7 +52,8 @@ using mach_header_platform = mach_header;
 #endif
 
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
-         void CONSUME_BLOCK(const void *start, uintptr_t size)>
+         void CONSUME_BLOCK(const void *imageAddress,
+                            const void *start, uintptr_t size)>
 void addImageCallback(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
   assert(mh->magic == MH_MAGIC_64 && "loaded non-64-bit image?!");
@@ -68,17 +69,19 @@ void addImageCallback(const mach_header *mh) {
   if (!section)
     return;
 
-  CONSUME_BLOCK(section, size);
+  CONSUME_BLOCK(mh, section, size);
 }
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
-         void CONSUME_BLOCK(const void *start, uintptr_t size)>
+         void CONSUME_BLOCK(const void *imageAddress,
+                            const void *start, uintptr_t size)>
 void addImageCallback(const mach_header *mh, intptr_t vmaddr_slide) {
   addImageCallback<SEGMENT_NAME, SECTION_NAME, CONSUME_BLOCK>(mh);
 }
 
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
           const char *SEGMENT_NAME2, const char *SECTION_NAME2,
-          void CONSUME_BLOCK(const void *start, uintptr_t size,
+          void CONSUME_BLOCK(const void *imageAddress,
+                             const void *start, uintptr_t size,
                              const void *start2, uintptr_t size2)>
 void addImageCallback2Sections(const mach_header *mh) {
 #if __POINTER_WIDTH__ == 64
@@ -104,11 +107,12 @@ void addImageCallback2Sections(const mach_header *mh) {
   if (!section2)
     size2 = 0;
 
-  CONSUME_BLOCK(section, size, section2, size2);
+  CONSUME_BLOCK(mh, section, size, section2, size2);
 }
 template <const char *SEGMENT_NAME, const char *SECTION_NAME,
           const char *SEGMENT_NAME2, const char *SECTION_NAME2,
-          void CONSUME_BLOCK(const void *start, uintptr_t size,
+          void CONSUME_BLOCK(const void *imageAddress,
+                             const void *start, uintptr_t size,
                              const void *start2, uintptr_t size2)>
 void addImageCallback2Sections(const mach_header *mh, intptr_t vmaddr_slide) {
   addImageCallback2Sections<SEGMENT_NAME, SECTION_NAME,