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,57 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+@_silgen_name("swift_enumerateConformingTypesFromImage")
+private func _swift_enumerateTypes(
+  fromImageAt imageAddress: UnsafeRawPointer?,
+  conformingTo protocolType: Any.Type,
+  _ body: (_ type: Any.Type, _ stop: inout Bool) throws -> Void
+) rethrows
+
+/// Enumerate all types in a given image that conform to the specified Swift
+/// protocol.
+///
+/// - 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.)
+///   - protocolType: The protocol of interest. This protocol does not need to
+///     be declared in the image referenced by `imageAddress`.
+///   - 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 checks each for
+/// conformance to `protocolType`. If the type conforms, it is passed to `body`.
+///
+/// Only Swift protocol types (for example, `Hashable.self` or
+/// `(Hashable & Codable).self`) should be passed to this function. All other
+/// types will produce a fatal error at runtime.
+///
+/// - Bug: Generic types that conform to `protocolType` are not enumerated by
+///   this function.
+///
+/// - Bug: Objective-C protocol conformance lookups are not supported yet.
+@available(macOS 9999.0, iOS 9999.0, watchOS 9999.0, tvOS 9999.0, *)
+public func _enumerateTypes(
+  fromImageAt imageAddress: UnsafeRawPointer? = nil,
+  conformingTo protocolType: Any.Type,
+  _ body: (_ type: Any.Type, _ stop: inout Bool) throws -> Void
+) rethrows {
+  try _swift_enumerateTypes(
+    fromImageAt: imageAddress,
+    conformingTo: protocolType,
+    body
+  )
+}
+

stdlib/public/runtime/ProtocolConformance.cpp

@@ -29,10 +29,17 @@
 #include "llvm/ADT/DenseMap.h"
 #include "../CompatibilityOverride/CompatibilityOverride.h"
 #include "ImageInspection.h"
+#include "ImageInspectionCommon.h"
 #include "Private.h"
+#include "TypeDiscovery.h"
 
 #include <vector>
 
+#if defined(__MACH__)
+#include <mach-o/dyld.h>
+#include <mach-o/getsect.h>
+#endif
+
 #if __has_include(<mach-o/dyld_priv.h>)
 #include <mach-o/dyld_priv.h>
 #define DYLD_EXPECTED_SWIFT_OPTIMIZATIONS_VERSION 1u
@@ -1271,5 +1278,188 @@ const Metadata *swift::findConformingSuperclass(
   return conformingType;
 }
 
+/// Walk the conformances in a \c ConformanceSection instance, check if they
+/// specify conformance to all the specified protocols, and call \a body for the
+/// types that match.
+SWIFT_CC(swift)
+static void _enumerateConformingTypesFromSection(
+  const ConformanceSection *section,
+  size_t count,
+  const llvm::ArrayRef<ProtocolDescriptorRef>& protocols,
+  TypeEnumerationFunction body,
+  void *bodyContext,
+  bool *stop,
+  SWIFT_CONTEXT void *context,
+  SWIFT_ERROR_RESULT SwiftError **error) {
+
+  for (size_t i = 0; i < count; i++) {
+    for (const auto& record : *section) {
+      auto protocolDescriptor = record->getProtocol();
+      auto metadata = record->getCanonicalTypeMetadata();
+      if (!protocolDescriptor || !metadata) {
+        // This record does not contain the info we're interested in.
+        continue;
+      }
+
+      for (const auto& protocol : protocols) {
+        const auto& swiftProtocol = protocol.getSwiftProtocol();
+        if (!swift_conformsToProtocol(metadata, swiftProtocol)) {
+          // This type does not conform to the specified protocol(s), so skip.
+          continue;
+        }
+      }
+
+      body(metadata, stop, bodyContext, error);
+      if (*stop || *error) {
+        return;
+      }
+    }
+  }
+}
+
+#if defined(__MACH__)
+/// Get the conformance section of a given Mach-O image.
+///
+/// If the image does not have a conformance section, returns \c llvm::None.
+static llvm::Optional<ConformanceSection>
+_getConformanceSectionFromMachImage(const void *imageAddress) {
+  unsigned long sectionSize = 0;
+#ifdef __LP64__
+  typedef const mach_header_64 *MachHeaderAddress;
+#else
+  typedef const mach_header *MachHeaderAddress;
+#endif
+  const auto& sectionData = getsectiondata((MachHeaderAddress)imageAddress,
+                                           SEG_TEXT,
+                                           MachOProtocolConformancesSection,
+                                           &sectionSize);
+  if (sectionData) {
+    return ConformanceSection(sectionData, sectionSize);
+  }
+
+  return llvm::None;
+}
+#endif
+
+SWIFT_CC(swift)
+void swift::swift_enumerateConformingTypesFromImage(
+  const void *imageAddress,
+  const Metadata *protocolMetadata,
+  TypeEnumerationFunction body,
+  void *bodyContext,
+  SWIFT_CONTEXT void *context,
+  SWIFT_ERROR_RESULT SwiftError **error) {
+
+  // Find the protocol(s) requested by the caller.
+  auto existentialType = dyn_cast<ExistentialTypeMetadata>(protocolMetadata);
+  if (!existentialType) {
+    auto typeName = swift_getTypeName(protocolMetadata, true);
+    swift::fatalError(
+      0,
+      "_enumerateTypes(fromImageAt:conformingTo:_:) can only test for "
+      "conformance to a protocol, but %.*s is not a protocol.",
+      (int)typeName.length, typeName.data);
+  }
+  const auto& protocols = existentialType->getProtocols();
+
+  if (protocols.empty()) {
+    swift::fatalError(
+      0,
+      "Any is not a protocol and _enumerateTypes(fromImageAt:conformingTo:_:) "
+      "cannot test for conformance to it.");
+  }
+
+#if SWIFT_OBJC_INTEROP
+  // Check for any Objective-C protocols in the input. We do not currently
+  // support searching for Objective-C protocol conformances because they do not
+  // produce conformance records in the compiled binary. In the future, if we
+  // detect any such protocols, we can use objc_copyClassList() to add to the
+  // enumerated set (taking care not to enumerate a class twice if it conforms
+  // to an Objective-C protocol *and* a Swift protocol.)
+  for (const auto& protocol : protocols) {
+    if (protocol.isObjC()) {
+      swift::fatalError(
+        0,
+        "_enumerateTypes(fromImageAt:conformingTo:_:) does not currently "
+        "support Objective-C protocols such as %s.",
+        protocol.getName());
+    }
+  }
+#endif
+
+#if defined(__WASM__)
+  // Swift+WASM does not support dynamic linking and #dsohandle is 0, so
+  // ignore the passed image address and iterate over any/all available
+  // sections. Presumably only one will be available.
+  imageAddress = nullptr;
+#endif
+
+  // Callback parameters (hoisted out of the loops.)
+  bool stop = false;
+  *error = nullptr;
+
+#if defined(__MACH__)
+  // This platform has dyld and we can use it instead of walking the section
+  // table. Takes care of the nuances of DYLD shared cache support.
+  if (imageAddress) {
+    if (auto section = _getConformanceSectionFromMachImage(imageAddress)) {
+      _enumerateConformingTypesFromSection(section.getPointer(), 1,
+                                           protocols, body, bodyContext,
+                                           &stop, context, error);
+    }
+
+  } else {
+    // Walk all loaded images and consider each one in turn. The calls into dyld
+    // aren't thread-safe (another thread could modify the set of loaded images)
+    // but not in a "I'm going to crash" sense, just a "missed an image" sense.
+    auto count = _dyld_image_count();
+    for (uint32_t i = 0; i < count; i++) {
+      const auto& imageAddress = _dyld_get_image_header(i);
+      if (!imageAddress) {
+        continue;
+      }
+
+      if (auto section = _getConformanceSectionFromMachImage(imageAddress)) {
+        _enumerateConformingTypesFromSection(section.getPointer(), 1,
+                                             protocols, body, bodyContext,
+                                             &stop, context, error);
+        if (stop || *error) {
+          return;
+        }
+      }
+    }
+  }
+
+#else
+  // This platform doesn't have getsectiondata(), so walk all sections in the
+  // conformances table.
+  auto snapshot = Conformances.get().SectionsToScan.snapshot();
+
+  if (imageAddress) {
+    // FIXME: This operation is somewhat costly and could probably be optimized on a per-platform basis.
+    for (const auto& section : snapshot) {
+      // Check if this section came from the image of interest.
+      SymbolInfo info;
+      if (!lookupSymbol(section.begin(), &info)) {
+        continue;
+      }
+
+      if (info.baseAddress == imageAddress) {
+        _enumerateConformingTypesFromSection(&section, 1,
+                                             protocols, body, bodyContext,
+                                             &stop, context, error);
+        return;
+      }
+    }
+
+  } else {
+    // Walk all sections.
+    _enumerateConformingTypesFromSection(snapshot.begin(), snapshot.count(),
+                                         protocols, body, bodyContext,
+                                         &stop, context, error);
+  }
+#endif
+}
+
 #define OVERRIDE_PROTOCOLCONFORMANCE COMPATIBILITY_OVERRIDE
 #include COMPATIBILITY_OVERRIDE_INCLUDE_PATH

stdlib/public/runtime/TypeDiscovery.h

@@ -0,0 +1,64 @@
+//===--- TypeDiscovery.h - Dynamic type lookup at runtime--------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Functions to look up types in the Swift type system at runtime.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_RUNTIME_TYPE_DISCOVERY_H
+#define SWIFT_RUNTIME_TYPE_DISCOVERY_H
+
+#include "swift/Runtime/Config.h"
+#include "swift/Runtime/Error.h"
+#include "swift/Runtime/Metadata.h"
+
+namespace swift {
+  /// The closure type used by \c swift_enumerateConformingTypesFromImage().
+  typedef void (* TypeEnumerationFunction)(
+    const Metadata *type,
+    bool *stop,
+    SWIFT_CONTEXT void *context,
+    SWIFT_ERROR_RESULT SwiftError **error) SWIFT_CC(swift);
+
+  /// Enumerate a list of all types (as Metadata) in a given image that conform
+  /// to the specified Swift protocol.
+  ///
+  /// \param imageAddress A platform-specific pointer to the image of interest.
+  ///   The image must have been loaded into the current process. Pass
+  ///   \c nullptr to search all loaded images.
+  /// \param protocolMetadata The protocol of interest. This protocol does not
+  ///   need to be declared in the image referenced by \a imageAddress.
+  /// \param body A closure to invoke once per matching type.
+  /// \param bodyContext The Swift context pointer for \a body.
+  /// \param context The Swift context pointer for this function. Ignored.
+  /// \param error An error pointer passed to \a body. If \a body throws an
+  /// 	error, enumeration is stopped and the error is rethrown.
+  ///
+  /// This function walks all known types in the given image and checks each for
+  /// conformance to \a protocol. Any types that do conform are passed to
+  /// \c body for evaluation.
+  ///
+  /// \bug Generic types with constrained conformance to \a protocol are not
+  ///   supported.
+  ///
+  /// \bug Objective-C protocol conformance lookups are not supported.
+  SWIFT_CC(swift) SWIFT_RUNTIME_STDLIB_INTERNAL
+  void swift_enumerateConformingTypesFromImage(
+    const void *imageAddress,
+    const Metadata *protocolMetadata,
+    TypeEnumerationFunction body,
+    void *bodyContext,
+    SWIFT_CONTEXT void *context,
+    SWIFT_ERROR_RESULT SwiftError **error);
+} // end namespace swift
+
+#endif /* SWIFT_RUNTIME_TYPE_DISCOVERY_H */

test/Runtime/type_discovery.swift

@@ -0,0 +1,46 @@
+// RUN: %target-run-simple-swift
+// REQUIRES: executable_test
+
+import StdlibUnittest
+
+protocol P { }
+struct A: P { }
+class B: P { }
+enum C: P { }
+struct D /* does not conform */ { }
+
+guard #available(macOS 9999.0, iOS 9999.0, watchOS 9999.0, tvOS 9999.0, *) else {
+  print("Skipping tests in \(#fileID) due to unavailability.")
+  return
+}
+
+// Don't use #dsohandle on Windows because __ImageBase is not linked properly
+// when building the test target and we get a link-time error trying to use it.
+let address: UnsafeRawPointer?
+#if os(Windows)
+address = nil
+#else
+address = #dsohandle
+#endif
+
+do {
+  var typeList = [P.Type]()
+
+  _enumerateTypes(fromImageAt: address, conformingTo: P.self) { type, _ in
+    typeList.append(type as! P.Type)
+  }
+  expectTrue(typeList.contains(where: { $0 == A.self }))
+  expectTrue(typeList.contains(where: { $0 == B.self }))
+  expectTrue(typeList.contains(where: { $0 == C.self }))
+  expectFalse(typeList.contains(where: { $0 == D.self }))
+}
+
+// Test that `stop` works.
+do {
+  var iterationCount = 0
+  _enumerateTypes(fromImageAt: address, conformingTo: P.self) { type, _ in
+    iterationCount += 1
+    stop = true
+  }
+  expectEqual(iterationCount, 1)
+}