[AutoDiff] Move AD leak checking utilities to shared location. (swiftlang#26373)

Move AD leak checking utilities to DifferentiableUnittest so
they can be used by multiple tests.

Author: dan-zheng

stdlib/private/DifferentiationUnittest/CMakeLists.txt

@@ -2,6 +2,6 @@ add_swift_target_library(swiftDifferentiationUnittest ${SWIFT_STDLIB_LIBRARY_BUI
   # This file should be listed first. Module name is inferred from the filename.
   DifferentiationUnittest.swift
 
-  SWIFT_COMPILE_FLAGS
+  SWIFT_MODULE_DEPENDS StdlibUnittest
   INSTALL_IN_COMPONENT stdlib-experimental
   DARWIN_INSTALL_NAME_DIR "${SWIFT_DARWIN_STDLIB_PRIVATE_INSTALL_NAME_DIR}")

stdlib/private/DifferentiationUnittest/DifferentiationUnittest.swift

@@ -10,10 +10,44 @@
 //
 //===----------------------------------------------------------------------===//
 
+import StdlibUnittest
+
 public enum _GlobalLeakCount {
   public static var count = 0
 }
 
+/// Execute body and check expected leak count.
+public func withLeakChecking(
+  expectedLeakCount: Int = 0, file: String = #file, line: UInt = #line,
+  _ body: () -> Void
+) {
+  // Note: compare expected leak count with relative leak count after
+  // running `body`.
+  // This approach is more robust than comparing leak count with zero
+  // and resetting leak count to zero, which is stateful and causes issues.
+  let beforeLeakCount = _GlobalLeakCount.count
+  body()
+  let leakCount = _GlobalLeakCount.count - beforeLeakCount
+  expectEqual(
+    expectedLeakCount, leakCount, "Leaks detected: \(leakCount)",
+    file: file, line: line)
+}
+
+public extension TestSuite {
+  /// Execute test function and check expected leak count.
+  func testWithLeakChecking(
+    _ name: String,
+    expectedLeakCount: Int = 0,
+    file: String = #file, line: UInt = #line,
+    _ testFunction: @escaping () -> Void
+  ) {
+    test(name, file: file, line: line) {
+      withLeakChecking(expectedLeakCount: expectedLeakCount, file: file,
+                       line: line, testFunction)
+    }
+  }
+}
+
 /// A type that tracks the number of live instances of a wrapped value type.
 ///
 /// `Tracked<T>` is used to check for memory leaks in functions created via

test/AutoDiff/leakchecking.swift

@@ -8,37 +8,6 @@ import DifferentiationUnittest
 
 var LeakCheckingTests = TestSuite("LeakChecking")
 
-/// Execute body and check expected leak count.
-func withLeakChecking(
-  expectedLeakCount: Int = 0, file: String = #file, line: UInt = #line,
-  _ body: () -> Void
-) {
-  // Note: compare expected leak count with relative leak count after
-  // running `body`.
-  // This approach is more robust than comparing leak count with zero
-  // and resetting leak count to zero, which is stateful and causes issues.
-  let beforeLeakCount = _GlobalLeakCount.count
-  body()
-  let leakCount = _GlobalLeakCount.count - beforeLeakCount
-  expectEqual(
-    expectedLeakCount, leakCount, "Leaks detected: \(leakCount)",
-    file: file, line: line)
-}
-
-extension TestSuite {
-  func testWithLeakChecking(
-    _ name: String,
-    expectedLeakCount: Int = 0,
-    file: String = #file, line: UInt = #line,
-    _ testFunction: @escaping () -> Void
-  ) {
-    test(name, file: file, line: line) {
-      withLeakChecking(expectedLeakCount: expectedLeakCount, file: file,
-                       line: line, testFunction)
-    }
-  }
-}
-
 struct ExampleLeakModel : Differentiable {
   var bias: Tracked<Float> = 2.0
   func applied(to input: Tracked<Float>) -> Tracked<Float> {