Override ObjC's class_getImageName to handle Swift classes

lib/Driver/DarwinToolChains.cpp

@@ -198,6 +198,26 @@ static void addVersionString(const ArgList &inputArgs, ArgStringList &arguments,
   arguments.push_back(inputArgs.MakeArgString(os.str()));
 }
 
+/// Returns true if the compiler depends on features provided by the ObjC
+/// runtime that are not present on the deployment target indicated by
+/// \p triple.
+static bool wantsObjCRuntime(const llvm::Triple &triple) {
+  assert((!triple.isTvOS() || triple.isiOS()) &&
+         "tvOS is considered a kind of iOS");
+
+  // When updating the versions listed here, please record the most recent
+  // feature being depended on and when it was introduced:
+  //
+  // - The hook to override class_getImageName (macOS 10.14 and equivalent)
+  if (triple.isiOS())
+    return triple.isOSVersionLT(12);
+  if (triple.isMacOSX())
+    return triple.isMacOSXVersionLT(10, 14);
+  if (triple.isWatchOS())
+    return triple.isOSVersionLT(5);
+  llvm_unreachable("unknown Darwin OS");
+}
+
 ToolChain::InvocationInfo
 toolchains::Darwin::constructInvocation(const LinkJobAction &job,
                                         const JobContext &context) const {
@@ -282,15 +302,9 @@ toolchains::Darwin::constructInvocation(const LinkJobAction &job,
   if (llvm::sys::fs::exists(CompilerRTPath))
     Arguments.push_back(context.Args.MakeArgString(CompilerRTPath));
 
-  bool wantsObjCRuntime = false;
-  if (Triple.isiOS())
-    wantsObjCRuntime = Triple.isOSVersionLT(9);
-  else if (Triple.isMacOSX())
-    wantsObjCRuntime = Triple.isMacOSXVersionLT(10, 11);
-
   if (context.Args.hasFlag(options::OPT_link_objc_runtime,
                            options::OPT_no_link_objc_runtime,
-                           /*Default=*/wantsObjCRuntime)) {
+                           /*Default=*/wantsObjCRuntime(Triple))) {
     llvm::SmallString<128> ARCLiteLib(D.getSwiftProgramPath());
     llvm::sys::path::remove_filename(ARCLiteLib); // 'swift'
     llvm::sys::path::remove_filename(ARCLiteLib); // 'bin'

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/Driver/linker-arclite.swift

@@ -11,15 +11,24 @@
 
 // RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-ios8.0 %S/../Inputs/empty.swift | %FileCheck -check-prefix IOS_ARCLITE %s
 
-// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-macosx10.11 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-ios9.0 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target arm64-apple-tvos9.0 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
-// RUN: %swiftc_driver -driver-print-jobs -target armv7k-apple-watchos2.0 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
-
 // IOS_ARCLITE: bin/ld{{"? }}
 // IOS_ARCLITE: -force_load {{[^ ]+/lib/arc/libarclite_iphonesimulator.a}}
 // IOS_ARCLITE: -o {{[^ ]+}}
 
+
+// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-macosx10.14 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-macosx10.13 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-ios12 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target x86_64-apple-ios11 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target arm64-apple-tvos12 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target arm64-apple-tvos11 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target armv7k-apple-watchos5 %S/../Inputs/empty.swift | %FileCheck -check-prefix NO_ARCLITE %s
+// RUN: %swiftc_driver -driver-print-jobs -target armv7k-apple-watchos4 %S/../Inputs/empty.swift | %FileCheck -check-prefix ANY_ARCLITE %s
+
 // NO_ARCLITE: bin/ld{{"? }}
 // NO_ARCLITE-NOT: arclite
 // NO_ARCLITE: -o {{[^ ]+}}
+
+// ANY_ARCLITE: bin/ld{{"? }}
+// ANY_ARCLITE: -force_load {{[^ ]+}}/lib/arc/libarclite_{{.+}}.a
+// ANY_ARCLITE: -o {{[^ ]+}}

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()