Support witness table elimination in relative protocol witness tables mode

Needs https://reviews.llvm.org/D143204 which as merged here:

llvm-project/pull/6198
llvm-project/pull/6197

rdar://98583748

Author: aschwaighofer

lib/IRGen/GenProto.cpp

@@ -2319,7 +2319,9 @@ static void addWTableTypeMetadata(IRGenModule &IGM,
     auto &fnProtoInfo =
         IGM.getProtocolInfo(conf->getProtocol(), ProtocolInfoKind::Full);
     auto index = fnProtoInfo.getFunctionIndex(member).forProtocolWitnessTable();
-    auto offset = index.getValue() * IGM.getPointerSize().getValue();
+    auto entrySize = IGM.IRGen.Opts.UseRelativeProtocolWitnessTables ?
+      4 : IGM.getPointerSize().getValue();
+    auto offset = index.getValue() * entrySize;
     global->addTypeMetadata(offset, typeIdForMethod(IGM, member));
 
     minOffset = std::min(minOffset, offset);
@@ -3810,9 +3812,11 @@ static llvm::Value *emitWTableSlotLoad(IRGenFunction &IGF, llvm::Value *wtable,
 
     // TODO/FIXME: Using @llvm.type.checked.load loses the "invariant" marker
     // which could mean redundant loads don't get removed.
-    llvm::Value *checkedLoad = IGF.Builder.CreateIntrinsicCall(
-        llvm::Intrinsic::type_checked_load, args);
-    assert(!isRelativeTable && "Not yet implemented");
+    llvm::Value *checkedLoad =
+        isRelativeTable ? IGF.Builder.CreateIntrinsicCall(
+                              llvm::Intrinsic::type_checked_load_relative, args)
+                        : IGF.Builder.CreateIntrinsicCall(
+                              llvm::Intrinsic::type_checked_load, args);
     return IGF.Builder.CreateExtractValue(checkedLoad, 0);
   }
 

test/IRGen/witness-method-elimination-exec-relative.swift

@@ -0,0 +1,38 @@
+// Tests that under -enable-llvm-wme, LLVM GlobalDCE is able to remove unused
+// witness methods, and that used witness methods are not removed (by running
+// the program).
+
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -Xfrontend -enable-llvm-wme %s -Onone -emit-ir -o %t/main.ll -Xfrontend -enable-relative-protocol-witness-tables
+// RUN: %target-clang %t/main.ll -isysroot %sdk -L%swift_obj_root/lib/swift/%target-sdk-name -flto -o %t/main
+// RUN: %target-codesign %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
+
+protocol TheProtocol {
+  func func1_live()
+  func func2_dead()
+}
+
+struct MyStruct : TheProtocol {
+  func func1_live() { print("MyStruct.func1_live") }
+  func func2_dead() { print("MyStruct.func2_dead") }
+}
+
+let x: TheProtocol = MyStruct()
+x.func1_live()
+// CHECK: MyStruct.func1_live
+
+// NM:     $s4main8MyStructV10func1_liveyyF
+// NM-NOT: $s4main8MyStructV10func2_deadyyF

test/IRGen/witness-method-elimination-ir-relative.swift

@@ -0,0 +1,51 @@
+// Tests that under -enable-llvm-wme, IRGen marks wtables and wcall sites with
+// the right attributes and intrinsics.
+
+// RUN: %target-build-swift -Xfrontend -disable-objc-interop -Xfrontend -enable-llvm-wme \
+// RUN:    -Xfrontend -enable-relative-protocol-witness-tables \
+// RUN:    %s -emit-ir -o - | %FileCheck %s --check-prefix=CHECK --check-prefix=CHECK-%target-ptrsize
+
+// REQUIRES: PTRSIZE=64
+
+protocol TheProtocol {
+  func func1_live()
+  func func2_dead()
+}
+
+struct MyStruct : TheProtocol {
+  func func1_live() { print("MyStruct.func1_live") }
+  func func2_dead() { print("MyStruct.func2_dead") }
+}
+
+
+// CHECK: @"$s4main8MyStructVAA11TheProtocolAAWP" =
+// CHECK-SAME: hidden constant [3 x i32]
+// CHECK-SAME:  i32 trunc {{.*}} @"$s4main8MyStructVAA11TheProtocolAAMc"
+// CHECK-SAME:  i32 trunc {{.*}} @"$s4main8MyStructVAA11TheProtocolA2aDP10func1_liveyyFTW"
+// CHECK-SAME:  i32 trunc {{.*}} @"$s4main8MyStructVAA11TheProtocolA2aDP10func2_deadyyFTW"
+// CHECK-SAME: ], align 8, !type !0, !type !1, !vcall_visibility !2, !typed_global_not_for_cfi !3
+
+func test1() {
+  // CHECK: define hidden swiftcc void @"$s4main5test1yyF"()
+  let x: MyStruct = MyStruct()
+  x.func1_live()
+  // CHECK:      call swiftcc void @"$s4main8MyStructVACycfC"()
+  // CHECK-NEXT: call swiftcc void @"$s4main8MyStructV10func1_liveyyF"()
+  // CHECK-NEXT: ret void
+}
+
+func test2() {
+  // CHECK: define hidden swiftcc void @"$s4main5test2yyF"()
+  let x: TheProtocol = MyStruct()
+  x.func1_live()
+  // CHECK: [[TBL:%.*]] = phi i8**
+  // CHECK: [[TBL2:%.*]] = bitcast i8** [[TBL]] to i32*
+  // CHECK: [[ENTRY:%.*]] = getelementptr inbounds i32, i32* [[TBL2]], i32 1
+  // CHECK: [[ENTRY2:%.*]] = bitcast i32* [[ENTRY]] to i8*
+  // CHECK: call { i8*, i1 } @llvm.type.checked.load.relative(i8* [[ENTRY2]], i32 0, metadata !"$s4main11TheProtocolP10func1_liveyyFTq")
+}
+
+// CHECK: !0 = !{i64 4, !"$s4main11TheProtocolP10func1_liveyyFTq"}
+// CHECK: !1 = !{i64 8, !"$s4main11TheProtocolP10func2_deadyyFTq"}
+// CHECK: !2 = !{i64 1, i64 4, i64 12}
+// CHECK: !3 = !{}