[IRGen] Properly handle ObjC class existentials in CVW generation

lib/IRGen/TypeLayout.cpp

@@ -20,6 +20,7 @@
 #include "IRGen.h"
 #include "IRGenFunction.h"
 #include "IRGenModule.h"
+#include "ReferenceTypeInfo.h"
 #include "SwitchBuilder.h"
 #include "swift/ABI/MetadataValues.h"
 #include "swift/AST/GenericEnvironment.h"
@@ -2408,7 +2409,8 @@ bool EnumTypeLayoutEntry::refCountString(IRGenModule &IGM,
     if (isMultiPayloadEnum() &&
         buildMultiPayloadRefCountString(IGM, B, genericSig)) {
       return true;
-    } else if (buildSinglePayloadRefCountString(IGM, B, genericSig)) {
+    } else if (!isMultiPayloadEnum() &&
+               buildSinglePayloadRefCountString(IGM, B, genericSig)) {
       return true;
     }
 
@@ -3829,6 +3831,13 @@ TypeInfoBasedTypeLayoutEntry::layoutString(IRGenModule &IGM,
 bool TypeInfoBasedTypeLayoutEntry::refCountString(
     IRGenModule &IGM, LayoutStringBuilder &B,
     GenericSignature genericSig) const {
+  if (auto *refTI = dyn_cast<ReferenceTypeInfo>(&typeInfo)) {
+    if (refTI->getReferenceCountingType() == ReferenceCounting::ObjC) {
+      B.addRefCount(LayoutStringBuilder::RefCountingKind::ObjC,
+                    typeInfo.getFixedSize().getValue());
+      return true;
+    }
+  }
   return false;
 }
 

test/IRGen/layout_string_witnesses_objc_existentials.swift

@@ -0,0 +1,30 @@
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -enable-experimental-feature LayoutStringValueWitnesses -enable-experimental-feature LayoutStringValueWitnessesInstantiation -enable-layout-string-value-witnesses -enable-layout-string-value-witnesses-instantiation -emit-ir -import-objc-header %S/Inputs/ObjCBaseClasses.h %s | %FileCheck %s
+
+// REQUIRES: PTRSIZE=64
+// REQUIRES: objc_interop
+
+protocol P {}
+
+extension ObjCBase: P {}
+
+// CHECK: @"type_layout_string 41layout_string_witnesses_objc_existentials27MultiPayloadObjCExistentialO" =
+// CHECK: internal constant <{ i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64 }>
+// CHECK: <{ i64 -9223372036854775808, i64 88, i64 1441151880758558720, i64 sub (i64 ptrtoint (ptr @"get_enum_tag_for_layout_string 41layout_string_witnesses_objc_existentials27MultiPayloadObjCExistentialO" to i64),
+// CHECK: i64 ptrtoint (ptr getelementptr inbounds (<{ i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64 }>,
+// CHECK: ptr @"type_layout_string 41layout_string_witnesses_objc_existentials27MultiPayloadObjCExistentialO", i32 0, i32 3) to i64)),
+// CHECK: i64 2, i64 32, i64 16, i64 0, i64 16, i64 360287970189639680, i64 0, i64 720575940379279360, i64 0, i64 0 }>
+enum MultiPayloadObjCExistential {
+    case x(AnyObject)
+    case y(P & ObjCBase)
+}
+
+// CHECK: @"type_layout_string 41layout_string_witnesses_objc_existentials34MultiPayloadObjCExistentialWrapperV" =
+// CHECK: internal constant <{ i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64 }>
+// CHECK: <{ i64 -9223372036854775808, i64 88, i64 1441151880758558720, i64 sub (i64 ptrtoint (ptr @"get_enum_tag_for_layout_string 41layout_string_witnesses_objc_existentials27MultiPayloadObjCExistentialO" to i64),
+// CHECK: i64 ptrtoint (ptr getelementptr inbounds (<{ i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64 }>,
+// CHECK: ptr @"type_layout_string 41layout_string_witnesses_objc_existentials34MultiPayloadObjCExistentialWrapperV", i32 0, i32 3) to i64)),
+// CHECK: i64 2, i64 32, i64 16, i64 0, i64 16, i64 360287970189639680, i64 0, i64 720575940379279360, i64 0, i64 8 }>
+struct MultiPayloadObjCExistentialWrapper {
+    let x: MultiPayloadObjCExistential
+    let y: Int = 0
+}

test/Interpreter/layout_string_witnesses_objc.swift

@@ -48,3 +48,45 @@ func testNestedResilientObjc() {
 }
 
 testNestedResilientObjc()
+
+protocol P {}
+
+extension ObjCPrintOnDealloc: P {}
+
+enum MultiPayloadObjCExistential {
+    case x(AnyObject)
+    case y(P & ObjCPrintOnDealloc)
+}
+
+struct MultiPayloadObjCExistentialWrapper {
+    let x: MultiPayloadObjCExistential
+    let y: Int = 0
+}
+
+func testMultiPayloadObjCExistentialWrapper() {
+    let ptr = allocateInternalGenericPtr(of: NestedWrapper<MultiPayloadObjCExistentialWrapper>.self)
+
+    do {
+        let x = MultiPayloadObjCExistentialWrapper(x: .y(ObjCPrintOnDealloc()))
+        testGenericInit(ptr, to: x)
+    }
+
+    do {
+        let y = MultiPayloadObjCExistentialWrapper(x: .y(ObjCPrintOnDealloc()))
+        // CHECK: Before deinit
+        print("Before deinit")
+
+        // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+        testGenericAssign(ptr, from: y)
+    }
+
+    // CHECK-NEXT: Before deinit
+    print("Before deinit")
+
+    // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+    testGenericDestroy(ptr, of: NestedWrapper<MultiPayloadObjCExistentialWrapper>.self)
+
+    ptr.deallocate()
+}
+
+testMultiPayloadObjCExistentialWrapper()