Override ObjC's class_getImageName to handle Swift classes

This not only restores the correct behavior for classes with generic
ancestry, but also handles actual generic classes as well. (This is
the function that backs Foundation's Bundle.init(for: AnyClass)
initializer.)

https://bugs.swift.org/browse/SR-1917
rdar://problem/33450609&40367300

stdlib/public/runtime/Metadata.cpp

@@ -24,6 +24,7 @@
 #include "swift/Runtime/ExistentialContainer.h"
 #include "swift/Runtime/HeapObject.h"
 #include "swift/Runtime/Mutex.h"
+#include "swift/Runtime/Once.h"
 #include "swift/Strings.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
@@ -54,6 +55,7 @@
 #endif
 
 #if SWIFT_OBJC_INTEROP
+#include <dlfcn.h>
 #include <objc/runtime.h>
 #endif
 
@@ -1974,6 +1976,45 @@ swift::swift_relocateClassMetadata(ClassMetadata *self,
   return self;
 }
 
+#if SWIFT_OBJC_INTEROP
+
+// FIXME: This is from a later version of <objc/runtime.h>. Once the declaration
+// is available in SDKs, we can remove this typedef.
+typedef BOOL (*objc_hook_getImageName)(
+    Class _Nonnull cls, const char * _Nullable * _Nonnull outImageName);
+
+/// \see customGetImageNameFromClass
+static objc_hook_getImageName defaultGetImageNameFromClass = nullptr;
+
+/// A custom implementation of Objective-C's class_getImageName for Swift
+/// classes, which knows how to handle dynamically-initialized class metadata.
+///
+/// Per the documentation for objc_setHook_getImageName, any non-Swift classes
+/// will still go through the normal implementation of class_getImageName,
+/// which is stored in defaultGetImageNameFromClass.
+static BOOL
+customGetImageNameFromClass(Class _Nonnull objcClass,
+                            const char * _Nullable * _Nonnull outImageName) {
+  auto *classAsMetadata = reinterpret_cast<const ClassMetadata *>(objcClass);
+
+  // Is this a Swift class?
+  if (classAsMetadata->isTypeMetadata() &&
+      !classAsMetadata->isArtificialSubclass()) {
+    const void *descriptor = classAsMetadata->getDescription();
+    assert(descriptor &&
+           "all non-artificial Swift classes should have a descriptor");
+    Dl_info imageInfo = {};
+    if (!dladdr(descriptor, &imageInfo))
+      return NO;
+    *outImageName = imageInfo.dli_fname;
+    return imageInfo.dli_fname != nullptr;
+  }
+
+  // If not, fall back to the default implementation.
+  return defaultGetImageNameFromClass(objcClass, outImageName);
+}
+#endif
+
 /// Initialize the field offset vector for a dependent-layout class, using the
 /// "Universal" layout strategy.
 void
@@ -1982,6 +2023,23 @@ swift::swift_initClassMetadata(ClassMetadata *self,
                                size_t numFields,
                                const TypeLayout * const *fieldTypes,
                                size_t *fieldOffsets) {
+#if SWIFT_OBJC_INTEROP
+  // Register our custom implementation of class_getImageName.
+  static swift_once_t onceToken;
+  swift_once(&onceToken, [](void *unused) {
+    (void)unused;
+    // FIXME: This is from a later version of <objc/runtime.h>. Once the
+    // declaration is available in SDKs, we can access this directly instead of
+    // using dlsym.
+    if (void *setHookPtr = dlsym(RTLD_DEFAULT, "objc_setHook_getImageName")) {
+      auto setHook = reinterpret_cast<
+          void(*)(objc_hook_getImageName _Nonnull,
+                  objc_hook_getImageName _Nullable * _Nonnull)>(setHookPtr);
+      setHook(customGetImageNameFromClass, &defaultGetImageNameFromClass);
+    }
+  }, nullptr);
+#endif
+
   _swift_initializeSuperclass(self);
 
   // Start layout by appending to a standard heap object header.

test/Interpreter/SDK/Inputs/SimpleNSObjectSubclass.swift

@@ -0,0 +1,3 @@
+import Foundation
+
+public class SimpleSubclass: NSObject {}

test/Interpreter/SDK/Inputs/class_getImageName-helper.swift

@@ -0,0 +1,23 @@
+import Foundation
+
+public class SimpleSwiftObject {}
+public class SimpleNSObject: NSObject {
+  @objc public dynamic var observableName: String = ""
+}
+
+public class GenericSwiftObject<T> {}
+public class GenericNSObject<T>: NSObject {}
+
+public class GenericAncestrySwiftObject: GenericSwiftObject<AnyObject> {}
+public class GenericAncestryNSObject: GenericNSObject<AnyObject> {
+  @objc public dynamic var observableName: String = ""
+}
+
+public class ResilientFieldSwiftObject {
+  public var url: URL?
+  public var data: Data?
+}
+public class ResilientFieldNSObject: NSObject {
+  public var url: URL?
+  public var data: Data?
+}

test/Interpreter/SDK/Inputs/class_getImageName-static-helper.h

@@ -0,0 +1,7 @@
+static inline const char *getNameOfClassToFind() {
+  return "SimpleNSObjectSubclass.SimpleSubclass";
+}
+
+static inline const char *getHookName() {
+  return "objc_setHook_getImageName";
+}

test/Interpreter/SDK/class_getImageName-static.swift

@@ -0,0 +1,61 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-library -o %t/libSimpleNSObjectSubclass.dylib %S/Inputs/SimpleNSObjectSubclass.swift
+// RUN: %target-codesign %t/libSimpleNSObjectSubclass.dylib
+
+// RUN: %target-build-swift %s -o %t/main -lSimpleNSObjectSubclass -L%t -import-objc-header %S/Inputs/class_getImageName-static-helper.h
+// RUN: %target-run %t/main %t/libSimpleNSObjectSubclass.dylib
+
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+import Darwin
+import ObjectiveC
+// import SimpleNSObjectSubclass // Deliberately omitted in favor of dynamic loads.
+
+// Note: The following typealias uses AnyObject instead of AnyClass so that the
+// function type is trivially bridgeable to Objective-C. (The representation of
+// AnyClass is not the same as Objective-C's 'Class' type.)
+typealias GetImageHook = @convention(c) (AnyObject, UnsafeMutablePointer<UnsafePointer<CChar>?>) -> ObjCBool
+var hook: GetImageHook?
+
+func checkThatSwiftHookWasNotInstalled() {
+  // Check that the Swift hook did not get installed.
+  guard let setHookPtr = dlsym(UnsafeMutableRawPointer(bitPattern: -2),
+                               getHookName()) else {
+    // If the version of the ObjC runtime we're using doesn't have the hook,
+    // we're good.
+    return
+  }
+
+  let setHook = unsafeBitCast(setHookPtr, to: (@convention(c) (GetImageHook, UnsafeMutablePointer<GetImageHook?>) -> Void).self)
+  setHook({ hook!($0, $1) }, &hook)
+
+  var info: Dl_info = .init()
+  guard 0 != dladdr(unsafeBitCast(hook, to: UnsafeRawPointer.self), &info) else {
+    fatalError("could not get dladdr info for objc_hook_getImageName")
+  }
+
+  precondition(String(cString: info.dli_fname).hasSuffix("libobjc.A.dylib"),
+               "hook was replaced")
+}
+
+// It's important that this test does not register any Swift classes with the
+// Objective-C runtime---that's where Swift sets up its custom hook, and we want
+// to check the behavior /without/ that hook. That includes the buffer types for
+// String and Array. Therefore, we get C strings directly from a bridging
+// header.
+
+guard let theClass = objc_getClass(getNameOfClassToFind()) as! AnyClass? else {
+  fatalError("could not find class")
+}
+
+guard let imageName = class_getImageName(theClass) else {
+  fatalError("could not find image")
+}
+
+checkThatSwiftHookWasNotInstalled()
+
+// Okay, now we can use String.
+
+precondition(String(cString: imageName).hasSuffix("libSimpleNSObjectSubclass.dylib"),
+             "found wrong image")

test/Interpreter/SDK/class_getImageName.swift

@@ -0,0 +1,119 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-library -o %t/libGetImageNameHelper.dylib -emit-module %S/Inputs/class_getImageName-helper.swift
+// RUN: %target-codesign %t/libGetImageNameHelper.dylib
+
+// RUN: %target-build-swift -g %s -I %t -o %t/main -L %t -lGetImageNameHelper
+// RUN: %target-run %t/main %t/libGetImageNameHelper.dylib
+
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+import Darwin
+import ObjectiveC
+import GetImageNameHelper
+import StdlibUnittest
+
+func check(_ cls: AnyClass, in library: String) {
+  guard let imageName = class_getImageName(cls) else {
+    expectUnreachable("could not find image for \(cls)")
+    return
+  }
+  expectTrue(String(cString: imageName).hasSuffix(library),
+             "wrong library for \(cls)")
+}
+
+let isMissingObjCRuntimeHook =
+    (nil == dlsym(UnsafeMutableRawPointer(bitPattern: -2),
+                  "objc_setHook_getImageName"))
+
+var testSuite = TestSuite("class_getImageName")
+
+testSuite.test("Simple") {
+  check(SimpleSwiftObject.self, in: "libGetImageNameHelper.dylib")
+  check(SimpleNSObject.self, in: "libGetImageNameHelper.dylib")
+}
+
+testSuite.test("Generic")
+    .xfail(.custom({ isMissingObjCRuntimeHook },
+                   reason: "hook for class_getImageName not present"))
+    .code {
+  check(GenericSwiftObject<Int>.self, in: "libGetImageNameHelper.dylib")
+  check(GenericSwiftObject<NSObject>.self, in: "libGetImageNameHelper.dylib")
+
+  check(GenericNSObject<Int>.self, in: "libGetImageNameHelper.dylib")
+  check(GenericNSObject<NSObject>.self, in: "libGetImageNameHelper.dylib")
+}
+
+testSuite.test("GenericAncestry")
+    .xfail(.custom({ isMissingObjCRuntimeHook },
+                   reason: "hook for class_getImageName not present"))
+    .code {
+  check(GenericAncestrySwiftObject.self, in: "libGetImageNameHelper.dylib")
+  check(GenericAncestryNSObject.self, in: "libGetImageNameHelper.dylib")
+}
+
+testSuite.test("Resilient") {
+  check(ResilientFieldSwiftObject.self, in: "libGetImageNameHelper.dylib")
+  check(ResilientFieldNSObject.self, in: "libGetImageNameHelper.dylib")
+}
+
+testSuite.test("ObjC") {
+  check(NSObject.self, in: "libobjc.A.dylib")
+}
+
+testSuite.test("KVO/Simple") {
+  // We use object_getClass in this test to not look through KVO's artificial
+  // subclass.
+  let obj = SimpleNSObject()
+  let observation = obj.observe(\.observableName) { _, _ in }
+  withExtendedLifetime(observation) {
+    let theClass = object_getClass(obj)
+    precondition(theClass !== SimpleNSObject.self, "no KVO subclass?")
+    expectNil(class_getImageName(theClass),
+              "should match what happens with NSObject (below)")
+  }
+}
+
+testSuite.test("KVO/GenericAncestry") {
+  // We use object_getClass in this test to not look through KVO's artificial
+  // subclass.
+  let obj = GenericAncestryNSObject()
+  let observation = obj.observe(\.observableName) { _, _ in }
+  withExtendedLifetime(observation) {
+    let theClass = object_getClass(obj)
+    precondition(theClass !== GenericAncestryNSObject.self, "no KVO subclass?")
+    expectNil(class_getImageName(theClass),
+              "should match what happens with NSObject (below)")
+  }
+}
+
+testSuite.test("KVO/ObjC") {
+  // We use object_getClass in this test to not look through KVO's artificial
+  // subclass.
+  let obj = NSObject()
+  let observation = obj.observe(\.description) { _, _ in }
+  withExtendedLifetime(observation) {
+    let theClass = object_getClass(obj)
+    precondition(theClass !== NSObject.self, "no KVO subclass?")
+    expectNil(class_getImageName(theClass),
+              "should match what happens with the Swift objects (above)")
+  }
+}
+
+testSuite.test("dynamic") {
+  let newClass: AnyClass = objc_allocateClassPair(/*superclass*/nil,
+                                                  "CompletelyDynamic",
+                                                  /*extraBytes*/0)!
+  objc_registerClassPair(newClass)
+
+  // We don't actually care what the result is; we just need to not crash.
+  _ = class_getImageName(newClass)
+}
+
+testSuite.test("nil") {
+  // The ObjC runtime should handle this before it even gets to Swift's custom
+  // implementation, but just in case.
+  expectNil(class_getImageName(nil))
+}
+
+runAllTests()