[Test] Simplified registration.

At the cost of adding an unsafe bitcast implementation detail,
simplified the code involved to register a new FunctionTest when adding
one.

Also simplifies how Swift native `FunctionTest`s are registered with the
C++ registry.

Now, the to-be-executed thin closure for native Swift `FunctionTest`s is
stored under within the swift::test::FunctionTest instance corresponding
to it.  Because its type isn't representable in C++, `void *` is used
instead.  When the FunctionTest is invoked, a thunk is called which
takes the actual test function and bridged versions of the arguments.
That thunk unwraps the arguments, casts the stored function to the
appropriate type, and invokes it.

Thanks to Andrew Trick for the idea.

Author: nate-chandler

SwiftCompilerSources/Sources/SIL/Test.swift

@@ -17,7 +17,7 @@
 //       }
 // - In SwiftCompilerSources/Sources/SIL/Test.swift's registerSILTests function,
 //   register the new test:
-//       registerFunctionTest(myNewTest, implementation: { myNewTest.run($0, $1, $2) })
+//       registerFunctionTest(myNewTest)
 //
 //===----------------------------------------------------------------------===//
 //
@@ -90,22 +90,21 @@
 import Basic
 import SILBridging
 
+/// The primary interface to in-IR tests.
 public struct FunctionTest {
-  public var name: String
-  public typealias FunctionTestImplementation = (Function, TestArguments, TestContext) -> ()
-  private var implementation: FunctionTestImplementation
-  init(name: String, implementation: @escaping FunctionTestImplementation) {
+  let name: String
+  let invocation: FunctionTestInvocation
+
+  public init(name: String, invocation: @escaping FunctionTestInvocation) {
     self.name = name
-    self.implementation = implementation
-  }
-  fileprivate func run(
-    _ function: BridgedFunction, 
-    _ arguments: BridgedTestArguments, 
-    _ test: BridgedTestContext) {
-    implementation(function.function, arguments.native, test.native)
+    self.invocation = invocation
   }
 }
 
+/// The type of the closure passed to a FunctionTest.
+public typealias FunctionTestInvocation = @convention(thin) (Function, TestArguments, TestContext) -> ()
+
+/// Wraps the arguments specified in the specify_test instruction.
 public struct TestArguments {
   public var bridged: BridgedTestArguments
   fileprivate init(bridged: BridgedTestArguments) {
@@ -128,6 +127,7 @@ extension BridgedTestArguments {
   public var native: TestArguments { TestArguments(bridged: self) }
 }
 
+/// An interface to the various analyses that are available.
 public struct TestContext {
   public var bridged: BridgedTestContext
   fileprivate init(bridged: BridgedTestContext) {
@@ -139,16 +139,52 @@ extension BridgedTestContext {
   public var native: TestContext { TestContext(bridged: self) }
 }
 
-private func registerFunctionTest(
-  _ test: FunctionTest,
-  implementation: BridgedFunctionTestThunk) {
+/// Registration of each test in the SIL module.
+public func registerSILTests() {
+  // Register each test.
+  registerFunctionTest(parseTestSpecificationTest)
+
+  // Finally register the thunk they all call through.
+  registerFunctionTestThunk(functionTestThunk)
+}
+
+
+private func registerFunctionTest(_ test: FunctionTest) {
   test.name._withStringRef { ref in
-    registerFunctionTest(ref, implementation)
+    registerFunctionTest(ref, eraseInvocation(test.invocation))
   }
 }
 
-public func registerSILTests() {
-  registerFunctionTest(parseTestSpecificationTest, implementation: { parseTestSpecificationTest.run($0, $1, $2) })
+/// The function called by the swift::test::FunctionTest which invokes the
+/// actual test function.
+///
+/// This function is necessary because tests need to be written in terms of
+/// native Swift types (Function, TestArguments, TestContext) rather than their
+/// bridged variants, but such a function isn't representable in C++.  This
+/// thunk unwraps the bridged types and invokes the real function.
+private func functionTestThunk(
+  _ erasedInvocation: UnsafeMutableRawPointer,
+  _ function: BridgedFunction, 
+  _ arguments: BridgedTestArguments, 
+  _ context: BridgedTestContext) {
+  let invocation = uneraseInvocation(erasedInvocation)
+  invocation(function.function, arguments.native, context.native)
+}
+
+/// Bitcast a thin test closure to void *.
+///
+/// Needed so that the closure can be represented in C++ for storage in the test
+/// registry.
+private func eraseInvocation(_ invocation: FunctionTestInvocation) -> UnsafeMutableRawPointer {
+  return unsafeBitCast(invocation, to: UnsafeMutableRawPointer.self)
+}
+
+/// Bitcast a void * to a thin test closure.
+///
+/// Needed so that the closure stored in the C++ test registry can be invoked
+/// via the functionTestThunk.
+private func uneraseInvocation(_ erasedInvocation: UnsafeMutableRawPointer) -> FunctionTestInvocation {
+  return unsafeBitCast(erasedInvocation, to: FunctionTestInvocation.self)
 }
 
 // Arguments:

include/swift/SIL/SILBridging.h

@@ -1527,6 +1527,10 @@ struct Arguments;
 class FunctionTest;
 } // namespace swift::test
 
+struct BridgedFunctionTest {
+  swift::test::FunctionTest *_Nonnull test;
+};
+
 struct BridgedTestArguments {
   swift::test::Arguments *_Nonnull arguments;
 
@@ -1562,10 +1566,15 @@ struct BridgedTestContext {
   swift::SwiftPassInvocation *_Nonnull invocation;
 };
 
-typedef void (*_Nonnull BridgedFunctionTestThunk)(BridgedFunction,
-                                                  BridgedTestArguments,
-                                                  BridgedTestContext);
-void registerFunctionTest(llvm::StringRef, BridgedFunctionTestThunk thunk);
+using SwiftNativeFunctionTestThunk = void (*_Nonnull)(void *_Nonnull,
+                                                      BridgedFunction,
+                                                      BridgedTestArguments,
+                                                      BridgedTestContext);
+
+void registerFunctionTestThunk(SwiftNativeFunctionTestThunk);
+
+void registerFunctionTest(llvm::StringRef,
+                          void *_Nonnull nativeSwiftInvocation);
 
 SWIFT_END_NULLABILITY_ANNOTATIONS
 

include/swift/SIL/Test.h

@@ -130,12 +130,11 @@ class FunctionTest final {
   ///                    values such as the results of analyses
   using Invocation = void (*)(SILFunction &, Arguments &, FunctionTest &);
 
-  using InvocationWithContext = void (*)(SILFunction &, Arguments &,
-                                         FunctionTest &, void *);
+  using NativeSwiftInvocation = void *;
 
 private:
   /// The lambda to be run.
-  TaggedUnion<Invocation, std::pair<InvocationWithContext, void *>> invocation;
+  TaggedUnion<Invocation, NativeSwiftInvocation> invocation;
 
 public:
   /// Creates a test that will run \p invocation and stores it in the global
@@ -151,7 +150,7 @@ class FunctionTest final {
   ///     } // end namespace swift::test
   FunctionTest(StringRef name, Invocation invocation);
 
-  FunctionTest(StringRef name, void *context, InvocationWithContext invocation);
+  FunctionTest(StringRef name, NativeSwiftInvocation invocation);
 
   /// Computes and returns the function's dominance tree.
   DominanceInfo *getDominanceInfo();

lib/SIL/Utils/SILBridging.cpp

@@ -133,14 +133,8 @@ void registerBridgedClass(StringRef className, SwiftMetatype metatype) {
 //                                Test
 //===----------------------------------------------------------------------===//
 
-void registerFunctionTest(llvm::StringRef name,
-                          BridgedFunctionTestThunk thunk) {
-  new swift::test::FunctionTest(
-      name, reinterpret_cast<void *>(thunk),
-      [](auto &function, auto &args, auto &test, void *ctx) {
-        auto thunk = reinterpret_cast<BridgedFunctionTestThunk>(ctx);
-        thunk({&function}, {&args}, test.getContext());
-      });
+void registerFunctionTest(llvm::StringRef name, void *nativeSwiftInvocation) {
+  new swift::test::FunctionTest(name, nativeSwiftInvocation);
 }
 
 bool BridgedTestArguments::hasUntaken() const {

lib/SIL/Utils/Test.cpp

@@ -26,6 +26,7 @@ namespace {
 
 class Registry {
   DenseMap<StringRef, FunctionTest *> registeredTests;
+  SwiftNativeFunctionTestThunk thunk;
 
 public:
   static Registry &get() {
@@ -39,6 +40,12 @@ class Registry {
     (void)inserted;
   }
 
+  void registerFunctionTestThunk(SwiftNativeFunctionTestThunk thunk) {
+    this->thunk = thunk;
+  }
+
+  SwiftNativeFunctionTestThunk getFunctionTestThunk() { return thunk; }
+
   FunctionTest *getFunctionTest(StringRef name) {
     auto *res = registeredTests[name];
     if (!res) {
@@ -62,15 +69,18 @@ class Registry {
 
 } // end anonymous namespace
 
+void registerFunctionTestThunk(SwiftNativeFunctionTestThunk thunk) {
+  Registry::get().registerFunctionTestThunk(thunk);
+}
+
 FunctionTest::FunctionTest(StringRef name, Invocation invocation)
     : invocation(invocation), pass(nullptr), function(nullptr),
       dependencies(nullptr) {
   Registry::get().registerFunctionTest(this, name);
 }
-FunctionTest::FunctionTest(StringRef name, void *context,
-                           InvocationWithContext invocation)
-    : invocation(std::make_pair(invocation, context)), pass(nullptr),
-      function(nullptr), dependencies(nullptr) {
+FunctionTest::FunctionTest(StringRef name, NativeSwiftInvocation invocation)
+    : invocation(invocation), pass(nullptr), function(nullptr),
+      dependencies(nullptr) {
   Registry::get().registerFunctionTest(this, name);
 }
 
@@ -84,11 +94,12 @@ void FunctionTest::run(SILFunction &function, Arguments &arguments,
   this->function = &function;
   this->dependencies = &dependencies;
   if (invocation.isa<Invocation>()) {
-    auto fn = this->invocation.get<Invocation>();
+    auto fn = invocation.get<Invocation>();
     fn(function, arguments, *this);
   } else {
-    auto pair = invocation.get<std::pair<InvocationWithContext, void *>>();
-    pair.first(function, arguments, *this, pair.second);
+    auto *fn = invocation.get<NativeSwiftInvocation>();
+    Registry::get().getFunctionTestThunk()(fn, {&function}, {&arguments},
+                                           getContext());
   }
   this->pass = nullptr;
   this->function = nullptr;