[Runtime] Fix _swift_refCountBytesForMetatype for reference types

_swift_addRefCountStringForMetatype and _swift_refCountBytesForMetatype diverged in the code that determines whether a type is a reference, causing the size number of ref count bytes to differ from the actually used bytes. This can cause early termination of the runtime interpreter functions, which in turn causes unbalanced reference counts.

rdar://112474091

Author: drexin

stdlib/public/runtime/Metadata.cpp

@@ -2734,7 +2734,18 @@ void swift::swift_initStructMetadataWithLayoutString(
     } else if (fieldType->hasLayoutString()) {
       refCountBytes += *(const size_t *)(fieldType->getLayoutString() +
                                          sizeof(uint64_t));
-    } else if (fieldType->isClassObject() || fieldType->isAnyExistentialType()) {
+    } else if (auto *cls = fieldType->getClassObject()) {
+      if (cls->isTypeMetadata()) {
+        auto *vwt = cls->getValueWitnesses();
+        if (vwt != &VALUE_WITNESS_SYM(Bo) &&
+            vwt != &VALUE_WITNESS_SYM(BO) &&
+            vwt != &VALUE_WITNESS_SYM(Bb)) {
+          refCountBytes += sizeof(uint64_t) + sizeof(uintptr_t);
+          continue;
+        }
+      }
+      refCountBytes += sizeof(uint64_t);
+    } else if (fieldType->isAnyExistentialType()) {
       refCountBytes += sizeof(uint64_t);
     } else {
       refCountBytes += sizeof(uint64_t) + sizeof(uintptr_t);

test/Interpreter/Inputs/ObjCClasses/ObjCClasses.h

@@ -118,6 +118,9 @@ __attribute__((swift_name("OuterType.InnerType")))
 @property NSArray<OuterType *> *things;
 @end
 
+@interface ObjCPrintOnDealloc : NSObject
+@end
+
 NS_ASSUME_NONNULL_END
 
 #endif

test/Interpreter/Inputs/ObjCClasses/ObjCClasses.m

@@ -191,3 +191,9 @@ - (id)init {
 }
 
 @end
+
+@implementation ObjCPrintOnDealloc
+- (void)dealloc {
+  printf("ObjCPrintOnDealloc deinitialized!\n");
+}
+@end

test/Interpreter/Inputs/layout_string_witnesses_types.swift

@@ -271,6 +271,16 @@ public struct Wrapper<T> {
     }
 }
 
+public struct NestedWrapper<T> {
+    public let x: Wrapper<T>
+    public let y: Wrapper<T>
+
+    public init(x: Wrapper<T>, y: Wrapper<T>) {
+        self.x = x
+        self.y = y
+    }
+}
+
 struct InternalGeneric<T> {
     let x: T
     let y: Int

test/Interpreter/layout_string_witnesses_objc.swift

@@ -0,0 +1,46 @@
+// RUN: %empty-directory(%t)
+//
+// RUN: %target-clang -fobjc-arc %S/Inputs/ObjCClasses/ObjCClasses.m -c -o %t/ObjCClasses.o
+// RUN: %target-swift-frontend -prespecialize-generic-metadata -enable-experimental-feature LayoutStringValueWitnesses -enable-experimental-feature LayoutStringValueWitnessesInstantiation -enable-layout-string-value-witnesses -enable-layout-string-value-witnesses-instantiation -enable-type-layout -enable-autolinking-runtime-compatibility-bytecode-layouts -parse-stdlib -emit-module -emit-module-path=%t/layout_string_witnesses_types.swiftmodule %S/Inputs/layout_string_witnesses_types.swift
+// RUN: %target-build-swift-dylib(%t/%target-library-name(layout_string_witnesses_types)) -Xfrontend -enable-experimental-feature -Xfrontend LayoutStringValueWitnesses -Xfrontend -enable-experimental-feature -Xfrontend LayoutStringValueWitnessesInstantiation -Xfrontend -enable-layout-string-value-witnesses -Xfrontend -enable-layout-string-value-witnesses-instantiation -Xfrontend -enable-type-layout -Xfrontend -parse-stdlib -parse-as-library %S/Inputs/layout_string_witnesses_types.swift
+// RUN: %target-build-swift -g -Xfrontend -enable-experimental-feature -Xfrontend LayoutStringValueWitnesses -Xfrontend -enable-experimental-feature -Xfrontend LayoutStringValueWitnessesInstantiation -Xfrontend -enable-layout-string-value-witnesses -Xfrontend -enable-layout-string-value-witnesses-instantiation -Xfrontend -enable-type-layout -parse-stdlib -module-name layout_string_witnesses_dynamic -llayout_string_witnesses_types -L%t -I %S/Inputs/ObjCClasses/ %t/ObjCClasses.o -I %t -o %t/main %s %target-rpath(%t)
+// RUN: %target-codesign %t/main
+// RUN: %target-run %t/main %t/%target-library-name(layout_string_witnesses_types) | %FileCheck %s --check-prefix=CHECK
+
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+import Swift
+import layout_string_witnesses_types
+import ObjCClasses
+import Foundation
+
+func testNestedResilientObjc() {
+    let ptr = allocateInternalGenericPtr(of: NestedWrapper<ObjCPrintOnDealloc>.self)
+
+    do {
+        let x = NestedWrapper<ObjCPrintOnDealloc>(x: .init(x: ObjCPrintOnDealloc()), y: .init(x: ObjCPrintOnDealloc()))
+        testGenericInit(ptr, to: x)
+    }
+
+    do {
+        let y = NestedWrapper<ObjCPrintOnDealloc>(x: .init(x: ObjCPrintOnDealloc()), y: .init(x: ObjCPrintOnDealloc()))
+        // CHECK: Before deinit
+        print("Before deinit")
+
+        // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+        // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+        testGenericAssign(ptr, from: y)
+    }
+
+    // CHECK-NEXT: Before deinit
+    print("Before deinit")
+
+    // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+    // CHECK-NEXT: ObjCPrintOnDealloc deinitialized!
+    testGenericDestroy(ptr, of: NestedWrapper<ObjCPrintOnDealloc>.self)
+
+    ptr.deallocate()
+}
+
+testNestedResilientObjc()