From 3301541be5b1906fff58776c7ed9c3c49438333f Mon Sep 17 00:00:00 2001
From: Dario Rexin <drexin@apple.com>
Date: Mon, 11 Nov 2024 16:19:58 -0800
Subject: [PATCH] [IRGen] Fix crash in CVW generation for ObjC references

rdar://139664644

The code that differentiates between regular ObjC and native Swift ObjC references could crash when generics were involved. Instead of through the TypeInfo, we are going directly throught the SILType to the type decl, which avoids the crash caused by casting the TypeInfo.
---
 lib/IRGen/TypeLayout.cpp                      |  14 +--
 .../layout_string_witnesses_objc.swift        | 100 ++++++++++++++++++
 2 files changed, 107 insertions(+), 7 deletions(-)

diff --git a/lib/IRGen/TypeLayout.cpp b/lib/IRGen/TypeLayout.cpp
index f7660b8833fa5..b5201c123d5f3 100644
--- a/lib/IRGen/TypeLayout.cpp
+++ b/lib/IRGen/TypeLayout.cpp
@@ -1295,14 +1295,14 @@ bool ScalarTypeLayoutEntry::refCountString(IRGenModule &IGM,
     B.addRefCount(LayoutStringBuilder::RefCountingKind::Block, size);
     break;
   case ScalarKind::ObjCReference: {
-    auto *classTI = dyn_cast<ClassTypeInfo>(&typeInfo);
-    assert(classTI);
-    if (!classTI->getClass()->hasClangNode()) {
-      B.addRefCount(LayoutStringBuilder::RefCountingKind::NativeSwiftObjC,
-                    size);
-    } else {
-      B.addRefCount(LayoutStringBuilder::RefCountingKind::ObjC, size);
+    if (auto *classDecl = representative.getClassOrBoundGenericClass()) {
+      if (!classDecl->hasClangNode()) {
+        B.addRefCount(LayoutStringBuilder::RefCountingKind::NativeSwiftObjC,
+                      size);
+        break;
+      }
     }
+    B.addRefCount(LayoutStringBuilder::RefCountingKind::ObjC, size);
     break;
   }
   case ScalarKind::ThickFunc:
diff --git a/test/Interpreter/layout_string_witnesses_objc.swift b/test/Interpreter/layout_string_witnesses_objc.swift
index 04870c97bd28a..5f49b7169e40a 100644
--- a/test/Interpreter/layout_string_witnesses_objc.swift
+++ b/test/Interpreter/layout_string_witnesses_objc.swift
@@ -190,3 +190,103 @@ func testMultiPayloadBlock() {
 }
 
 testMultiPayloadBlock()
+
+class GenericOuterClassNSObject<T: NSObject> {
+    enum InnerEnum {
+        case x(T.Type)
+        case y(T)
+    }
+}
+
+func testNestedGenericEnumNSObject() {
+    let ptr = UnsafeMutablePointer<GenericOuterClassNSObject<ObjCPrintOnDealloc>.InnerEnum>.allocate(capacity: 1)
+
+    // initWithCopy
+    do {
+        let x = GenericOuterClassNSObject<ObjCPrintOnDealloc>.InnerEnum.y(ObjCPrintOnDealloc())
+        testInit(ptr, to: x)
+    }
+
+    // assignWithTake
+    do {
+        let y = GenericOuterClassNSObject<ObjCPrintOnDealloc>.InnerEnum.y(ObjCPrintOnDealloc())
+
+        // CHECK-NEXT: Before deinit
+        print("Before deinit")
+
+        // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+        testAssign(ptr, from: y)
+    }
+
+    // assignWithCopy
+    do {
+        var z = GenericOuterClassNSObject<ObjCPrintOnDealloc>.InnerEnum.y(ObjCPrintOnDealloc())
+
+        // CHECK-NEXT: Before deinit
+        print("Before deinit")
+
+        // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+        testAssignCopy(ptr, from: &z)
+    }
+
+    // CHECK-NEXT: Before deinit
+    print("Before deinit")
+
+    // destroy
+    // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+    testDestroy(ptr)
+
+    ptr.deallocate()
+}
+
+testNestedGenericEnumNSObject()
+
+class GenericOuterClassSwiftObjC<T: SwiftObjC> {
+    enum InnerEnum {
+        case x(T.Type)
+        case y(T)
+    }
+}
+
+func testNestedGenericEnumSwiftObjC() {
+    let ptr = UnsafeMutablePointer<GenericOuterClassSwiftObjC<SwiftObjC>.InnerEnum>.allocate(capacity: 1)
+
+    // initWithCopy
+    do {
+        let x = GenericOuterClassSwiftObjC<SwiftObjC>.InnerEnum.y(SwiftObjC())
+        testInit(ptr, to: x)
+    }
+
+    // assignWithTake
+    do {
+        let y = GenericOuterClassSwiftObjC<SwiftObjC>.InnerEnum.y(SwiftObjC())
+
+        // CHECK-NEXT: Before deinit
+        print("Before deinit")
+
+        // CHECK-NEXT: SwiftObjC deinitialized!
+        testAssign(ptr, from: y)
+    }
+
+    // assignWithCopy
+    do {
+        var z = GenericOuterClassSwiftObjC<SwiftObjC>.InnerEnum.y(SwiftObjC())
+
+        // CHECK-NEXT: Before deinit
+        print("Before deinit")
+
+        // CHECK-NEXT: SwiftObjC deinitialized!
+        testAssignCopy(ptr, from: &z)
+    }
+
+    // CHECK-NEXT: Before deinit
+    print("Before deinit")
+
+    // destroy
+    // CHECK-NEXT: SwiftObjC deinitialized!
+    testDestroy(ptr)
+
+    ptr.deallocate()
+}
+
+testNestedGenericEnumSwiftObjC()