Refactor type enumeration loop. (#388)

Author: grynspan

Sources/Testing/ExitTests/ExitTest.swift

@@ -79,33 +79,20 @@ extension ExitTest {
   /// - Returns: The specified exit test function, or `nil` if no such exit test
   ///   could be found.
   public static func find(at sourceLocation: SourceLocation) -> Self? {
-    struct Context {
-      var sourceLocation: SourceLocation
-      var result: ExitTest?
-    }
-    var context = Context(sourceLocation: sourceLocation)
-    withUnsafeMutablePointer(to: &context) { context in
-      swt_enumerateTypes(context) { type, context in
-        let context = context!.assumingMemoryBound(to: (Context).self)
-        if let type = unsafeBitCast(type, to: Any.Type.self) as? any __ExitTestContainer.Type,
-           type.__sourceLocation == context.pointee.sourceLocation {
-          context.pointee.result = ExitTest(
-            expectedExitCondition: type.__expectedExitCondition,
-            body: type.__body,
-            sourceLocation: type.__sourceLocation
-          )
-          return false
-        }
-        return true
-      } withNamesMatching: { typeName, _ in
-        // strstr() lets us avoid copying either string before comparing.
-        Self._exitTestContainerTypeNameMagic.withCString { testContainerTypeNameMagic in
-          nil != strstr(typeName, testContainerTypeNameMagic)
-        }
+    var result: Self?
+
+    enumerateTypes(withNamesContaining: _exitTestContainerTypeNameMagic) { type, stop in
+      if let type = type as? any __ExitTestContainer.Type, type.__sourceLocation == sourceLocation {
+        result = ExitTest(
+          expectedExitCondition: type.__expectedExitCondition,
+          body: type.__body,
+          sourceLocation: type.__sourceLocation
+        )
+        stop = true
       }
     }
 
-    return context.result
+    return result
   }
 }
 

Sources/Testing/Test+Discovery.swift

@@ -48,19 +48,12 @@ extension Test {
   private static var _all: some Sequence<Self> {
     get async {
       await withTaskGroup(of: [Self].self) { taskGroup in
-        swt_enumerateTypes(&taskGroup) { type, context in
-          if let type = unsafeBitCast(type, to: Any.Type.self) as? any __TestContainer.Type {
-            let taskGroup = context!.assumingMemoryBound(to: TaskGroup<[Self]>.self)
-            taskGroup.pointee.addTask {
+        enumerateTypes(withNamesContaining: _testContainerTypeNameMagic) { type, _ in
+          if let type = type as? any __TestContainer.Type {
+            taskGroup.addTask {
               await type.__tests
             }
           }
-          return true
-        } withNamesMatching: { typeName, _ in
-          // strstr() lets us avoid copying either string before comparing.
-          Self._testContainerTypeNameMagic.withCString { testContainerTypeNameMagic in
-            nil != strstr(typeName, testContainerTypeNameMagic)
-          }
         }
 
         return await taskGroup.reduce(into: [], +=)
@@ -115,3 +108,34 @@ extension Test {
     return tests.count - originalCount
   }
 }
+
+// MARK: -
+
+/// The type of callback called by ``enumerateTypes(withNamesContaining:_:)``.
+///
+/// - Parameters:
+///   - type: A Swift type.
+///   - stop: An `inout` boolean variable indicating whether type enumeration
+///     should stop after the function returns. Set `stop` to `true` to stop
+///     type enumeration.
+typealias TypeEnumerator = (_ type: Any.Type, _ stop: inout Bool) -> Void
+
+/// Enumerate all types known to Swift found in the current process whose names
+/// contain a given substring.
+///
+/// - Parameters:
+///   - nameSubstring: A string which the names of matching classes all contain.
+///   - body: A function to invoke, once per matching type.
+func enumerateTypes(withNamesContaining nameSubstring: String, _ typeEnumerator: TypeEnumerator) {
+  withoutActuallyEscaping(typeEnumerator) { typeEnumerator in
+    withUnsafePointer(to: typeEnumerator) { context in
+      swt_enumerateTypes(withNamesContaining: nameSubstring, .init(mutating: context)) { type, stop, context in
+        let typeEnumerator = context!.load(as: TypeEnumerator.self)
+        let type = unsafeBitCast(type, to: Any.Type.self)
+        var stop2 = false
+        typeEnumerator(type, &stop2)
+        stop.pointee = stop2
+      }
+    }
+  }
+}

Sources/TestingInternals/Discovery.cpp

@@ -11,6 +11,7 @@
 #include "Discovery.h"
 
 #include <atomic>
+#include <cstring>
 #include <iterator>
 #include <type_traits>
 #include <vector>
@@ -328,13 +329,13 @@ static void enumerateTypeMetadataSections(const SectionEnumerator& body) {}
 
 #pragma mark -
 
-void swt_enumerateTypes(void *context, SWTTypeEnumerator body, SWTTypeNameFilter nameFilter) {
+void swt_enumerateTypesWithNamesContaining(const char *nameSubstring, void *context, SWTTypeEnumerator body) {
   enumerateTypeMetadataSections([=] (const void *section, size_t size) {
     auto records = reinterpret_cast<const SWTTypeMetadataRecord *>(section);
     size_t recordCount = size / sizeof(SWTTypeMetadataRecord);
 
-    bool keepGoing = true;
-    for (size_t i = 0; i < recordCount && keepGoing; i++) {
+    bool stop = false;
+    for (size_t i = 0; i < recordCount && !stop; i++) {
       const auto& record = records[i];
 
       auto contextDescriptor = record.getContextDescriptor();
@@ -348,19 +349,16 @@ void swt_enumerateTypes(void *context, SWTTypeEnumerator body, SWTTypeNameFilter
         continue;
       }
 
-      // If the caller supplied a name filtering function, check that the type's
-      // name passes. This will be more expensive than the checks above, but
-      // should be cheaper than realizing the metadata.
-      if (nameFilter) {
-        const char *typeName = contextDescriptor->getName();
-        bool nameOK = typeName && (* nameFilter)(typeName, context);
-        if (!nameOK) {
-          continue;
-        }
+      // Check that the type's name passes. This will be more expensive than the
+      // checks above, but should be cheaper than realizing the metadata.
+      const char *typeName = contextDescriptor->getName();
+      bool nameOK = typeName && nullptr != std::strstr(typeName, nameSubstring);
+      if (!nameOK) {
+        continue;
       }
 
       if (void *typeMetadata = contextDescriptor->getMetadata()) {
-        keepGoing = body(typeMetadata, context);
+        body(typeMetadata, &stop, context);
       }
     }
   });

Sources/TestingInternals/include/Discovery.h

@@ -17,45 +17,28 @@
 
 SWT_ASSUME_NONNULL_BEGIN
 
-/// The type of callback that is called by `swt_enumerateTypes()`.
+/// The type of callback called by `swt_enumerateTypes()`.
 ///
 /// - Parameters:
 ///   - typeMetadata: A type metadata pointer that can be bitcast to `Any.Type`.
+///   - stop: A pointer to a boolean variable indicating whether type
+///     enumeration should stop after the function returns. Set `*stop` to
+///     `true` to stop type enumeration.
 ///   - context: An arbitrary pointer passed by the caller to
 ///     `swt_enumerateTypes()`.
-///
-/// - Returns: Whether or not to continue enumeration.
-typedef bool (* SWTTypeEnumerator)(void *typeMetadata, void *_Null_unspecified context);
-
-/// The type name filter that is called by `swt_enumerateTypes()`.
-///
-/// - Parameters:
-///   - typeName: The name of the type being considered, as a C string.
-///   - context: An arbitrary pointer passed by the caller to
-///     `swt_enumerateTypes()`.
-///
-/// - Returns: Whether or not the type named by `typeName` should be passed to
-///   the corresponding enumerator function.
-typedef bool (* SWTTypeNameFilter)(const char *typeName, void *_Null_unspecified context);
+typedef void (* SWTTypeEnumerator)(void *typeMetadata, bool *stop, void *_Null_unspecified context);
 
 /// Enumerate all types known to Swift found in the current process.
 ///
 /// - Parameters:
-///   - nameFilter: If not `nullptr`, a filtering function that checks if a type
-///     name is valid before realizing the type.
-///   - body: A function to invoke. `context` is passed to it along with a
-///     type metadata pointer (which can be bitcast to `Any.Type`.)
+///   - nameSubstring: A string which the names of matching classes all contain.
 ///   - context: An arbitrary pointer to pass to `body`.
-///
-/// This function may enumerate the same type more than once (for instance, if
-/// it is present in an image's metadata table multiple times, or if it is an
-/// Objective-C class implemented in Swift.) Callers are responsible for
-/// deduping type metadata pointers passed to `body`.
-SWT_EXTERN void swt_enumerateTypes(
+///   - body: A function to invoke, once per matching type.
+SWT_EXTERN void swt_enumerateTypesWithNamesContaining(
+  const char *nameSubstring,
   void *_Null_unspecified context,
-  SWTTypeEnumerator body,
-  SWTTypeNameFilter _Nullable nameFilter
-) SWT_SWIFT_NAME(swt_enumerateTypes(_:_:withNamesMatching:));
+  SWTTypeEnumerator body
+) SWT_SWIFT_NAME(swt_enumerateTypes(withNamesContaining:_:_:));
 
 SWT_ASSUME_NONNULL_END