Refactor test content discovery to produce a sequence.

Sources/Testing/Discovery+Platform.swift

@@ -35,7 +35,7 @@ struct SectionBounds: Sendable {
   ///
   /// - Returns: A sequence of structures describing the bounds of metadata
   ///   sections of the given kind found in the current process.
-  static func all(_ kind: Kind) -> some RandomAccessCollection<SectionBounds> {
+  static func all(_ kind: Kind) -> some Sequence<SectionBounds> {
     _sectionBounds(kind)
   }
 }
@@ -237,7 +237,7 @@ private func _sectionBounds(_ kind: SectionBounds.Kind) -> [SectionBounds] {
   case .typeMetadata:
     ".sw5tymd"
   }
-  return HMODULE.all.compactMap { _findSection(named: sectionName, in: $0) }
+  return HMODULE.all.lazy.compactMap { _findSection(named: sectionName, in: $0) }
 }
 #else
 /// The fallback implementation of ``SectionBounds/all(_:)`` for platforms that

Sources/Testing/Discovery.swift

@@ -29,21 +29,6 @@ public typealias __TestContentRecord = (
   reserved2: UInt
 )
 
-/// Resign any pointers in a test content record.
-///
-/// - Parameters:
-///   - record: The test content record to resign.
-///
-/// - Returns: A copy of `record` with its pointers resigned.
-///
-/// On platforms/architectures without pointer authentication, this function has
-/// no effect.
-private func _resign(_ record: __TestContentRecord) -> __TestContentRecord {
-  var record = record
-  record.accessor = record.accessor.map(swt_resign)
-  return record
-}
-
 // MARK: -
 
 /// A protocol describing a type that can be stored as test content at compile
@@ -79,42 +64,68 @@ protocol TestContent: ~Copyable {
   associatedtype TestContentAccessorHint: Sendable = Never
 }
 
-extension TestContent where Self: ~Copyable {
-  /// Enumerate all test content records found in the given test content section
-  /// in the current process that match this ``TestContent`` type.
+// MARK: - Individual test content records
+
+/// A type describing a test content record of a particular (known) type.
+///
+/// Instances of this type can be created by calling
+/// ``TestContent/allTestContentRecords()`` on a type that conforms to
+/// ``TestContent``.
+///
+/// This type is not part of the public interface of the testing library. In the
+/// future, we could make it public if we want to support runtime discovery of
+/// test content by second- or third-party code.
+struct TestContentRecord<T>: Sendable where T: ~Copyable {
+  /// The base address of the image containing this instance, if known.
   ///
-  /// - Parameters:
-  ///   - sectionBounds: The bounds of the section to inspect.
+  /// This property is not available on platforms such as WASI that statically
+  /// link to the testing library.
   ///
-  /// - Returns: A sequence of tuples. Each tuple contains an instance of
-  ///   `__TestContentRecord` and the base address of the image containing that
-  ///   test content record. Only test content records matching this
-  ///   ``TestContent`` type's requirements are included in the sequence.
-  private static func _testContentRecords(in sectionBounds: SectionBounds) -> some Sequence<(imageAddress: UnsafeRawPointer?, record: __TestContentRecord)> {
-    sectionBounds.buffer.withMemoryRebound(to: __TestContentRecord.self) { records in
-      records.lazy
-        .filter { $0.kind == testContentKind }
-        .map(_resign)
-        .map { (sectionBounds.imageAddress, $0) }
-    }
+  /// - Note: The value of this property is distinct from the pointer returned
+  ///   by `dlopen()` (on platforms that have that function) and cannot be used
+  ///   with interfaces such as `dlsym()` that expect such a pointer.
+#if SWT_NO_DYNAMIC_LINKING
+  @available(*, unavailable, message: "Image addresses are not available on this platform.")
+#endif
+  nonisolated(unsafe) var imageAddress: UnsafeRawPointer?
+
+  /// The underlying test content record loaded from a metadata section.
+  private var _record: __TestContentRecord
+
+  fileprivate init(imageAddress: UnsafeRawPointer?, record: __TestContentRecord) {
+#if !SWT_NO_DYNAMIC_LINKING
+    self.imageAddress = imageAddress
+#endif
+    self._record = record
   }
+}
 
-  /// Call the given accessor function.
+// This `T: TestContent` constraint is in an extension in order to work around a
+// compiler crash. SEE: rdar://143049814
+extension TestContentRecord where T: TestContent & ~Copyable {
+  /// The context value for this test content record.
+  var context: UInt {
+    _record.context
+  }
+
+  /// Load the value represented by this record.
   ///
   /// - Parameters:
-  ///   - accessor: The C accessor function of a test content record matching
-  ///     this type.
-  ///   - hint: A pointer to a kind-specific hint value. If not `nil`, this
-  ///     value is passed to `accessor`, allowing that function to determine if
-  ///     its record matches before initializing its out-result.
+  ///   - hint: An optional hint value. If not `nil`, this value is passed to
+  ///     the accessor function of the underlying test content record.
   ///
-  /// - Returns: An instance of this type's accessor result or `nil` if an
-  ///   instance could not be created (or if `hint` did not match.)
+  /// - Returns: An instance of the associated ``TestContentAccessorResult``
+  ///   type, or `nil` if the underlying test content record did not match
+  ///   `hint` or otherwise did not produce a value.
   ///
-  /// The caller is responsible for ensuring that `accessor` corresponds to a
-  /// test content record of this type.
-  private static func _callAccessor(_ accessor: SWTTestContentAccessor, withHint hint: TestContentAccessorHint?) -> TestContentAccessorResult? {
-    withUnsafeTemporaryAllocation(of: TestContentAccessorResult.self, capacity: 1) { buffer in
+  /// If this function is called more than once on the same instance, a new
+  /// value is created on each call.
+  func load<R>(withHint hint: T.TestContentAccessorHint? = nil) -> R? where R == T.TestContentAccessorResult, R: ~Copyable {
+    guard let accessor = _record.accessor.map(swt_resign) else {
+      return nil
+    }
+
+    return withUnsafeTemporaryAllocation(of: T.TestContentAccessorResult.self, capacity: 1) { buffer in
       let initialized = if let hint {
         withUnsafePointer(to: hint) { hint in
           accessor(buffer.baseAddress!, hint)
@@ -128,45 +139,23 @@ extension TestContent where Self: ~Copyable {
       return buffer.baseAddress!.move()
     }
   }
+}
 
-  /// The type of callback called by ``enumerateTestContent(withHint:_:)``.
-  ///
-  /// - Parameters:
-  ///   - imageAddress: A pointer to the start of the image. This value is _not_
-  ///     equal to the value returned from `dlopen()`. On platforms that do not
-  ///     support dynamic loading (and so do not have loadable images), the
-  ///     value of this argument is unspecified.
-  ///   - content: The value produced by the test content record's accessor.
-  ///   - context: Context associated with `content`. The value of this argument
-  ///     is dependent on the type of test content being enumerated.
-  ///   - stop: An `inout` boolean variable indicating whether test content
-  ///     enumeration should stop after the function returns. Set `stop` to
-  ///     `true` to stop test content enumeration.
-  typealias TestContentEnumerator = (_ imageAddress: UnsafeRawPointer?, _ content: borrowing TestContentAccessorResult, _ context: UInt, _ stop: inout Bool) -> Void
+// MARK: - Enumeration of test content records
 
-  /// Enumerate all test content of this type known to Swift and found in the
-  /// current process.
-  ///
-  /// - Parameters:
-  ///   - hint: An optional hint value. If not `nil`, this value is passed to
-  ///     the accessor function of each test content record whose `kind` field
-  ///     matches this type's ``testContentKind`` property.
-  ///   - body: A function to invoke, once per matching test content record.
+extension TestContent where Self: ~Copyable {
+  /// Get all test content of this type known to Swift and found in the current
+  /// process.
   ///
-  /// This function uses a callback instead of producing a sequence because it
-  /// is used with move-only types (specifically ``ExitTest``) and
-  /// `Sequence.Element` must be copyable.
-  static func enumerateTestContent(withHint hint: TestContentAccessorHint? = nil, _ body: TestContentEnumerator) {
-    let testContentRecords = SectionBounds.all(.testContent).lazy.flatMap(_testContentRecords(in:))
-
-    var stop = false
-    for (imageAddress, record) in testContentRecords {
-      if let accessor = record.accessor, let result = _callAccessor(accessor, withHint: hint) {
-        // Call the callback.
-        body(imageAddress, result, record.context, &stop)
-        if stop {
-          break
-        }
+  /// - Returns: A sequence of instances of ``TestContentRecord``. Only test
+  ///   content records matching this ``TestContent`` type's requirements are
+  ///   included in the sequence.
+  static func allTestContentRecords() -> some Sequence<TestContentRecord<Self>> {
+    SectionBounds.all(.testContent).lazy.flatMap { sb in
+      sb.buffer.withMemoryRebound(to: __TestContentRecord.self) { records in
+        records.lazy
+          .filter { $0.kind == testContentKind }
+          .map { TestContentRecord<Self>(imageAddress: sb.imageAddress, record: $0) }
       }
     }
   }

Sources/Testing/ExitTests/ExitTest.swift

@@ -242,24 +242,17 @@ extension ExitTest {
   /// - Returns: The specified exit test function, or `nil` if no such exit test
   ///   could be found.
   public static func find(identifiedBy id: ExitTest.ID) -> Self? {
-    var result: Self?
-
-    enumerateTestContent(withHint: id) { _, exitTest, _, stop in
-      if exitTest.id == id {
-        result = ExitTest(__identifiedBy: id, body: exitTest.body)
-        stop = true
+    for record in Self.allTestContentRecords() {
+      if let exitTest = record.load(withHint: id) {
+        return exitTest
       }
     }
 
-    if result == nil {
-      // Call the legacy lookup function that discovers tests embedded in types.
-      result = types(withNamesContaining: exitTestContainerTypeNameMagic).lazy
-        .compactMap { $0 as? any __ExitTestContainer.Type }
-        .first { $0.__id == id }
-        .map { ExitTest(__identifiedBy: $0.__id, body: $0.__body) }
-    }
-
-    return result
+    // Call the legacy lookup function that discovers tests embedded in types.
+    return types(withNamesContaining: exitTestContainerTypeNameMagic).lazy
+      .compactMap { $0 as? any __ExitTestContainer.Type }
+      .first { $0.__id == id }
+      .map { ExitTest(__identifiedBy: $0.__id, body: $0.__body) }
   }
 }
 

Sources/Testing/Test+Discovery.swift

@@ -22,59 +22,51 @@ extension Test: TestContent {
   /// The order of values in this sequence is unspecified.
   static var all: some Sequence<Self> {
     get async {
-      var generators = [@Sendable () async -> [Self]]()
+      // The result is a set rather than an array to deduplicate tests that were
+      // generated multiple times (e.g. from multiple discovery modes or from
+      // defective test records.)
+      var result = Set<Self>()
 
       // Figure out which discovery mechanism to use. By default, we'll use both
       // the legacy and new mechanisms, but we can set an environment variable
       // to explicitly select one or the other. When we remove legacy support,
       // we can also remove this enumeration and environment variable check.
-      enum DiscoveryMode {
-        case tryBoth
-        case newOnly
-        case legacyOnly
-      }
-      let discoveryMode: DiscoveryMode = switch Environment.flag(named: "SWT_USE_LEGACY_TEST_DISCOVERY") {
+      let (useNewMode, useLegacyMode) = switch Environment.flag(named: "SWT_USE_LEGACY_TEST_DISCOVERY") {
       case .none:
-        .tryBoth
+        (true, true)
       case .some(true):
-        .legacyOnly
+        (false, true)
       case .some(false):
-        .newOnly
+        (true, false)
       }
 
-      // Walk all test content and gather generator functions. Note we don't
-      // actually call the generators yet because enumerating test content may
-      // involve holding some internal lock such as the ones in libobjc or
-      // dl_iterate_phdr(), and we don't want to accidentally deadlock if the
-      // user code we call ends up loading another image.
-      if discoveryMode != .legacyOnly {
-        enumerateTestContent { imageAddress, generator, _, _ in
-          generators.append { @Sendable in
-            await [generator()]
+      // Walk all test content and gather generator functions, then call them in
+      // a task group and collate their results.
+      if useNewMode {
+        let generators = Self.allTestContentRecords().lazy.compactMap { $0.load() }
+        await withTaskGroup(of: Self.self) { taskGroup in
+          for generator in generators {
+            taskGroup.addTask(operation: generator)
           }
+          result = await taskGroup.reduce(into: result) { $0.insert($1) }
         }
       }
 
-      if discoveryMode != .newOnly && generators.isEmpty {
-        generators += types(withNamesContaining: testContainerTypeNameMagic).lazy
+      // Perform legacy test discovery if needed.
+      if useLegacyMode && result.isEmpty {
+        let types = types(withNamesContaining: testContainerTypeNameMagic).lazy
           .compactMap { $0 as? any __TestContainer.Type }
-          .map { type in
-            { @Sendable in await type.__tests }
-          }
-      }
-
-      // *Now* we call all the generators and return their results.
-      // Reduce into a set rather than an array to deduplicate tests that were
-      // generated multiple times (e.g. from multiple discovery modes or from
-      // defective test records.)
-      return await withTaskGroup(of: [Self].self) { taskGroup in
-        for generator in generators {
-          taskGroup.addTask {
-            await generator()
+        await withTaskGroup(of: [Self].self) { taskGroup in
+          for type in types {
+            taskGroup.addTask {
+              await type.__tests
+            }
           }
+          result = await taskGroup.reduce(into: result) { $0.formUnion($1) }
         }
-        return await taskGroup.reduce(into: Set()) { $0.formUnion($1) }
       }
+
+      return result
     }
   }
 }

Tests/TestingTests/MiscellaneousTests.swift

@@ -615,43 +615,45 @@ struct MiscellaneousTests {
       0xABCD1234,
       0,
       { outValue, hint in
-        if let hint, hint.loadUnaligned(as: TestContentAccessorHint.self) != expectedHint {
+        if let hint, hint.load(as: TestContentAccessorHint.self) != expectedHint {
           return false
         }
         _ = outValue.initializeMemory(as: TestContentAccessorResult.self, to: expectedResult)
         return true
       },
-      UInt(UInt64(0x0204060801030507) & UInt64(UInt.max)),
+      UInt(truncatingIfNeeded: UInt64(0x0204060801030507)),
       0
     )
   }
 
   @Test func testDiscovery() async {
-    await confirmation("Can find a single test record") { found in
-      DiscoverableTestContent.enumerateTestContent { _, value, context, _ in
-        if value == DiscoverableTestContent.expectedResult && context == DiscoverableTestContent.expectedContext {
-          found()
-        }
-      }
-    }
-
-    await confirmation("Can find a test record with matching hint") { found in
+    // Check the type of the test record sequence (it should be lazy.)
+    let allRecords = DiscoverableTestContent.allTestContentRecords()
+    #expect(allRecords is any LazySequenceProtocol)
+    #expect(!(allRecords is [TestContentRecord<DiscoverableTestContent>]))
+
+    // It should have exactly one matching record (because we only emitted one.)
+    #expect(Array(allRecords).count == 1)
+
+    // Can find a single test record
+    #expect(allRecords.contains { record in
+      record.load() == DiscoverableTestContent.expectedResult
+        && record.context == DiscoverableTestContent.expectedContext
+    })
+
+    // Can find a test record with matching hint
+    #expect(allRecords.contains { record in
       let hint = DiscoverableTestContent.expectedHint
-      DiscoverableTestContent.enumerateTestContent(withHint: hint) { _, value, context, _ in
-        if value == DiscoverableTestContent.expectedResult && context == DiscoverableTestContent.expectedContext {
-          found()
-        }
-      }
-    }
+      return record.load(withHint: hint) == DiscoverableTestContent.expectedResult
+        && record.context == DiscoverableTestContent.expectedContext
+    })
 
-    await confirmation("Doesn't find a test record with a mismatched hint", expectedCount: 0) { found in
+    // Doesn't find a test record with a mismatched hint
+    #expect(!allRecords.contains { record in
       let hint = ~DiscoverableTestContent.expectedHint
-      DiscoverableTestContent.enumerateTestContent(withHint: hint) { _, value, context, _ in
-        if value == DiscoverableTestContent.expectedResult && context == DiscoverableTestContent.expectedContext {
-          found()
-        }
-      }
-    }
+      return record.load(withHint: hint) == DiscoverableTestContent.expectedResult
+        && record.context == DiscoverableTestContent.expectedContext
+    })
   }
 #endif
 }