IRGen: Conditionally test weakly linked enum cases

rdar://46438568

Author: aschwaighofer

lib/IRGen/GenEnum.cpp

@@ -5412,13 +5412,50 @@ namespace {
                                        enumAddr);
     }
 
+    llvm::Value *testResilientTag(IRGenFunction &IGF, llvm::Value *tag,
+                                  EnumElementDecl *Case) const {
+      auto &C = IGM.getLLVMContext();
+
+      // If the enum case is weakly linked check the address of the case
+      // first.
+      llvm::BasicBlock *conditionalBlock = nullptr;
+      llvm::BasicBlock *afterConditionalBlock = nullptr;
+      llvm::BasicBlock *beforeNullPtrCheck = nullptr;
+      if (Case->isWeakImported(IGM.getSwiftModule())) {
+        beforeNullPtrCheck = IGF.Builder.GetInsertBlock();
+        auto address = IGM.getAddrOfEnumCase(Case, NotForDefinition);
+        conditionalBlock = llvm::BasicBlock::Create(C);
+        afterConditionalBlock = llvm::BasicBlock::Create(C);
+        auto *addressVal =
+            IGF.Builder.CreatePtrToInt(address.getAddress(), IGM.IntPtrTy);
+        auto isNullPtr = IGF.Builder.CreateICmpEQ(
+            addressVal, llvm::ConstantInt::get(IGM.IntPtrTy, 0));
+        IGF.Builder.CreateCondBr(isNullPtr, afterConditionalBlock,
+                                 conditionalBlock);
+      }
+      if (conditionalBlock)
+        IGF.Builder.emitBlock(conditionalBlock);
+
+      // Check the tag.
+      auto tagVal = loadResilientTagIndex(IGF, Case);
+      auto matchesTag = IGF.Builder.CreateICmpEQ(tag, tagVal);
+      if (conditionalBlock) {
+        IGF.Builder.CreateBr(afterConditionalBlock);
+        IGF.Builder.emitBlock(afterConditionalBlock);
+        auto phi = IGF.Builder.CreatePHI(IGM.Int1Ty, 2);
+        phi->addIncoming(IGF.Builder.getInt1(false), beforeNullPtrCheck);
+        phi->addIncoming(matchesTag, conditionalBlock);
+        matchesTag = phi;
+      }
+      return matchesTag;
+    }
+
     llvm::Value *
     emitIndirectCaseTest(IRGenFunction &IGF, SILType T,
                          Address enumAddr,
                          EnumElementDecl *Case) const override {
       llvm::Value *tag = emitGetEnumTagCall(IGF, T, enumAddr);
-      llvm::Value *expectedTag = loadResilientTagIndex(IGF, Case);
-      return IGF.Builder.CreateICmpEQ(tag, expectedTag);
+      return testResilientTag(IGF, tag, Case);
     }
 
     void emitIndirectSwitch(IRGenFunction &IGF,
@@ -5453,8 +5490,8 @@ namespace {
           if (continuationBB)
             IGF.Builder.emitBlock(continuationBB);
 
-          auto tagVal = loadResilientTagIndex(IGF, elt.decl);
-          auto matchesTag = IGF.Builder.CreateICmpEQ(tag, tagVal);
+          // Check the tag.
+          auto matchesTag = testResilientTag(IGF, tag, elt.decl);
 
           // If we are not the last block create a continuation block.
           if (++numCasesEmitted < dests.size())

test/IRGen/weak_import_native.swift

@@ -230,3 +230,30 @@ public func test_not_hoist_weakly_linked4() {
      var _ = One(elt:(ResilientStruct(), 1))
   }
 }
+
+// CHECK-DAG-LABEL: define{{.*}} @"$s18weak_import_native29test_weakly_linked_enum_cases1eSi0a1_b1_C7_helper1EO_t
+// CHECK:  [[TAG:%.*]] = call i32 %getEnumTag(
+// CHECK:  [[STRONG_CASE:%.*]] = load i32, i32* @"$s25weak_import_native_helper1EO6strongyA2CmFWC"
+// CHECK:  [[IS_STRONG:%.*]] = icmp eq i32 [[TAG]], [[STRONG_CASE]]
+// CHECK:  br i1 [[IS_STRONG]], label %[[BB0:[0-9]+]], label %[[BB1:[0-9]+]]
+//
+// CHECK:  <label>:[[BB1]]:
+// CHECK:  br i1 icmp eq ({{.*}} ptrtoint (i32* @"$s25weak_import_native_helper1EO0A0yA2CmFWC" to {{.*}}), {{.*}} 0), label %[[BB2:[0-9]+]], label %[[BB3:[0-9]+]]
+//
+// CHECK:; <label>:[[BB3]]:
+// CHECK:  [[WEAK_CASE:%.*]] = load i32, i32* @"$s25weak_import_native_helper1EO0A0yA2CmFWC"
+// CHECK:  [[IS_WEAK:%.*]] = icmp eq i32 [[TAG]], [[WEAK_CASE]]
+// CHECK:  br label %21
+//
+// CHECK:; <label>:[[BB2]]:
+// CHECK:  = phi i1 [ false, %[[BB1]] ], [ [[IS_WEAK]], %[[BB3]] ]
+public func test_weakly_linked_enum_cases(e: E) -> Int {
+  switch e {
+    case .strong:
+      return 1
+    case .weak:
+      return 2
+    default:
+      return 3
+  }
+}

validation-test/Evolution/test_backward_deploy_enum.swift

@@ -4,9 +4,6 @@
 import StdlibUnittest
 import backward_deploy_enum
 
-// <rdar://problem/46438568>
-// REQUIRES: rdar46438568
-
 var BackwardDeployEnumTest = TestSuite("BackwardDeployEnum")
 
 func checkIt(_ e: ResilientEnum) -> Int {