When ptrauth-copying vtable/wtables, allow NULL entries (due to VFE) (swiftlang#40578)

* When ptrauth-copying vtable/wtables, allow NULL entries (due to VFE)

* Mark virtual-function-elimination-generics-exec.swift UNSUPPORTED: arm64e until the rebranch

* Fix test expectations

Author: kubamracek

include/swift/Runtime/Config.h

@@ -330,9 +330,14 @@ extern uintptr_t __COMPATIBILITY_LIBRARIES_CANNOT_CHECK_THE_IS_SWIFT_BIT_DIRECTL
 
 /// Copy an address-discriminated signed pointer from the source to the dest.
 template <class T>
-SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE
-static inline void swift_ptrauth_copy(T *dest, const T *src, unsigned extra) {
+SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE static inline void
+swift_ptrauth_copy(T *dest, const T *src, unsigned extra, bool allowNull) {
 #if SWIFT_PTRAUTH
+  if (allowNull && *src == nullptr) {
+    *dest = nullptr;
+    return;
+  }
+
   *dest = ptrauth_auth_and_resign(*src,
                                   ptrauth_key_function_pointer,
                                   ptrauth_blend_discriminator(src, extra),
@@ -348,8 +353,13 @@ static inline void swift_ptrauth_copy(T *dest, const T *src, unsigned extra) {
 template <class T>
 SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE
 static inline void swift_ptrauth_copy_data(T *dest, const T *src,
-                                           unsigned extra) {
+                                           unsigned extra, bool allowNull) {
 #if SWIFT_PTRAUTH
+  if (allowNull && *src == nullptr) {
+    *dest = nullptr;
+    return;
+  }
+
   *dest = ptrauth_auth_and_resign(*src,
                                   ptrauth_key_process_independent_data,
                                   ptrauth_blend_discriminator(src, extra),
@@ -365,11 +375,11 @@ static inline void swift_ptrauth_copy_data(T *dest, const T *src,
 template <class T>
 SWIFT_RUNTIME_ATTRIBUTE_ALWAYS_INLINE static inline void
 swift_ptrauth_copy_code_or_data(T *dest, const T *src, unsigned extra,
-                                bool isCode) {
+                                bool isCode, bool allowNull) {
   if (isCode) {
-    return swift_ptrauth_copy(dest, src, extra);
+    return swift_ptrauth_copy(dest, src, extra, allowNull);
   } else {
-    return swift_ptrauth_copy_data(dest, src, extra);
+    return swift_ptrauth_copy_data(dest, src, extra, allowNull);
   }
 }
 

stdlib/public/runtime/Metadata.cpp

@@ -440,7 +440,8 @@ static void swift_objc_classCopyFixupHandler(Class oldClass, Class newClass) {
             reinterpret_cast<void **>(&dest[i]),
             reinterpret_cast<void *const *>(&src[i]),
             descriptors[i].Flags.getExtraDiscriminator(),
-            !descriptors[i].Flags.isAsync());
+            !descriptors[i].Flags.isAsync(),
+            /*allowNull*/ false); // Don't allow NULL for Obj-C classes
       }
     }
 
@@ -2686,7 +2687,9 @@ static void copySuperclassMetadataToSubclass(ClassMetadata *theClass,
             reinterpret_cast<void **>(&dest[i]),
             reinterpret_cast<void *const *>(&src[i]),
             descriptors[i].Flags.getExtraDiscriminator(),
-            !descriptors[i].Flags.isAsync());
+            !descriptors[i].Flags.isAsync(),
+            /*allowNull*/ true); // NULL allowed for VFE (methods in the vtable
+                                 // might be proven unused and null'ed)
       }
 #else
       memcpy(dest, src, vtable->VTableSize * sizeof(uintptr_t));
@@ -4746,15 +4749,19 @@ static void copyProtocolWitness(void **dest, void * const *src,
   case ProtocolRequirementFlags::Kind::Setter:
   case ProtocolRequirementFlags::Kind::ReadCoroutine:
   case ProtocolRequirementFlags::Kind::ModifyCoroutine:
-    swift_ptrauth_copy_code_or_data(dest, src,
-                                    reqt.Flags.getExtraDiscriminator(),
-                                    !reqt.Flags.isAsync());
+    swift_ptrauth_copy_code_or_data(
+        dest, src, reqt.Flags.getExtraDiscriminator(), !reqt.Flags.isAsync(),
+        /*allowNull*/ true); // NULL allowed for VFE (methods in the vtable
+                             // might be proven unused and null'ed)
     return;
 
   // FIXME: these should both use ptrauth_key_process_independent_data now.
   case ProtocolRequirementFlags::Kind::AssociatedConformanceAccessFunction:
   case ProtocolRequirementFlags::Kind::AssociatedTypeAccessFunction:
-    swift_ptrauth_copy(dest, src, reqt.Flags.getExtraDiscriminator());
+    swift_ptrauth_copy(
+        dest, src, reqt.Flags.getExtraDiscriminator(),
+        /*allowNull*/ true); // NULL allowed for VFE (methods in the vtable
+                             // might be proven unused and null'ed)
     return;
   }
   swift_unreachable("bad witness kind");

stdlib/public/runtime/MetadataLookup.cpp

@@ -2341,7 +2341,7 @@ void DynamicReplacementDescriptor::enableReplacement() const {
         reinterpret_cast<void **>(&chainRoot->implementationFunction),
         reinterpret_cast<void *const *>(&previous->implementationFunction),
         replacedFunctionKey->getExtraDiscriminator(),
-        !replacedFunctionKey->isAsync());
+        !replacedFunctionKey->isAsync(), /*allowNull*/ false);
   }
 
   // First populate the current replacement's chain entry.
@@ -2352,7 +2352,7 @@ void DynamicReplacementDescriptor::enableReplacement() const {
       reinterpret_cast<void **>(&currentEntry->implementationFunction),
       reinterpret_cast<void *const *>(&chainRoot->implementationFunction),
       replacedFunctionKey->getExtraDiscriminator(),
-      !replacedFunctionKey->isAsync());
+      !replacedFunctionKey->isAsync(), /*allowNull*/ false);
 
   currentEntry->next = chainRoot->next;
 
@@ -2388,7 +2388,7 @@ void DynamicReplacementDescriptor::disableReplacement() const {
       reinterpret_cast<void **>(&previous->implementationFunction),
       reinterpret_cast<void *const *>(&thisEntry->implementationFunction),
       replacedFunctionKey->getExtraDiscriminator(),
-      !replacedFunctionKey->isAsync());
+      !replacedFunctionKey->isAsync(), /*allowNull*/ false);
 }
 
 /// An automatic dymamic replacement entry.

test/IRGen/virtual-function-elimination-generics-exec.swift

@@ -0,0 +1,42 @@
+// Tests that under -enable-llvm-vfe, LLVM GlobalDCE is able to remove unused
+// virtual methods, and that used virtual methods are not removed (by running
+// the program).
+
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -Xfrontend -enable-llvm-vfe %s -Onone -emit-ir -o %t/main.ll
+// RUN: %target-clang %t/main.ll -isysroot %sdk -L%swift_obj_root/lib/swift/%target-sdk-name -flto -o %t/main
+// RUN: %target-run %t/main | %FileCheck %s
+
+// RUN: %llvm-nm --defined-only %t/main | %FileCheck %s --check-prefix=NM
+
+// REQUIRES: executable_test
+
+// FIXME(mracek): More work needed to get this to work on non-Apple platforms.
+// REQUIRES: VENDOR=apple
+
+// For LTO, the linker dlopen()'s the libLTO library, which is a scenario that
+// ASan cannot work in ("Interceptors are not working, AddressSanitizer is
+// loaded too late").
+// REQUIRES: no_asan
+
+// Requires rdar://85970023 which is only fixed in the 20211026 branch of LLVM.
+// UNSUPPORTED: CPU=arm64e
+
+class MyClass {
+  func foo() { print("MyClass.foo") }
+  func bar() { print("MyClass.bar") }
+}
+
+class MyDerivedClass<T>: MyClass {
+  override func foo() { print("MyDerivedClass<T>.foo, T=\(type(of: T.self))") }
+  override func bar() { print("MyDerivedClass<T>.bar, T=\(type(of: T.self))") }
+}
+
+let o: MyClass = MyDerivedClass<Int>()
+o.foo()
+// CHECK: MyDerivedClass<T>.foo, T=Int.Type
+
+// NM-NOT: $s4main14MyDerivedClassC3baryyF
+// NM:     $s4main14MyDerivedClassC3fooyyF
+// NM-NOT: $s4main7MyClassC3baryyF
+// NM:     $s4main7MyClassC3fooyyF