Fix sharing for instances of Kotlin subclasses of Obj-C classes

Make Obj-C references to these objects properly shareable,
fix deallocation on other threads.

To achieve that, move Kotlin fields to Kotlin counterpart
and make its lifetime equal to Obj-C counterpart lifetime,
like with KotlinBase.

 #KT-37225 Fixed.

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/descriptors/DescriptorUtils.kt

@@ -262,8 +262,8 @@ fun IrDeclaration.findTopLevelDeclaration(): IrDeclaration = when {
 
 internal val IrClass.isFrozen: Boolean
     get() = annotations.hasAnnotation(KonanFqNames.frozen) ||
-            // RTTI is used for non-reference type box or Objective-C object wrapper:
-            !this.defaultType.binaryTypeIsReference() || this.isObjCClass()
+            // RTTI is used for non-reference type box:
+            !this.defaultType.binaryTypeIsReference()
 
 fun IrConstructorCall.getAnnotationStringValue() = (getValueArgument(0) as? IrConst<String>)?.value
 

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/CodeGenerator.kt

@@ -1005,7 +1005,7 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
             }
 
             val classInfo = codegen.kotlinObjCClassInfo(irClass)
-            val classPointerGlobal = load(structGep(classInfo, 12 /* createdClass */))
+            val classPointerGlobal = load(structGep(classInfo, KotlinObjCClassInfoGenerator.createdClassFieldIndex))
 
             val storedClass = this.load(classPointerGlobal)
 

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/IrToBitcode.kt

@@ -1614,29 +1614,11 @@ internal class CodeGeneratorVisitor(val context: Context, val lifetimes: Map<IrE
 
         val typePtr = pointerType(fieldInfo.classBodyType)
 
-        val bodyPtr = getObjectBodyPtr(value.parentAsClass, thisPtr)
-
-        val typedBodyPtr = functionGenerationContext.bitcast(typePtr, bodyPtr)
+        val typedBodyPtr = functionGenerationContext.bitcast(typePtr, thisPtr)
         val fieldPtr = LLVMBuildStructGEP(functionGenerationContext.builder, typedBodyPtr, fieldInfo.index, "")
         return fieldPtr!!
     }
 
-    private fun getObjectBodyPtr(irClass: IrClass, objectPtr: LLVMValueRef): LLVMValueRef {
-        return if (irClass.isObjCClass()) {
-            assert(irClass.isKotlinObjCClass())
-
-            val objCPtr = callDirect(context.ir.symbols.interopObjCObjectRawValueGetter.owner,
-                    listOf(objectPtr), Lifetime.IRRELEVANT)
-
-            val objCDeclarations = context.llvmDeclarations.forClass(irClass).objCDeclarations!!
-            val bodyOffset = functionGenerationContext.load(objCDeclarations.bodyOffsetGlobal.llvmGlobal)
-
-            functionGenerationContext.gep(objCPtr, bodyOffset)
-        } else {
-            objectPtr
-        }
-    }
-
     //-------------------------------------------------------------------------//
     private fun evaluateStringConst(value: IrConst<String>) =
             context.llvm.staticData.kotlinStringLiteral(value.value).llvm

backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/llvm/KotlinObjCClassInfoGenerator.kt

@@ -39,9 +39,6 @@ internal class KotlinObjCClassInfoGenerator(override val context: Context) : Con
             it.name.asString().removeSuffix("Protocol")
         }
 
-        val bodySize =
-                LLVMStoreSizeOfType(llvmTargetData, context.llvmDeclarations.forClass(irClass).bodyType).toInt()
-
         val exportedClassName = selectExportedClassName(irClass)
         val className = exportedClassName ?: selectInternalClassName(irClass)
 
@@ -60,7 +57,6 @@ internal class KotlinObjCClassInfoGenerator(override val context: Context) : Con
                 staticData.placeGlobalConstArray("", runtime.objCMethodDescription, classMethods),
                 Int32(classMethods.size),
 
-                Int32(bodySize),
                 objCLLvmDeclarations.bodyOffsetGlobal.pointer,
 
                 irClass.typeInfoPtr,
@@ -134,6 +130,10 @@ internal class KotlinObjCClassInfoGenerator(override val context: Context) : Con
                         it.llvmFunction
                 )
             }
+
+    companion object {
+        const val createdClassFieldIndex = 11
+    }
 }
 
 internal fun CodeGenerator.kotlinObjCClassInfo(irClass: IrClass): LLVMValueRef {

backend.native/tests/interop/objc/smoke.h

@@ -273,3 +273,16 @@ struct MyStruct myStruct = {11, 12, 13, 14};
 @interface FooMangled : NSObject<Proto>
 //- (void) CompanionS;  // mangleSimple does not support this: it may clash after mangling
 @end;
+
+
+@class DeallocListener;
+
+@interface DeallocExecutor : NSObject
+@property DeallocListener* deallocListener;
+@end;
+
+@interface DeallocListener : NSObject
+@property (weak) DeallocExecutor* deallocExecutor;
+@property BOOL deallocated;
+-(BOOL)deallocExecutorIsNil;
+@end;

backend.native/tests/interop/objc/smoke.kt

@@ -5,6 +5,7 @@
 
 import kotlinx.cinterop.*
 import objcSmoke.*
+import kotlin.native.concurrent.*
 import kotlin.native.ref.*
 import kotlin.test.*
 
@@ -34,6 +35,8 @@ fun run() {
     testConstructorReturnsNull()
     testCallableReferences()
     testMangling()
+    testSharing()
+    testObjCWeakRef()
 
     assertEquals(2, ForwardDeclaredEnum.TWO.value)
 
@@ -103,6 +106,7 @@ fun run() {
         override fun description() = "global object"
     }
     println(globalObject)
+    globalObject = null // Prevent Kotlin object above from leaking.
 
     println(formatStringLength("%d %d", 42, 17))
 
@@ -530,6 +534,80 @@ fun testMangling() {
     assertEquals(Companion, enumMangledStruct.smth)
 }
 
+private class NSObjectImpl : NSObject() {
+    var x = 111
+}
+
+fun testSharing() = withWorker {
+    val obj = NSObjectImpl()
+    val array = nsArrayOf(obj)
+
+    assertFalse(obj.isFrozen)
+
+    runInWorker {
+        assertFailsWith<IncorrectDereferenceException> {
+            array.objectAtIndex(0)
+        }
+    }
+
+    obj.x = 222
+    obj.freeze()
+    assertTrue(obj.isFrozen)
+
+    runInWorker {
+        val obj1 = array.objectAtIndex(0) as NSObjectImpl
+        assertFailsWith<InvalidMutabilityException> {
+            obj1.x = 333
+        }
+    }
+
+    assertEquals(222, obj.x)
+
+    // TODO: test [obj release] etc.
+}
+
+fun testObjCWeakRef() {
+    val deallocListener = DeallocListener()
+    assertFalse(deallocListener.deallocated)
+
+    testObjCWeakRef0(deallocListener)
+
+    kotlin.native.internal.GC.collect()
+    assertTrue(deallocListener.deallocated)
+    assertTrue(deallocListener.deallocExecutorIsNil())
+}
+
+fun testObjCWeakRef0(deallocListener: DeallocListener) = withWorker {
+    assertTrue(deallocListener.deallocExecutorIsNil())
+
+    val obj = object : DeallocExecutor() {}
+    deallocListener.deallocExecutor = obj
+    obj.deallocListener = deallocListener
+
+    assertFalse(deallocListener.deallocExecutorIsNil())
+
+//    TODO: can't actually test, Obj-C runtime doesn't expect _tryRetain throwing an exception.
+//    runInWorker {
+//        assertFailsWith<IncorrectDereferenceException> {
+//            deallocListener.deallocExecutorIsNil()
+//        }
+//    }
+
+    obj.freeze()
+
+    runInWorker {
+        assertFalse(deallocListener.deallocExecutorIsNil())
+    }
+}
+
+private fun Worker.runInWorker(block: () -> Unit) {
+    block.freeze()
+    val future = this.execute(TransferMode.SAFE, { block }) {
+        it()
+    }
+    future.result // Throws on failure.
+}
+
 private val Any.objCClassName: String
     get() = object_getClassName(this)!!.toKString()
 

backend.native/tests/interop/objc/smoke.m

@@ -342,3 +342,16 @@ BOOL customStringDeallocated = NO;
 @implementation FooMangled : NSObject
 @synthesize Companion;
 @end;
+
+
+@implementation DeallocExecutor
+-(void)dealloc {
+    self.deallocListener.deallocated = YES;
+}
+@end;
+
+@implementation DeallocListener
+-(BOOL)deallocExecutorIsNil {
+    return self.deallocExecutor == nil;
+}
+@end;

runtime/src/main/cpp/Memory.cpp

@@ -1655,14 +1655,6 @@ void garbageCollect() {
 
 #endif  // USE_GC
 
-void deinitInstanceBody(const TypeInfo* typeInfo, void* body) {
-  for (int index = 0; index < typeInfo->objOffsetsCount_; index++) {
-    ObjHeader** location = reinterpret_cast<ObjHeader**>(
-        reinterpret_cast<uintptr_t>(body) + typeInfo->objOffsets_[index]);
-    ZeroHeapRef(location);
-  }
-}
-
 ForeignRefManager* initLocalForeignRef(ObjHeader* object) {
   if (!IsStrictMemoryModel) return nullptr;
 
@@ -2745,10 +2737,6 @@ void ReleaseHeapRefRelaxed(const ObjHeader* object) {
   releaseHeapRef<false>(const_cast<ObjHeader*>(object));
 }
 
-void DeinitInstanceBody(const TypeInfo* typeInfo, void* body) {
-  deinitInstanceBody(typeInfo, body);
-}
-
 ForeignRefContext InitLocalForeignRef(ObjHeader* object) {
   return initLocalForeignRef(object);
 }

runtime/src/main/cpp/MemoryPrivate.hpp

@@ -25,8 +25,6 @@ bool TryAddHeapRef(const ObjHeader* object);
 
 MODEL_VARIANTS(void, ReleaseHeapRef, const ObjHeader* object);
 
-void DeinitInstanceBody(const TypeInfo* typeInfo, void* body);
-
 void Kotlin_ObjCExport_releaseAssociatedObject(void* associatedObject);
 
 ForeignRefContext InitLocalForeignRef(ObjHeader* object);

runtime/src/main/cpp/ObjCExport.h

@@ -4,6 +4,7 @@
 #if KONAN_OBJC_INTEROP
 
 #import <objc/runtime.h>
+#import <Foundation/NSString.h>
 
 #import "Types.h"
 #import "Memory.h"

runtime/src/main/cpp/ObjCExport.mm

@@ -639,19 +639,26 @@ static const TypeInfo* createTypeInfo(
   const KStdOrderedMap<ClassId, KStdVector<VTableElement>>& interfaceVTables,
   const InterfaceTableRecord* superItable,
   int superItableSize,
-  bool itableEqualsSuper
+  bool itableEqualsSuper,
+  const TypeInfo* fieldsInfo
 ) {
   TypeInfo* result = (TypeInfo*)konanAllocMemory(sizeof(TypeInfo) + vtable.size() * sizeof(void*));
   result->typeInfo_ = result;
 
   result->flags_ = TF_OBJC_DYNAMIC;
 
-  result->instanceSize_ = superType->instanceSize_;
   result->superType_ = superType;
-  result->objOffsets_ = superType->objOffsets_;
-  result->objOffsetsCount_ = superType->objOffsetsCount_; // So TF_IMMUTABLE can also be inherited:
-  if ((superType->flags_ & TF_IMMUTABLE) != 0) {
-    result->flags_ |= TF_IMMUTABLE;
+  if (fieldsInfo == nullptr) {
+    result->instanceSize_ = superType->instanceSize_;
+    result->objOffsets_ = superType->objOffsets_;
+    result->objOffsetsCount_ = superType->objOffsetsCount_; // So TF_IMMUTABLE can also be inherited:
+    if ((superType->flags_ & TF_IMMUTABLE) != 0) {
+      result->flags_ |= TF_IMMUTABLE;
+    }
+  } else {
+    result->instanceSize_ = fieldsInfo->instanceSize_;
+    result->objOffsets_ = fieldsInfo->objOffsets_;
+    result->objOffsetsCount_ = fieldsInfo->objOffsetsCount_;
   }
 
   result->classId_ = superType->classId_;
@@ -793,7 +800,7 @@ static void throwIfCantBeOverridden(Class clazz, const KotlinToObjCMethodAdapter
   }
 }
 
-static const TypeInfo* createTypeInfo(Class clazz, const TypeInfo* superType) {
+static const TypeInfo* createTypeInfo(Class clazz, const TypeInfo* superType, const TypeInfo* fieldsInfo) {
   Class superClass = class_getSuperclass(clazz);
 
   KStdUnorderedSet<SEL> definedSelectors;
@@ -931,7 +938,8 @@ static const TypeInfo* createTypeInfo(Class clazz, const TypeInfo* superType) {
   // TODO: consider forbidding the class being abstract.
 
   const TypeInfo* result = createTypeInfo(superType, addedInterfaces, vtable, methodTable,
-                                          interfaceVTables, superITable, superITableSize, itableEqualsSuper);
+                                          interfaceVTables, superITable, superITableSize, itableEqualsSuper,
+                                          fieldsInfo);
 
   // TODO: it will probably never be requested, since such a class can't be instantiated in Kotlin.
   result->writableInfo_->objCExport.objCClass = clazz;
@@ -956,13 +964,22 @@ static const TypeInfo* getOrCreateTypeInfo(Class clazz) {
 
   result = Kotlin_ObjCExport_getAssociatedTypeInfo(clazz); // double-checking.
   if (result == nullptr) {
-    result = createTypeInfo(clazz, superType);
+    result = createTypeInfo(clazz, superType, nullptr);
     setAssociatedTypeInfo(clazz, result);
   }
 
   return result;
 }
 
+const TypeInfo* Kotlin_ObjCExport_createTypeInfoWithKotlinFieldsFrom(Class clazz, const TypeInfo* fieldsInfo) {
+  Class superClass = class_getSuperclass(clazz);
+  RuntimeCheck(superClass != nullptr, "");
+
+  const TypeInfo* superType = getOrCreateTypeInfo(superClass);
+
+  return createTypeInfo(clazz, superType, fieldsInfo);
+}
+
 static SimpleMutex classCreationMutex;
 static int anonymousClassNextId = 0;
 

runtime/src/main/cpp/ObjCExportPrivate.h

@@ -24,6 +24,8 @@ extern "C" OBJ_GETTER(Kotlin_ObjCExport_convertUnmappedObjCObject, id obj);
 extern "C" SEL Kotlin_ObjCExport_toKotlinSelector;
 extern "C" SEL Kotlin_ObjCExport_releaseAsAssociatedObjectSelector;
 
+const TypeInfo* Kotlin_ObjCExport_createTypeInfoWithKotlinFieldsFrom(Class clazz, const TypeInfo* fieldsInfo);
+
 #endif // KONAN_OBJC_INTEROP
 
 #endif // RUNTIME_OBJCEXPORTPRIVATE_H

runtime/src/main/cpp/ObjCInterop.mm

@@ -23,10 +23,11 @@
 #include <cstdint>
 
 #include "Memory.h"
-#include "MemoryPrivate.hpp"
+#include "MemorySharedRefs.hpp"
 
 #include "Natives.h"
 #include "ObjCInterop.h"
+#include "ObjCExportPrivate.h"
 #include "Types.h"
 #include "Utils.h"
 
@@ -39,9 +40,7 @@ const char* Kotlin_ObjCInterop_getUniquePrefix() {
   return result;
 }
 
-extern "C" {
-
-Class Kotlin_Interop_getObjCClass(const char* name);
+extern "C" id objc_msgSendSuper2(struct objc_super *super, SEL op, ...);
 
 struct KotlinClassData {
   const TypeInfo* typeInfo;
@@ -53,6 +52,70 @@ static inline struct KotlinClassData* GetKotlinClassData(Class clazz) {
   return static_cast<struct KotlinClassData*>(ivars);
 }
 
+namespace {
+
+BackRefFromAssociatedObject* getBackRef(id obj, KotlinClassData* classData) {
+  void* body = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(obj) + classData->bodyOffset);
+  return reinterpret_cast<BackRefFromAssociatedObject*>(body);
+}
+
+BackRefFromAssociatedObject* getBackRef(id obj) {
+  // TODO: suboptimal; consider specializing methods for each class.
+  auto* classData = GetKotlinClassData(object_getClass(obj));
+  return getBackRef(obj, classData);
+}
+
+OBJ_GETTER(toKotlinImp, id self, SEL _cmd) {
+  RETURN_OBJ(getBackRef(self)->ref());
+}
+
+id allocWithZoneImp(Class self, SEL _cmd, void* zone) {
+  // [super allocWithZone:zone]
+  struct objc_super s = {(id)self, object_getClass((id)self)};
+  auto messenger = reinterpret_cast<id (*) (struct objc_super*, SEL _cmd, void* zone)>(objc_msgSendSuper2);
+  id result = messenger(&s, _cmd, zone);
+
+  auto* classData = GetKotlinClassData(self); // TODO: suboptimal; consider specializing.
+  auto* typeInfo = classData->typeInfo;
+  ObjHolder holder;
+  auto kotlinObj = AllocInstanceWithAssociatedObject(typeInfo, result, holder.slot());
+
+  getBackRef(result, classData)->initAndAddRef(kotlinObj);
+
+  return result;
+}
+
+id retainImp(id self, SEL _cmd) {
+  getBackRef(self)->addRef();
+  return self;
+}
+
+BOOL _tryRetainImp(id self, SEL _cmd) {
+  // TODO: [tryAddRef] currently works only on the owner thread for non-shared objects;
+  // this is a regression for instances of Kotlin subclasses of Obj-C classes:
+  // loading a reference to such an object from Obj-C weak reference now fails on "wrong" thread
+  // unless the object is frozen.
+  return getBackRef(self)->tryAddRef();
+}
+
+void releaseImp(id self, SEL _cmd) {
+  getBackRef(self)->releaseRef();
+}
+
+void releaseAsAssociatedObjectImp(id self, SEL _cmd) {
+  // [super release]
+  Class clazz = object_getClass(self);
+  struct objc_super s = {self, clazz};
+  auto messenger = reinterpret_cast<void (*) (struct objc_super*, SEL _cmd)>(objc_msgSendSuper2);
+  messenger(&s, @selector(release));
+}
+
+}
+
+extern "C" {
+
+Class Kotlin_Interop_getObjCClass(const char* name);
+
 static inline void SetKotlinTypeInfo(Class clazz, const TypeInfo* typeInfo) {
   GetKotlinClassData(clazz)->typeInfo = typeInfo;
 }
@@ -67,39 +130,22 @@ RUNTIME_NOTHROW const TypeInfo* GetObjCKotlinTypeInfo(ObjHeader* obj) {
   return GetKotlinClassData(clazz)->typeInfo;
 }
 
-id objc_msgSendSuper2(struct objc_super *super, SEL op, ...);
 
-static void DeallocImp(id self, SEL _cmd) {
-  // TODO: doesn't support overriding Kotlin classes.
-  Class clazz = object_getClass(self);
-  RuntimeAssert(clazz != nullptr, "Must not be null");
-
-  struct KotlinClassData* classData = GetKotlinClassData(clazz);
-  void* body = reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(self) + classData->bodyOffset);
-
-  const TypeInfo* typeInfo = classData->typeInfo;
-  DeinitInstanceBody(typeInfo, body);
-
-  // Call super.dealloc:
-  struct objc_super s = {self, clazz};
-  auto messenger = reinterpret_cast<void (*) (struct objc_super*, SEL _cmd)>(objc_msgSendSuper2);
-  messenger(&s, _cmd);
-}
-
-static void AddDeallocMethod(Class clazz) {
+static void AddNSObjectOverride(bool isClassMethod, Class clazz, SEL selector, void* imp) {
   Class nsObjectClass = Kotlin_Interop_getObjCClass("NSObject");
 
-  SEL deallocSelector = sel_registerName("dealloc");
-  Method nsObjectDeallocMethod = class_getInstanceMethod(nsObjectClass, deallocSelector);
-  RuntimeAssert(nsObjectDeallocMethod != nullptr, "[NSObject dealloc] method not found");
+  Method nsObjectMethod = class_getInstanceMethod(
+      isClassMethod ? object_getClass((id)nsObjectClass) : nsObjectClass, selector);
+  RuntimeCheck(nsObjectMethod != nullptr, "NSObject method not found");
 
-  const char* nsObjectDeallocMethodTypeEncoding = method_getTypeEncoding(nsObjectDeallocMethod);
-  RuntimeAssert(nsObjectDeallocMethodTypeEncoding != nullptr, "[NSObject dealloc] method has no encoding provided");
+  const char* nsObjectMethodTypeEncoding = method_getTypeEncoding(nsObjectMethod);
+  RuntimeCheck(nsObjectMethodTypeEncoding != nullptr, "NSObject method has no encoding provided");
 
   // TODO: something of the above can be cached.
 
-  BOOL added = class_addMethod(clazz, deallocSelector, (IMP)DeallocImp, nsObjectDeallocMethodTypeEncoding);
-  RuntimeAssert(added, "Unable to add dealloc method to Objective-C class");
+  BOOL added = class_addMethod(
+      isClassMethod ? object_getClass((id)clazz) : clazz, selector, (IMP)imp, nsObjectMethodTypeEncoding);
+  RuntimeCheck(added, "Unable to add method to Objective-C class");
 }
 
 struct ObjCMethodDescription {
@@ -121,7 +167,6 @@ struct KotlinObjCClassInfo {
   const struct ObjCMethodDescription* classMethods;
   int32_t classMethodsNum;
 
-  int32_t bodySize;
   int32_t* bodyOffset;
 
   const TypeInfo* typeInfo;
@@ -195,17 +240,23 @@ void* CreateKotlinObjCClass(const KotlinObjCClassInfo* info) {
     }
   }
 
-  AddDeallocMethod(newClass);
+  AddNSObjectOverride(false, newClass, Kotlin_ObjCExport_toKotlinSelector, (void*)&toKotlinImp);
+  AddNSObjectOverride(true, newClass, @selector(allocWithZone:), (void*)&allocWithZoneImp);
+  AddNSObjectOverride(false, newClass, @selector(retain), (void*)&retainImp);
+  AddNSObjectOverride(false, newClass, @selector(_tryRetain), (void*)&_tryRetainImp);
+  AddNSObjectOverride(false, newClass, @selector(release), (void*)&releaseImp);
+  AddNSObjectOverride(false, newClass, Kotlin_ObjCExport_releaseAsAssociatedObjectSelector,
+      (void*)&releaseAsAssociatedObjectImp);
 
   AddMethods(newClass, info->instanceMethods, info->instanceMethodsNum);
   AddMethods(newMetaclass, info->classMethods, info->classMethodsNum);
 
-  SetKotlinTypeInfo(newClass, info->typeInfo);
-  SetKotlinTypeInfo(newMetaclass, info->metaTypeInfo);
+  SetKotlinTypeInfo(newClass, Kotlin_ObjCExport_createTypeInfoWithKotlinFieldsFrom(newClass, info->typeInfo));
 
+  int bodySize = sizeof(BackRefFromAssociatedObject);
   char bodyTypeEncoding[16];
-  snprintf(bodyTypeEncoding, sizeof(bodyTypeEncoding), "[%dc]", info->bodySize);
-  BOOL added = class_addIvar(newClass, "kotlinBody", info->bodySize, /* log2(align) = */ 3, bodyTypeEncoding);
+  snprintf(bodyTypeEncoding, sizeof(bodyTypeEncoding), "[%dc]", bodySize);
+  BOOL added = class_addIvar(newClass, "kotlinBody", bodySize, /* log2(align) = */ 3, bodyTypeEncoding);
   RuntimeAssert(added == YES, "Unable to add ivar to Objective-C class");
 
   objc_registerClassPair(newClass);