Fix isa swizzling for Kotlin subclasses of Obj-C classes

Don't rely on indexed ivars of instance's isa since isa swizzling
replaces it. Use method instead, it works on the replacement too
(if it is a subclass of the original isa).

 #KT-42482 Fixed.

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

@@ -5,9 +5,9 @@
 
 package org.jetbrains.kotlin.backend.konan.llvm
 
-import llvm.LLVMStoreSizeOfType
+import llvm.LLVMLinkage
+import llvm.LLVMSetLinkage
 import llvm.LLVMValueRef
-import org.jetbrains.kotlin.backend.common.atMostOne
 import org.jetbrains.kotlin.backend.konan.*
 import org.jetbrains.kotlin.backend.konan.descriptors.getAnnotationStringValue
 import org.jetbrains.kotlin.backend.konan.ir.*
@@ -65,7 +65,9 @@ internal class KotlinObjCClassInfoGenerator(override val context: Context) : Con
                 staticData.placeGlobal(
                         "kobjcclassptr:${irClass.fqNameForIrSerialization}#internal",
                         NullPointer(int8Type)
-                ).pointer
+                ).pointer,
+
+                generateClassDataImp(irClass)
         )
 
         objCLLvmDeclarations.classInfoGlobal.setInitializer(info)
@@ -131,6 +133,26 @@ internal class KotlinObjCClassInfoGenerator(override val context: Context) : Con
                 )
             }
 
+    private fun generateClassDataImp(irClass: IrClass): ConstPointer {
+        val classDataPointer = staticData.placeGlobal(
+                "kobjcclassdata:${irClass.fqNameForIrSerialization}#internal",
+                Zero(runtime.kotlinObjCClassData)
+        ).pointer
+
+        val functionType = functionType(classDataPointer.llvmType, false, int8TypePtr, int8TypePtr)
+        val functionName = "kobjcclassdataimp:${irClass.fqNameForIrSerialization}#internal"
+
+        val function = generateFunction(codegen, functionType, functionName) {
+            ret(classDataPointer.llvm)
+        }.also {
+            LLVMSetLinkage(it, LLVMLinkage.LLVMPrivateLinkage)
+        }
+
+        return constPointer(function)
+    }
+
+    private val codegen = CodeGenerator(context)
+
     companion object {
         const val createdClassFieldIndex = 11
     }

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

@@ -44,6 +44,7 @@ class Runtime(bitcodeFile: String) {
 
     val targetData = LLVMCreateTargetData(dataLayout)!!
 
+    val kotlinObjCClassData by lazy { getStructType("KotlinObjCClassData") }
     val kotlinObjCClassInfo by lazy { getStructType("KotlinObjCClassInfo") }
     val objCMethodDescription by lazy { getStructType("ObjCMethodDescription") }
     val objCTypeAdapter by lazy { getStructType("ObjCTypeAdapter") }

kotlin-native/backend.native/tests/interop/objc/tests/kt42482.h

@@ -0,0 +1,10 @@
+#import <Foundation/NSObject.h>
+
+extern BOOL kt42482Deallocated;
+extern id kt42482Global;
+
+@interface KT42482 : NSObject
+-(int)fortyTwo;
+@end;
+
+void kt42482Swizzle(id obj);

kotlin-native/backend.native/tests/interop/objc/tests/kt42482.kt

@@ -0,0 +1,35 @@
+import kotlin.native.ref.WeakReference
+import kotlinx.cinterop.*
+import kotlin.test.*
+import objcTests.*
+
+@Test
+fun testKT42482() {
+    // Attempt to make the state predictable:
+    kotlin.native.internal.GC.collect()
+
+    kt42482Deallocated = false
+    assertFalse(kt42482Deallocated);
+
+    {
+        assertEquals(41, KT42482().fortyTwo())
+
+        val obj: KT42482 = KT42482Impl()
+        assertEquals(42, obj.fortyTwo())
+
+        kt42482Swizzle(obj)
+        assertEquals(43, obj.fortyTwo())
+
+        // Test retain and release on swizzled object:
+        kt42482Global = obj
+        kt42482Global = null
+    }()
+
+    kotlin.native.internal.GC.collect()
+
+    assertTrue(kt42482Deallocated)
+}
+
+class KT42482Impl : KT42482() {
+    override fun fortyTwo() = 42
+}

kotlin-native/backend.native/tests/interop/objc/tests/kt42482.m

@@ -0,0 +1,36 @@
+#import "kt42482.h"
+#import "assert.h"
+#import <objc/runtime.h>
+
+BOOL kt42482Deallocated = NO;
+id kt42482Global = nil;
+
+@implementation KT42482
+-(int)fortyTwo {
+    return 41;
+}
+
+-(void)dealloc {
+    kt42482Deallocated = YES;
+}
+@end;
+
+int fortyTwoSwizzledImp(id self, SEL _cmd) {
+    return 43;
+}
+
+void kt42482Swizzle(id obj) {
+    Class oldClass = object_getClass(obj);
+
+    SEL selector = @selector(fortyTwo);
+
+    Class newClass = objc_allocateClassPair(oldClass, "KT42482Swizzled", 0);
+    assert(newClass != nil);
+    objc_registerClassPair(newClass);
+
+    Method method = class_getInstanceMethod([KT42482 class], selector);
+    assert(method != nil);
+    class_addMethod(newClass, selector, (IMP)&fortyTwoSwizzledImp, method_getTypeEncoding(method));
+
+    object_setClass(obj, newClass);
+}

kotlin-native/runtime/src/main/cpp/ObjCInterop.mm

@@ -46,26 +46,33 @@ const char* Kotlin_ObjCInterop_getUniquePrefix() {
 
 extern "C" id objc_msgSendSuper2(struct objc_super *super, SEL op, ...);
 
-struct KotlinClassData {
+struct KotlinObjCClassData {
   const TypeInfo* typeInfo;
+  Class objcClass;
   int32_t bodyOffset;
 };
 
-static inline struct KotlinClassData* GetKotlinClassData(Class clazz) {
-  void* ivars = object_getIndexedIvars(reinterpret_cast<id>(clazz));
-  return static_cast<struct KotlinClassData*>(ivars);
+// Acts only as container for the method, not actually applied to any class.
+@protocol HasKotlinObjCClassData
+@required
+-(void*)_kotlinObjCClassData;
+@end;
+
+static inline struct KotlinObjCClassData* GetKotlinClassData(id objOrClass) {
+  void* ptr = [(id<HasKotlinObjCClassData>)objOrClass _kotlinObjCClassData];
+  return static_cast<struct KotlinObjCClassData*>(ptr);
 }
 
 namespace {
 
-BackRefFromAssociatedObject* getBackRef(id obj, KotlinClassData* classData) {
+BackRefFromAssociatedObject* getBackRef(id obj, KotlinObjCClassData* 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));
+  auto* classData = GetKotlinClassData(obj);
   return getBackRef(obj, classData);
 }
 
@@ -75,11 +82,11 @@ OBJ_GETTER(toKotlinImp, id self, SEL _cmd) {
 
 id allocWithZoneImp(Class self, SEL _cmd, void* zone) {
   // [super allocWithZone:zone]
-  struct objc_super s = {(id)self, object_getClass((id)self)};
+  auto* classData = GetKotlinClassData(self); // TODO: suboptimal; consider specializing.
+  struct objc_super s = {(id)self, object_getClass(classData->objcClass)};
   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());
@@ -119,6 +126,8 @@ void releaseImp(id self, SEL _cmd) {
 }
 
 void releaseAsAssociatedObjectImp(id self, SEL _cmd) {
+  auto* classData = GetKotlinClassData(self);
+
   // This function is called by the GC. It made a decision to reclaim Kotlin object, and runs
   // deallocation hooks at the moment, including deallocation of the "associated object" ([self])
   // using the [super release] call below.
@@ -131,14 +140,14 @@ void releaseAsAssociatedObjectImp(id self, SEL _cmd) {
   // Generally retaining and releasing Kotlin object that is being deallocated would lead to
   // use-after-dispose and double-dispose problems (with unpredictable consequences) or to an assertion failure.
   // To workaround this, detach the back ref from the Kotlin object:
-  getBackRef(self)->detach();
+  getBackRef(self, classData)->detach();
   // So retain/release/etc. on [self] won't affect the Kotlin object, and an attempt to get
   // the reference to it (e.g. when calling Kotlin method on [self]) would crash.
   // The latter is generally ok, because by the time superclass dealloc gets launched, subclass state
   // should already be deinitialized, and Kotlin methods operate on the subclass.
 
   // [super release]
-  Class clazz = object_getClass(self);
+  Class clazz = classData->objcClass;
   struct objc_super s = {self, clazz};
   auto messenger = reinterpret_cast<void (*) (struct objc_super*, SEL _cmd)>(objc_msgSendSuper2);
   messenger(&s, @selector(release));
@@ -150,17 +159,12 @@ extern "C" {
 
 Class Kotlin_Interop_getObjCClass(const char* name);
 
-static inline void SetKotlinTypeInfo(Class clazz, const TypeInfo* typeInfo) {
-  GetKotlinClassData(clazz)->typeInfo = typeInfo;
-}
-
 const TypeInfo* GetObjCKotlinTypeInfo(ObjHeader* obj) RUNTIME_NOTHROW;
 
 RUNTIME_NOTHROW const TypeInfo* GetObjCKotlinTypeInfo(ObjHeader* obj) {
     void* objcPtr = obj->GetAssociatedObject();
     RuntimeAssert(objcPtr != nullptr, "");
-    Class clazz = object_getClass(reinterpret_cast<id>(objcPtr));
-    return GetKotlinClassData(clazz)->typeInfo;
+    return GetKotlinClassData(reinterpret_cast<id>(objcPtr))->typeInfo;
 }
 
 
@@ -181,6 +185,24 @@ static void AddNSObjectOverride(bool isClassMethod, Class clazz, SEL selector, v
   RuntimeCheck(added, "Unable to add method to Objective-C class");
 }
 
+static void AddKotlinClassData(bool isClassMethod, Class clazz, void* imp) {
+  SEL selector = @selector(_kotlinObjCClassData);
+
+  auto methodDescription = protocol_getMethodDescription(
+      @protocol(HasKotlinObjCClassData),
+      selector,
+      YES, YES
+  );
+
+  const char* typeEncoding = methodDescription.types;
+
+  RuntimeCheck(typeEncoding != nullptr, "unable to find method in Objective-C protocol");
+
+  BOOL added = class_addMethod(
+      isClassMethod ? object_getClass((id)clazz) : clazz, selector, (IMP)imp, typeEncoding);
+  RuntimeCheck(added, "Unable to add method to Objective-C class");
+}
+
 struct ObjCMethodDescription {
   void* (*imp)(void*, void*, ...);
   const char* selector;
@@ -206,6 +228,8 @@ struct KotlinObjCClassInfo {
   const TypeInfo* metaTypeInfo;
 
   void** createdClass;
+
+  KotlinObjCClassData* (*classDataImp)(void*, void*);
 };
 
 static void AddMethods(Class clazz, const struct ObjCMethodDescription* methods, int32_t methodsNum) {
@@ -221,11 +245,10 @@ static int anonymousClassNextId = 0;
 
 static Class allocateClass(const KotlinObjCClassInfo* info) {
   Class superclass = Kotlin_Interop_getObjCClass(info->superclassName);
-  size_t extraBytes = sizeof(struct KotlinClassData);
 
   if (info->exported) {
     RuntimeCheck(info->name != nullptr, "exported Objective-C class must have a name");
-    Class result = objc_allocateClassPair(superclass, info->name, extraBytes);
+    Class result = objc_allocateClassPair(superclass, info->name, 0);
     if (result != nullptr) return result;
     // Similar to how Objective-C runtime handles this:
     fprintf(stderr, "Class %s has multiple implementations. Which one will be used is undefined.\n", info->name);
@@ -242,7 +265,7 @@ static Class allocateClass(const KotlinObjCClassInfo* info) {
   int classId = anonymousClassNextId++;
   className += std::to_string(classId);
 
-  Class result = objc_allocateClassPair(superclass, className.c_str(), extraBytes);
+  Class result = objc_allocateClassPair(superclass, className.c_str(), 0);
   RuntimeCheck(result != nullptr, "Failed to allocate Objective-C class");
   return result;
 }
@@ -284,7 +307,12 @@ void* CreateKotlinObjCClass(const KotlinObjCClassInfo* info) {
   AddMethods(newClass, info->instanceMethods, info->instanceMethodsNum);
   AddMethods(newMetaclass, info->classMethods, info->classMethodsNum);
 
-  SetKotlinTypeInfo(newClass, Kotlin_ObjCExport_createTypeInfoWithKotlinFieldsFrom(newClass, info->typeInfo));
+  // Adding both instance and class methods to make [GetKotlinClassData] work
+  // for instances as well as the class itself.
+  AddKotlinClassData(false, newClass, (void*)info->classDataImp);
+  AddKotlinClassData(true, newClass, (void*)info->classDataImp);
+
+  const TypeInfo* actualTypeInfo = Kotlin_ObjCExport_createTypeInfoWithKotlinFieldsFrom(newClass, info->typeInfo);
 
   int bodySize = sizeof(BackRefFromAssociatedObject);
   char bodyTypeEncoding[16];
@@ -297,9 +325,15 @@ void* CreateKotlinObjCClass(const KotlinObjCClassInfo* info) {
   Ivar body = class_getInstanceVariable(newClass, "kotlinBody");
   RuntimeAssert(body != nullptr, "Unable to get ivar added to Objective-C class");
   int32_t offset = (int32_t)ivar_getOffset(body);
-  GetKotlinClassData(newClass)->bodyOffset = offset;
   *info->bodyOffset = offset;
 
+  // Doing this after objc_registerClassPair because it is not clear whether calling class methods
+  // is safe before that.
+  auto* classData = GetKotlinClassData(newClass);
+  classData->typeInfo = actualTypeInfo;
+  classData->objcClass = newClass;
+  classData->bodyOffset = offset;
+
   *info->createdClass = newClass;
   return newClass;
 }