Implement thread-safe tracking of Obj-C references to Kotlin objects

Fix #2052
Fix #3204

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

@@ -442,6 +442,11 @@ internal class Llvm(val context: Context, val llvmModule: LLVMModuleRef) {
     val freezeSubgraph = importRtFunction("FreezeSubgraph")
     val checkMainThread = importRtFunction("CheckIsMainThread")
 
+    val kRefSharedHolderInitLocal = importRtFunction("KRefSharedHolder_initLocal")
+    val kRefSharedHolderInit = importRtFunction("KRefSharedHolder_init")
+    val kRefSharedHolderDispose = importRtFunction("KRefSharedHolder_dispose")
+    val kRefSharedHolderRef = importRtFunction("KRefSharedHolder_ref")
+
     val createKotlinObjCClass by lazy { importRtFunction("CreateKotlinObjCClass") }
     val getObjCKotlinTypeInfo by lazy { importRtFunction("GetObjCKotlinTypeInfo") }
     val missingInitImp by lazy { importRtFunction("MissingInitImp") }

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

@@ -136,9 +136,11 @@ private val BlockPointerBridge.nameSuffix: String
 internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjCExportCodeGenerator) {
     private val codegen get() = objCExportCodeGenerator.codegen
 
+    private val kRefSharedHolderType = LLVMGetTypeByName(codegen.runtime.llvmModule, "class.KRefSharedHolder")!!
+
     private val blockLiteralType = structType(
             codegen.runtime.getStructType("Block_literal_1"),
-            codegen.kObjHeaderPtr
+            kRefSharedHolderType
     )
 
     private val blockDescriptorType = codegen.runtime.getStructType("Block_descriptor_1")
@@ -149,8 +151,8 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
             "blockDisposeHelper"
     ) {
         val blockPtr = bitcast(pointerType(blockLiteralType), param(0))
-        val slot = structGep(blockPtr, 1)
-        storeHeapRef(kNullObjHeaderPtr, slot) // TODO: can dispose_helper write to the block?
+        val refHolder = structGep(blockPtr, 1)
+        call(context.llvm.kRefSharedHolderDispose, listOf(refHolder))
 
         ret(null)
     }.also {
@@ -163,15 +165,22 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
             "blockCopyHelper"
     ) {
         val dstBlockPtr = bitcast(pointerType(blockLiteralType), param(0))
-        val dstSlot = structGep(dstBlockPtr, 1)
+        val dstRefHolder = structGep(dstBlockPtr, 1)
 
         val srcBlockPtr = bitcast(pointerType(blockLiteralType), param(1))
-        val srcSlot = structGep(srcBlockPtr, 1)
+        val srcRefHolder = structGep(srcBlockPtr, 1)
 
-        // Kotlin reference was `memcpy`ed from src to dst, "revert" this:
-        storeRefUnsafe(kNullObjHeaderPtr, dstSlot)
-        // and copy properly:
-        storeHeapRef(loadSlot(srcSlot, isVar = false), dstSlot)
+        // Note: in current implementation copy helper is invoked only for stack-allocated blocks from the same thread,
+        // so it is technically not necessary to check owner.
+        // However this is not guaranteed by Objective-C runtime, so keep it suboptimal but reliable:
+        val ref = call(
+                context.llvm.kRefSharedHolderRef,
+                listOf(srcRefHolder),
+                exceptionHandler = ExceptionHandler.Caller,
+                verbatim = true
+        )
+
+        call(context.llvm.kRefSharedHolderInit, listOf(dstRefHolder, ref))
 
         ret(null)
     }.also {
@@ -211,22 +220,17 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
 
 
 
-    private fun FunctionGenerationContext.storeRefUnsafe(value: LLVMValueRef, slot: LLVMValueRef) {
-        assert(value.type == kObjHeaderPtr)
-        assert(slot.type == kObjHeaderPtrPtr)
-
-        store(
-                bitcast(int8TypePtr, value),
-                bitcast(pointerType(int8TypePtr), slot)
-        )
-    }
-
     private fun ObjCExportCodeGenerator.generateInvoke(bridge: BlockPointerBridge): ConstPointer {
         val numberOfParameters = bridge.numberOfParameters
 
         val result = generateFunction(codegen, bridge.blockInvokeLlvmType, "invokeBlock${bridge.nameSuffix}") {
             val blockPtr = bitcast(pointerType(blockLiteralType), param(0))
-            val kotlinFunction = loadSlot(structGep(blockPtr, 1), isVar = false)
+            val kotlinFunction = call(
+                    context.llvm.kRefSharedHolderRef,
+                    listOf(structGep(blockPtr, 1)),
+                    exceptionHandler = ExceptionHandler.Caller,
+                    verbatim = true
+            )
 
             val args = (1 .. numberOfParameters).map { index ->
                 objCReferenceToKotlin(param(index), Lifetime.ARGUMENT)
@@ -282,7 +286,7 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
 
             val blockOnStack = alloca(blockLiteralType)
             val blockOnStackBase = structGep(blockOnStack, 0)
-            val slot = structGep(blockOnStack, 1)
+            val refHolder = structGep(blockOnStack, 1)
 
             listOf(bitcast(int8TypePtr, isa), flags, reserved, invoke, descriptor).forEachIndexed { index, value ->
                 // Although value is actually on the stack, it's not in normal slot area, so we cannot handle it
@@ -290,8 +294,7 @@ internal class BlockAdapterToFunctionGenerator(val objCExportCodeGenerator: ObjC
                 store(value, structGep(blockOnStackBase, index))
             }
 
-            // Note: it is the slot in the block located on stack, so no need to manage it properly:
-            storeRefUnsafe(kotlinRef, slot)
+            call(context.llvm.kRefSharedHolderInitLocal, listOf(refHolder, kotlinRef))
 
             val copiedBlock = callFromBridge(retainBlock, listOf(bitcast(int8TypePtr, blockOnStack)))
 

backend.native/tests/framework/values/expectedLazy.h

@@ -960,6 +960,29 @@ __attribute__((swift_name("TestWeakRefs")))
 - (NSArray<id> *)createCycle __attribute__((swift_name("createCycle()")));
 @end;
 
+__attribute__((objc_subclassing_restricted))
+__attribute__((swift_name("SharedRefs")))
+@interface ValuesSharedRefs : KotlinBase
+- (instancetype)init __attribute__((swift_name("init()"))) __attribute__((objc_designated_initializer));
++ (instancetype)new __attribute__((availability(swift, unavailable, message="use object initializers instead")));
+- (ValuesSharedRefsMutableData *)createRegularObject __attribute__((swift_name("createRegularObject()")));
+- (void (^)(void))createLambda __attribute__((swift_name("createLambda()")));
+- (NSMutableArray<id> *)createCollection __attribute__((swift_name("createCollection()")));
+- (ValuesSharedRefsMutableData *)createFrozenRegularObject __attribute__((swift_name("createFrozenRegularObject()")));
+- (void (^)(void))createFrozenLambda __attribute__((swift_name("createFrozenLambda()")));
+- (NSMutableArray<id> *)createFrozenCollection __attribute__((swift_name("createFrozenCollection()")));
+- (BOOL)hasAliveObjects __attribute__((swift_name("hasAliveObjects()")));
+@end;
+
+__attribute__((objc_subclassing_restricted))
+__attribute__((swift_name("SharedRefs.MutableData")))
+@interface ValuesSharedRefsMutableData : KotlinBase
+- (instancetype)init __attribute__((swift_name("init()"))) __attribute__((objc_designated_initializer));
++ (instancetype)new __attribute__((availability(swift, unavailable, message="use object initializers instead")));
+- (void)update __attribute__((swift_name("update()")));
+@property int32_t x __attribute__((swift_name("x")));
+@end;
+
 @interface ValuesEnumeration (ValuesKt)
 - (ValuesEnumeration *)getAnswer __attribute__((swift_name("getAnswer()")));
 @end;

backend.native/tests/framework/values/values.kt

@@ -10,6 +10,7 @@ package conversions
 
 import kotlin.native.concurrent.freeze
 import kotlin.native.concurrent.isFrozen
+import kotlin.native.ref.WeakReference
 import kotlin.properties.ReadWriteProperty
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
@@ -738,3 +739,38 @@ class TestWeakRefs(private val frozen: Boolean) {
 
     private class Node(var next: Node?)
 }
+
+class SharedRefs {
+    class MutableData {
+        var x = 0
+
+        fun update() { x += 1 }
+    }
+
+    fun createRegularObject(): MutableData = create { MutableData() }
+
+    fun createLambda(): () -> Unit = create {
+        var mutableData = 0
+        {
+            println(mutableData++)
+        }
+    }
+
+    fun createCollection(): MutableList<Any> = create {
+        mutableListOf()
+    }
+
+    fun createFrozenRegularObject() = createRegularObject().freeze()
+    fun createFrozenLambda() = createLambda().freeze()
+    fun createFrozenCollection() = createCollection().freeze()
+
+    fun hasAliveObjects(): Boolean {
+        kotlin.native.internal.GC.collect()
+        return mustBeRemoved.any { it.get() != null }
+    }
+
+    private fun <T : Any> create(block: () -> T) = block()
+            .also { mustBeRemoved += WeakReference(it) }
+
+    private val mustBeRemoved = mutableListOf<WeakReference<*>>()
+}

backend.native/tests/framework/values/values.swift

@@ -808,6 +808,175 @@ func testWeakRefs0(frozen: Bool) throws {
     // try test4()
 }
 
+var falseFlag = false
+
+class TestSharedRefs {
+    private func testLambdaSimple() throws {
+        func getClosure() -> (() -> Void) {
+            let lambda = autoreleasepool {
+                SharedRefs().createLambda()
+            }
+            return { if falseFlag { lambda() } }
+        }
+
+        DispatchQueue.global().async(execute: getClosure())
+    }
+
+    private static func runInNewThread(initializeKotlinRuntime: Bool, block: @escaping () -> Void) {
+        class Closure {
+            static var currentBlock: (() -> Void)? = nil
+            static var initializeKotlinRuntime: Bool = false
+        }
+
+        Closure.currentBlock = block
+        Closure.initializeKotlinRuntime = initializeKotlinRuntime
+
+        var thread: pthread_t? = nil
+        let createCode = pthread_create(&thread, nil, { _ in
+            if Closure.initializeKotlinRuntime {
+                let ignore = SharedRefs() // Ensures that Kotlin runtime gets initialized.
+            }
+
+            Closure.currentBlock!()
+            Closure.currentBlock = nil
+
+            return nil
+        }, nil)
+        try! assertEquals(actual: createCode, expected: 0)
+
+        let joinCode = pthread_join(thread!, nil)
+        try! assertEquals(actual: joinCode, expected: 0)
+    }
+
+    private func runInNewThread(initializeKotlinRuntime: Bool, block: @escaping () -> Void) {
+        return TestSharedRefs.runInNewThread(initializeKotlinRuntime: initializeKotlinRuntime, block: block)
+    }
+
+    private func testObjectPartialRelease() {
+        let object = autoreleasepool { SharedRefs().createRegularObject() }
+        var objectVar: AnyObject? = object
+
+        runInNewThread(initializeKotlinRuntime: true) {
+            objectVar = nil
+        }
+    }
+
+    private func testRunRefCount<T>(
+        run: (@escaping () -> Void) -> Void,
+        createObject: @escaping (SharedRefs) -> T
+    ) throws {
+        let refs = SharedRefs()
+
+        var objectVar1: T? = autoreleasepool { createObject(refs) }
+        var objectVar2: T? = nil
+
+        try assertTrue(refs.hasAliveObjects())
+
+        run {
+            objectVar2 = objectVar1!
+            objectVar1 = nil
+        }
+
+        try assertTrue(refs.hasAliveObjects())
+
+        run {
+            objectVar2 = nil
+        }
+
+        try assertFalse(refs.hasAliveObjects())
+    }
+
+    private func testBackgroundRefCount<T>(createObject: @escaping (SharedRefs) -> T) throws {
+        try testRunRefCount(
+            run: { runInNewThread(initializeKotlinRuntime: false, block: $0) },
+            createObject: createObject
+        )
+
+        try testRunRefCount(
+            run: { runInNewThread(initializeKotlinRuntime: true, block: $0) },
+            createObject: createObject
+        )
+    }
+
+    private func testReferenceOutlivesThread(releaseWithKotlinRuntime: Bool) throws {
+        var objectVar: AnyObject? = nil
+        weak var objectWeakVar: AnyObject? = nil
+        var collection: AnyObject? = nil
+
+        runInNewThread(initializeKotlinRuntime: false) {
+            autoreleasepool {
+                let refs = SharedRefs()
+                collection = refs.createCollection()
+
+                let object = refs.createRegularObject()
+                objectVar = object
+                objectWeakVar = object
+
+                try! assertTrue(objectWeakVar === object)
+            }
+        }
+
+        runInNewThread(initializeKotlinRuntime: releaseWithKotlinRuntime) {
+            objectVar = nil
+            collection = nil
+            ValuesKt.gc()
+            try! assertTrue(objectWeakVar === nil)
+        }
+
+    }
+
+    private func testMoreWorkBeforeThreadExit() throws {
+        class Deinit {
+            static var object1: AnyObject? = nil
+            static var object2: AnyObject? = nil
+            static weak var weakVar2: AnyObject? = nil
+
+            deinit {
+                TestSharedRefs.runInNewThread(initializeKotlinRuntime: false) {
+                    Deinit.object2 = nil
+                }
+            }
+        }
+
+        runInNewThread(initializeKotlinRuntime: false) {
+            autoreleasepool {
+                let object1 = SharedRefs.MutableData()
+                Deinit.object1 = object1
+                setAssociatedObject(object: object1, value: Deinit())
+
+                let object2 = SharedRefs.MutableData()
+                Deinit.object2 = object2
+                Deinit.weakVar2 = object2
+            }
+
+            TestSharedRefs.runInNewThread(initializeKotlinRuntime: false) {
+                Deinit.object1 = nil
+            }
+        }
+
+        try assertTrue(Deinit.weakVar2 === nil)
+    }
+
+    func test() throws {
+        try testLambdaSimple()
+        try testObjectPartialRelease()
+
+        try testBackgroundRefCount(createObject: { $0.createLambda() })
+        try testBackgroundRefCount(createObject: { $0.createRegularObject() })
+        try testBackgroundRefCount(createObject: { $0.createCollection() })
+
+        try testBackgroundRefCount(createObject: { $0.createFrozenLambda() })
+        try testBackgroundRefCount(createObject: { $0.createFrozenRegularObject() })
+        try testBackgroundRefCount(createObject: { $0.createFrozenCollection() })
+
+        try testReferenceOutlivesThread(releaseWithKotlinRuntime: false)
+        try testReferenceOutlivesThread(releaseWithKotlinRuntime: true)
+        try testMoreWorkBeforeThreadExit()
+
+        usleep(300 * 1000)
+    }
+}
+
 // See https://github.com/JetBrains/kotlin-native/issues/2931
 func testGH2931() throws {
     for i in 0..<50000 {
@@ -877,6 +1046,7 @@ class ValuesTests : TestProvider {
             TestCase(name: "TestInvalidIdentifiers", method: withAutorelease(testInvalidIdentifiers)),
             TestCase(name: "TestDeprecation", method: withAutorelease(testDeprecation)),
             TestCase(name: "TestWeakRefs", method: withAutorelease(testWeakRefs)),
+            TestCase(name: "TestSharedRefs", method: withAutorelease(TestSharedRefs().test)),
             TestCase(name: "TestGH2931", method: withAutorelease(testGH2931)),
         ]
     }

runtime/src/main/cpp/Memory.cpp

@@ -466,31 +466,6 @@ struct MemoryState {
 #endif // COLLECT_STATISTIC
 };
 
-ObjHeader* KRefSharedHolder::ref() const {
-  verifyRefOwner();
-  return obj_;
-}
-
-void KRefSharedHolder::initRefOwner() {
-  RuntimeAssert(owner_ == nullptr, "Must be uninitialized");
-  owner_ = memoryState;
-}
-
-void KRefSharedHolder::verifyRefOwner() const {
-  // Note: checking for 'shareable()' and retrieving 'type_info()'
-  // are supposed to be correct even for unowned object.
-  if (owner_ != memoryState) {
-    // Initialized runtime is required to throw the exception below
-    // or to provide proper execution context for shared objects:
-    if (memoryState == nullptr) Kotlin_initRuntimeIfNeeded();
-    auto* container = obj_->container();
-    if (!isShareable(container)) {
-      // TODO: add some info about the owner.
-      ThrowIllegalObjectSharingException(obj_->type_info(), obj_);
-    }
-  }
-}
-
 namespace {
 
 #if TRACE_MEMORY
@@ -2626,10 +2601,6 @@ ArrayHeader* ArenaContainer::PlaceArray(const TypeInfo* type_info, uint32_t coun
 extern "C" {
 
 // Private memory interface.
-void AddRefFromAssociatedObject(const ObjHeader* object) {
-  addHeapRef(const_cast<ObjHeader*>(object));
-}
-
 void ReleaseHeapRefStrict(const ObjHeader* object) {
   releaseHeapRef<true>(const_cast<ObjHeader*>(object));
 }
@@ -2637,10 +2608,6 @@ void ReleaseHeapRefRelaxed(const ObjHeader* object) {
   releaseHeapRef<false>(const_cast<ObjHeader*>(object));
 }
 
-void ReleaseRefFromAssociatedObject(const ObjHeader* object) {
-  ReleaseHeapRef(object);
-}
-
 void DeinitInstanceBody(const TypeInfo* typeInfo, void* body) {
   deinitInstanceBody(typeInfo, body);
 }

runtime/src/main/cpp/Memory.h

@@ -568,34 +568,6 @@ class ExceptionObjHolder {
    ObjHeader* obj_;
 };
 
-class KRefSharedHolder {
- public:
-  inline ObjHeader** slotToInit() {
-    initRefOwner();
-    return &obj_;
-  }
-
-  inline void init(ObjHeader* obj) {
-    SetHeapRef(slotToInit(), obj);
-  }
-
-  ObjHeader* ref() const;
-
-  inline void dispose() {
-    verifyRefOwner();
-    ZeroHeapRef(&obj_);
-  }
-
- private:
-  typedef MemoryState* RefOwner;
-
-  ObjHeader* obj_;
-  RefOwner owner_;
-
-  void initRefOwner();
-  void verifyRefOwner() const;
-};
-
 class ForeignRefManager;
 typedef ForeignRefManager* ForeignRefContext;
 

runtime/src/main/cpp/MemoryPrivate.hpp

@@ -23,8 +23,6 @@ extern "C" {
 
 MODEL_VARIANTS(void, ReleaseHeapRef, const ObjHeader* object);
 
-void AddRefFromAssociatedObject(const ObjHeader* object) RUNTIME_NOTHROW;
-void ReleaseRefFromAssociatedObject(const ObjHeader* object) RUNTIME_NOTHROW;
 void DeinitInstanceBody(const TypeInfo* typeInfo, void* body);
 
 void Kotlin_ObjCExport_releaseAssociatedObject(void* associatedObject);

runtime/src/main/cpp/MemorySharedRefs.cpp

@@ -0,0 +1,119 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#include "Exceptions.h"
+#include "MemoryPrivate.hpp"
+#include "MemorySharedRefs.hpp"
+#include "Runtime.h"
+
+void KRefSharedHolder::initLocal(ObjHeader* obj) {
+  RuntimeAssert(obj != nullptr, "must not be null");
+  context_ = InitLocalForeignRef(obj);
+  obj_ = obj;
+}
+
+void KRefSharedHolder::init(ObjHeader* obj) {
+  RuntimeAssert(obj != nullptr, "must not be null");
+  context_ = InitForeignRef(obj);
+  obj_ = obj;
+}
+
+ObjHeader* KRefSharedHolder::ref() const {
+  ensureRefAccessible();
+  return obj_;
+}
+
+static inline void ensureForeignRefAccessible(ObjHeader* object, ForeignRefContext context) {
+  if (!Kotlin_hasRuntime()) {
+    // So the object is either unowned or shared.
+    // In the former case initialized runtime is required to throw the exception below,
+    // in the latter case -- to provide proper execution context for caller.
+    Kotlin_initRuntimeIfNeeded();
+  }
+
+  if (!IsForeignRefAccessible(object, context)) {
+    // TODO: add some info about the context.
+    // Note: retrieving 'type_info()' is supposed to be correct even for unowned object.
+    ThrowIllegalObjectSharingException(object->type_info(), object);
+  }
+}
+
+void KRefSharedHolder::dispose() const {
+  if (obj_ == nullptr) {
+    // To handle the case when it is not initialized. See [KotlinMutableSet/Dictionary dealloc].
+    return;
+  }
+
+  DeinitForeignRef(obj_, context_);
+}
+
+void KRefSharedHolder::ensureRefAccessible() const {
+  ensureForeignRefAccessible(obj_, context_);
+}
+
+void BackRefFromAssociatedObject::initAndAddRef(ObjHeader* obj) {
+  RuntimeAssert(obj != nullptr, "must not be null");
+  obj_ = obj;
+
+  // Generally a specialized addRef below:
+  context_ = InitForeignRef(obj);
+  refCount = 1;
+}
+
+void BackRefFromAssociatedObject::addRef() {
+  if (atomicAdd(&refCount, 1) == 1) {
+    // There are no references to the associated object itself, so Kotlin object is being passed from Kotlin,
+    // and it is owned therefore.
+    ensureRefAccessible(); // TODO: consider removing explicit verification.
+
+    // Foreign reference has already been deinitialized (see [releaseRef]).
+    // Create a new one:
+    context_ = InitForeignRef(obj_);
+  }
+}
+
+bool BackRefFromAssociatedObject::tryAddRef() {
+  ObjHeader* obj = this->ref();
+  if (!obj->has_meta_object()) {
+    // Then object is being deallocated.
+    return false;
+  } else {
+    this->addRef();
+    return true;
+  }
+}
+
+void BackRefFromAssociatedObject::releaseRef() {
+  ForeignRefContext context = context_;
+  if (atomicAdd(&refCount, -1) == 0) {
+    // Note: by this moment "subsequent" addRef may have already happened and patched context_.
+    // So use the value loaded before refCount update:
+    DeinitForeignRef(obj_, context);
+    // From this moment [context] is generally a dangling pointer.
+    // This is handled in [IsForeignRefAccessible] and [addRef].
+  }
+}
+
+void BackRefFromAssociatedObject::ensureRefAccessible() const {
+  ensureForeignRefAccessible(obj_, context_);
+}
+
+extern "C" {
+RUNTIME_NOTHROW void KRefSharedHolder_initLocal(KRefSharedHolder* holder, ObjHeader* obj) {
+  holder->initLocal(obj);
+}
+
+RUNTIME_NOTHROW void KRefSharedHolder_init(KRefSharedHolder* holder, ObjHeader* obj) {
+  holder->init(obj);
+}
+
+RUNTIME_NOTHROW void KRefSharedHolder_dispose(const KRefSharedHolder* holder) {
+  holder->dispose();
+}
+
+ObjHeader* KRefSharedHolder_ref(const KRefSharedHolder* holder) {
+  return holder->ref();
+}
+} // extern "C"

runtime/src/main/cpp/MemorySharedRefs.hpp

@@ -0,0 +1,55 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the LICENSE file.
+ */
+
+#ifndef RUNTIME_MEMORYSHAREDREFS_HPP
+#define RUNTIME_MEMORYSHAREDREFS_HPP
+
+#include "Memory.h"
+
+class KRefSharedHolder {
+ public:
+  void initLocal(ObjHeader* obj);
+
+  void init(ObjHeader* obj);
+
+  ObjHeader* ref() const;
+
+  void dispose() const;
+
+ private:
+  ObjHeader* obj_;
+  ForeignRefContext context_;
+
+  void ensureRefAccessible() const;
+};
+
+class BackRefFromAssociatedObject {
+ public:
+  void initAndAddRef(ObjHeader* obj);
+
+  void addRef();
+
+  bool tryAddRef();
+
+  void releaseRef();
+
+  ObjHeader* ref() const {
+    ensureRefAccessible();
+    return obj_;
+  }
+
+  inline bool permanent() const {
+    return obj_->permanent(); // Safe to query from any thread.
+  }
+
+ private:
+  ObjHeader* obj_;
+  ForeignRefContext context_;
+  volatile int refCount;
+
+  void ensureRefAccessible() const;
+};
+
+#endif // RUNTIME_MEMORYSHAREDREFS_HPP

runtime/src/main/cpp/ObjCExport.mm

@@ -14,8 +14,10 @@
  * limitations under the License.
  */
 
+#import "Atomic.h"
 #import "Types.h"
 #import "Memory.h"
+#import "MemorySharedRefs.hpp"
 #include "Natives.h"
 
 #if KONAN_OBJC_INTEROP
@@ -140,7 +142,7 @@ extern "C" id objc_autoreleaseReturnValue(id self);
 static void initializeObjCExport();
 
 @implementation KotlinBase {
-  KRefSharedHolder refHolder;
+  BackRefFromAssociatedObject refHolder;
 }
 
 -(KRef)toKotlin:(KRef*)OBJ_RESULT {
@@ -173,14 +175,14 @@ static void initializeObjCExport();
   }
   ObjHolder holder;
   AllocInstanceWithAssociatedObject(typeInfo, result, holder.slot());
-  UpdateHeapRef(result->refHolder.slotToInit(), holder.obj());
+  result->refHolder.initAndAddRef(holder.obj());
   return result;
 }
 
 +(instancetype)createWrapper:(ObjHeader*)obj {
   KotlinBase* candidate = [super allocWithZone:nil];
   // TODO: should we call NSObject.init ?
-  candidate->refHolder.init(obj);
+  candidate->refHolder.initAndAddRef(obj);
 
   if (!obj->permanent()) { // TODO: permanent objects should probably be supported as custom types.
     if (!obj->container()->shareable()) {
@@ -188,7 +190,7 @@ static void initializeObjCExport();
     } else {
       id old = AtomicCompareAndSwapAssociatedObject(obj, nullptr, candidate);
       if (old != nullptr) {
-        candidate->refHolder.dispose();
+        candidate->refHolder.releaseRef();
         [candidate releaseAsAssociatedObject];
         return objc_retainAutoreleaseReturnValue(old);
       }
@@ -199,35 +201,27 @@ static void initializeObjCExport();
 }
 
 -(instancetype)retain {
-  ObjHeader* obj = refHolder.ref();
-  if (obj->permanent()) { // TODO: consider storing `isPermanent` to self field.
+  if (refHolder.permanent()) { // TODO: consider storing `isPermanent` to self field.
     [super retain];
   } else {
-    AddRefFromAssociatedObject(obj);
+    refHolder.addRef();
   }
   return self;
 }
 
 -(BOOL)_tryRetain {
-  ObjHeader* obj = refHolder.ref();
-  if (obj->permanent()) {
+  if (refHolder.permanent()) {
     return [super _tryRetain];
-  } else if (!obj->has_meta_object()) {
-    // Then object is being deallocated;
-    // return `NO` as required by _tryRetain semantics:
-    return NO;
   } else {
-    AddRefFromAssociatedObject(obj);
-    return YES;
+    return refHolder.tryAddRef();
   }
 }
 
 -(oneway void)release {
-  ObjHeader* obj = refHolder.ref();
-  if (obj->permanent()) {
+  if (refHolder.permanent()) {
     [super release];
   } else {
-    ReleaseRefFromAssociatedObject(obj);
+    refHolder.releaseRef();
   }
 }
 

runtime/src/main/cpp/ObjCExportCollections.mm

@@ -15,6 +15,7 @@
  */
 
 #import "Memory.h"
+#import "MemorySharedRefs.hpp"
 #import "Types.h"
 
 #if KONAN_OBJC_INTEROP
@@ -400,6 +401,9 @@ static inline id KSet_getElement(KRef set, id object) {
 // ?
 
 -(void)dealloc {
+  // Note: since setHolder initialization is not performed directly with alloc,
+  // it is possible that it wasn't initialized properly.
+  // Fortunately setHolder.dispose() handles the zero-initialized case too.
   setHolder.dispose();
   [super dealloc];
 }
@@ -508,6 +512,9 @@ static inline id KMap_get(KRef map, id aKey) {
 }
 
 -(void)dealloc {
+  // Note: since mapHolder initialization is not performed directly with alloc,
+  // it is possible that it wasn't initialized properly.
+  // Fortunately mapHolder.dispose() handles the zero-initialized case too.
   mapHolder.dispose();
   [super dealloc];
 }

runtime/src/main/cpp/ObjCInteropUtils.mm

@@ -22,6 +22,7 @@
 #import <Foundation/NSException.h>
 #import <Foundation/NSString.h>
 #import "Memory.h"
+#import "MemorySharedRefs.hpp"
 
 namespace {
   Class nsStringClass = nullptr;

runtime/src/main/cpp/Runtime.cpp

@@ -176,6 +176,10 @@ RuntimeState* RUNTIME_USED Kotlin_getRuntime() {
   return ::runtimeState;
 }
 
+bool Kotlin_hasRuntime() {
+  return isValidRuntime();
+}
+
 void CheckIsMainThread() {
   if (!isMainThread)
     ThrowIncorrectDereferenceException();

runtime/src/main/cpp/Runtime.h

@@ -42,9 +42,11 @@ RuntimeState* RUNTIME_USED Kotlin_suspendRuntime();
 // the runtime. After resume, current thread could be used for executing Kotlin code.
 void RUNTIME_USED Kotlin_resumeRuntime(RuntimeState*);
 
-// Gets currently active runtime, nullptr if no runtime is currently available.
+// Gets currently active runtime, fails if no runtime is currently available.
 RuntimeState* RUNTIME_USED Kotlin_getRuntime();
 
+bool Kotlin_hasRuntime();
+
 // Appends given node to an initializer list.
 void AppendToInitializersTail(struct InitNode*);