[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

@@ -2785,15 +2785,26 @@ size_t swift::_swift_refCountBytesForMetatype(const Metadata *type) {
     size_t offset = sizeof(uint64_t);
     return LayoutStringReader{type->getLayoutString(), offset}
         .readBytes<size_t>();
-  } else if (type->isClassObject() || type->isAnyExistentialType()) {
-    return sizeof(uint64_t);
   } else if (auto *tuple = dyn_cast<TupleTypeMetadata>(type)) {
     size_t res = 0;
     for (InProcess::StoredSize i = 0; i < tuple->NumElements; i++) {
       res += _swift_refCountBytesForMetatype(tuple->getElement(i).Type);
     }
     return res;
+  } else if (auto *cls = type->getClassObject()) {
+    if (cls->isTypeMetadata()) {
+      auto *vwt = cls->getValueWitnesses();
+      if (vwt != &VALUE_WITNESS_SYM(Bo) &&
+          vwt != &VALUE_WITNESS_SYM(BO) &&
+          vwt != &VALUE_WITNESS_SYM(Bb)) {
+        goto metadata;
+      }
+    }
+    return sizeof(uint64_t);
+  } else if (type->isAnyExistentialType()) {
+    return sizeof(uint64_t);
   } else {
+  metadata:
     return 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

@@ -276,6 +276,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 -check-prefix=CHECK-%target-os
+
+// 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-macosx: Before deinit
+        print("Before deinit")
+
+        // CHECK-macosx-NEXT: ObjCPrintOnDealloc deinitialized!
+        // CHECK-macosx-NEXT: ObjCPrintOnDealloc deinitialized!
+        testGenericAssign(ptr, from: y)
+    }
+
+    // CHECK-macosx-NEXT: Before deinit
+    print("Before deinit")
+
+    // CHECK-macosx-NEXT: ObjCPrintOnDealloc deinitialized!
+    // CHECK-macosx-NEXT: ObjCPrintOnDealloc deinitialized!
+    testGenericDestroy(ptr, of: NestedWrapper<ObjCPrintOnDealloc>.self)
+
+    ptr.deallocate()
+}
+
+testNestedResilientObjc()