Native: improve ObjCExport thread state switching

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

@@ -133,6 +133,7 @@ internal inline fun<R> generateFunction(codegen: CodeGenerator,
             codegen,
             startLocation,
             endLocation,
+            switchToRunnable = function.origin == CBridgeOrigin.C_TO_KOTLIN_BRIDGE,
             function)
     try {
         generateFunctionBody(functionGenerationContext, code)
@@ -148,8 +149,15 @@ internal inline fun<R> generateFunction(codegen: CodeGenerator,
 
 internal inline fun<R> generateFunction(codegen: CodeGenerator, function: LLVMValueRef,
                                         startLocation: LocationInfo? = null, endLocation: LocationInfo? = null,
+                                        switchToRunnable: Boolean = false,
                                         code:FunctionGenerationContext.(FunctionGenerationContext) -> R) {
-    val functionGenerationContext = FunctionGenerationContext(function, codegen, startLocation, endLocation)
+    val functionGenerationContext = FunctionGenerationContext(
+            function,
+            codegen,
+            startLocation,
+            endLocation,
+            switchToRunnable = switchToRunnable
+    )
     try {
         generateFunctionBody(functionGenerationContext, code)
     } finally {
@@ -161,6 +169,7 @@ internal inline fun generateFunction(
         codegen: CodeGenerator,
         functionType: LLVMTypeRef,
         name: String,
+        switchToRunnable: Boolean = false,
         block: FunctionGenerationContext.(FunctionGenerationContext) -> Unit
 ): LLVMValueRef {
     val function = addLlvmFunctionWithDefaultAttributes(
@@ -169,7 +178,7 @@ internal inline fun generateFunction(
             name,
             functionType
     )
-    generateFunction(codegen, function, startLocation = null, endLocation = null, code = block)
+    generateFunction(codegen, function, startLocation = null, endLocation = null, switchToRunnable = switchToRunnable, code = block)
     return function
 }
 
@@ -179,7 +188,7 @@ internal inline fun <R> generateFunctionNoRuntime(
         function: LLVMValueRef,
         code: FunctionGenerationContext.(FunctionGenerationContext) -> R,
 ) {
-    val functionGenerationContext = FunctionGenerationContext(function, codegen, null, null)
+    val functionGenerationContext = FunctionGenerationContext(function, codegen, null, null, switchToRunnable = false)
     try {
         functionGenerationContext.forbidRuntime = true
         require(!functionGenerationContext.isObjectType(functionGenerationContext.returnType!!)) {
@@ -367,6 +376,7 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
                                          val codegen: CodeGenerator,
                                          startLocation: LocationInfo?,
                                          private val endLocation: LocationInfo?,
+                                         private val switchToRunnable: Boolean,
                                          internal val irFunction: IrFunction? = null): ContextUtils {
 
     override val context = codegen.context
@@ -1301,8 +1311,7 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
             returnSlot = LLVMGetParam(function, numParameters(function.type) - 1)
         }
 
-        if (context.memoryModel == MemoryModel.EXPERIMENTAL &&
+        if (context.memoryModel == MemoryModel.EXPERIMENTAL && switchToRunnable) {
-                irFunction?.origin == CBridgeOrigin.C_TO_KOTLIN_BRIDGE) {
             check(!forbidRuntime) { "Attempt to switch the thread state when runtime is forbidden" }
             positionAtEnd(prologueBb)
             // TODO: Do we need to do it for every c->kotlin bridge?
@@ -1320,10 +1329,9 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
 
     internal fun epilogue() {
         appendingTo(prologueBb) {
-            if (needsRuntimeInit) {
+            if (needsRuntimeInit && context.config.memoryModel != MemoryModel.EXPERIMENTAL) {
                 check(!forbidRuntime) { "Attempt to init runtime where runtime usage is forbidden" }
                 call(context.llvm.initRuntimeIfNeeded, emptyList())
-                // TODO: Do we also need to switch to runnable mode?
             }
             val slots = if (needSlotsPhi)
                 LLVMBuildArrayAlloca(builder, kObjHeaderPtr, Int32(slotCount).llvm, "")!!
@@ -1391,7 +1399,7 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
         val shouldHaveCleanupLandingpad = forwardingForeignExceptionsTerminatedWith != null ||
                 needSlots ||
                 !stackLocalsManager.isEmpty() ||
-                (context.memoryModel == MemoryModel.EXPERIMENTAL && irFunction?.origin == CBridgeOrigin.C_TO_KOTLIN_BRIDGE)
+                (context.memoryModel == MemoryModel.EXPERIMENTAL && switchToRunnable)
 
         if (shouldHaveCleanupLandingpad) {
             appendingTo(cleanupLandingpad) {
@@ -1455,7 +1463,7 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
             if (!forbidRuntime) {
                 call(safePointFunction, emptyList())
             }
-            if (irFunction?.origin == CBridgeOrigin.C_TO_KOTLIN_BRIDGE) {
+            if (switchToRunnable) {
                 check(!forbidRuntime) { "Generating a bridge when runtime is forbidden" }
                 switchThreadState(Native)
             }

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

@@ -52,7 +52,7 @@ internal fun ObjCExportCodeGeneratorBase.generateBlockToKotlinFunctionConverter(
         )
 
         val invoke = loadBlockInvoke(blockPtr, bridge)
-        val result = callFromBridge(invoke, listOf(blockPtr) + args)
+        val result = callFromBridge(invoke, listOf(blockPtr) + args, toNative = true)
 
         val kotlinResult = if (bridge.returnsVoid) {
             theUnitInstanceRef.llvm
@@ -154,7 +154,8 @@ internal class BlockGenerator(private val codegen: CodeGenerator) {
     val disposeHelper = generateFunction(
             codegen,
             functionType(voidType, false, int8TypePtr),
-            "blockDisposeHelper"
+            "blockDisposeHelper",
+            switchToRunnable = true
     ) {
         val blockPtr = bitcast(pointerType(blockLiteralType), param(0))
         val refHolder = structGep(blockPtr, 1)
@@ -234,7 +235,7 @@ internal class BlockGenerator(private val codegen: CodeGenerator) {
             invokeName: String,
             genBody: FunctionGenerationContext.(LLVMValueRef, List<LLVMValueRef>) -> Unit
     ): ConstPointer {
-        val result = generateFunction(codegen, blockType.blockInvokeLlvmType, invokeName) {
+        val result = generateFunction(codegen, blockType.blockInvokeLlvmType, invokeName, switchToRunnable = true) {
             val blockPtr = bitcast(pointerType(blockLiteralType), param(0))
             val kotlinObject = call(
                     context.llvm.kRefSharedHolderRef,

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

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.*
+import org.jetbrains.kotlin.konan.ForeignExceptionMode
 import org.jetbrains.kotlin.konan.target.AppleConfigurables
 import org.jetbrains.kotlin.konan.target.LinkerOutputKind
 import org.jetbrains.kotlin.name.Name
@@ -67,13 +68,31 @@ internal open class ObjCExportCodeGeneratorBase(codegen: CodeGenerator) : ObjCCo
     fun FunctionGenerationContext.callFromBridge(
             function: LLVMValueRef,
             args: List<LLVMValueRef>,
-            resultLifetime: Lifetime = Lifetime.IRRELEVANT
+            resultLifetime: Lifetime = Lifetime.IRRELEVANT,
+            toNative: Boolean = false
     ): LLVMValueRef {
 
         // TODO: it is required only for Kotlin-to-Objective-C bridges.
         this.forwardingForeignExceptionsTerminatedWith = objcTerminate
 
-        return call(function, args, resultLifetime, ExceptionHandler.Caller)
+        val switchStateToNative = toNative && context.config.memoryModel == MemoryModel.EXPERIMENTAL
+        val exceptionHandler: ExceptionHandler
+
+        if (switchStateToNative) {
+            switchThreadState(ThreadState.Native)
+            // Note: this is suboptimal. We should forbid Kotlin exceptions thrown from native code, and use simple fatal handler here.
+            exceptionHandler = filteringExceptionHandler(ExceptionHandler.Caller, ForeignExceptionMode.default, switchThreadState = true)
+        } else {
+            exceptionHandler = ExceptionHandler.Caller
+        }
+
+        val result = call(function, args, resultLifetime, exceptionHandler)
+
+        if (switchStateToNative) {
+            switchThreadState(ThreadState.Runnable)
+        }
+
+        return result
     }
 
     fun FunctionGenerationContext.kotlinReferenceToObjC(value: LLVMValueRef) =
@@ -134,7 +153,8 @@ internal class ObjCExportCodeGenerator(
             returnType: LLVMTypeRef,
             receiver: LLVMValueRef,
             selector: String,
-            vararg args: LLVMValueRef
+            switchToNative: Boolean,
+            vararg args: LLVMValueRef,
     ): LLVMValueRef {
 
         val objcMsgSendType = functionType(
@@ -143,7 +163,7 @@ internal class ObjCExportCodeGenerator(
                 listOf(int8TypePtr, int8TypePtr) + args.map { it.type }
         )
 
-        return callFromBridge(msgSender(objcMsgSendType), listOf(receiver, genSelector(selector)) + args)
+        return callFromBridge(msgSender(objcMsgSendType), listOf(receiver, genSelector(selector)) + args, toNative = switchToNative)
     }
 
     fun FunctionGenerationContext.kotlinToObjC(
@@ -621,7 +641,8 @@ private fun ObjCExportCodeGenerator.emitBoxConverter(
         val value = kotlinToObjC(kotlinValue, objCValueType)
 
         val nsNumberSubclass = genGetLinkedClass(namer.numberBoxName(boxClass.classId!!).binaryName)
-        ret(genSendMessage(int8TypePtr, nsNumberSubclass, nsNumberFactorySelector, value))
+        val switchToNative = false // We consider these methods fast enough.
+        ret(genSendMessage(int8TypePtr, nsNumberSubclass, nsNumberFactorySelector, switchToNative, value))
     }
 
     LLVMSetLinkage(converter, LLVMLinkage.LLVMPrivateLinkage)
@@ -748,7 +769,7 @@ private inline fun ObjCExportCodeGenerator.generateObjCImpBy(
         null
     }
 
-    generateFunction(codegen, result, startLocation = location, endLocation = location) {
+    generateFunction(codegen, result, startLocation = location, endLocation = location, switchToRunnable = true) {
         genBody()
     }
 
@@ -861,7 +882,7 @@ private fun ObjCExportCodeGenerator.generateObjCImp(
                 is MethodBridge.ReturnValue.WithError.ZeroForError -> {
                     if (returnType.successBridge == MethodBridge.ReturnValue.Instance.InitResult) {
                         // Release init receiver, as required by convention.
-                        callFromBridge(objcRelease, listOf(param(0)))
+                        callFromBridge(objcRelease, listOf(param(0)), toNative = true)
                     }
                     Zero(returnType.objCType(context)).llvm
                 }
@@ -1040,7 +1061,7 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
             }
         }
 
-        val targetResult = callFromBridge(objcMsgSend, objCArgs)
+        val targetResult = callFromBridge(objcMsgSend, objCArgs, toNative = true)
 
         assert(baseMethod.symbol !is IrConstructorSymbol)
 
@@ -1496,7 +1517,7 @@ private inline fun ObjCExportCodeGenerator.generateObjCToKotlinSyntheticGetter(
             MethodBridgeReceiver.Static, valueParameters = emptyList()
     )
 
-    val imp = generateFunction(codegen, objCFunctionType(context, methodBridge), "objc2kotlin") {
+    val imp = generateFunction(codegen, objCFunctionType(context, methodBridge), "objc2kotlin", switchToRunnable = true) {
         block()
     }
 
@@ -1517,6 +1538,7 @@ private fun ObjCExportCodeGenerator.objCToKotlinMethodAdapter(
 
 private fun ObjCExportCodeGenerator.createUnitInstanceAdapter(selector: String) =
         generateObjCToKotlinSyntheticGetter(selector) {
+            // Note: generateObjCToKotlinSyntheticGetter switches to Runnable, which is probably not required here and thus suboptimal.
             initRuntimeIfNeeded() // For instance methods it gets called when allocating.
 
             ret(callFromBridge(context.llvm.Kotlin_ObjCExport_convertUnit, listOf(codegen.theUnitInstanceRef.llvm)))

kotlin-native/runtime/src/legacymm/cpp/Memory.cpp

@@ -3746,4 +3746,8 @@ kotlin::ThreadState kotlin::GetThreadState(MemoryState* thread) noexcept {
     return ThreadState::kRunnable;
 }
 
+ALWAYS_INLINE kotlin::CalledFromNativeGuard::CalledFromNativeGuard() noexcept {
+    // no-op, used by the new MM only.
+}
+
 const bool kotlin::kSupportsMultipleMutators = true;

kotlin-native/runtime/src/main/cpp/Memory.h

@@ -442,6 +442,20 @@ private:
     bool reentrant_;
 };
 
+// Scopely sets the kRunnable thread state for the current thread,
+// and initializes runtime if needed for new MM.
+// No-op for old GC.
+class CalledFromNativeGuard final : private Pinned {
+public:
+    ALWAYS_INLINE CalledFromNativeGuard() noexcept;
+
+    ~CalledFromNativeGuard() noexcept {
+        SwitchThreadState(thread_, ThreadState::kNative);
+    }
+private:
+    MemoryState* thread_;
+};
+
 template <ThreadState state, typename R, typename... Args>
 ALWAYS_INLINE inline R CallWithThreadState(R(*function)(Args...), Args... args) {
     ThreadStateGuard guard(state);

kotlin-native/runtime/src/main/cpp/MemorySharedRefs.cpp

@@ -19,7 +19,7 @@ inline bool isForeignRefAccessible(ObjHeader* object, ForeignRefContext context)
     // If runtime has not been initialized on this thread, then the object is either unowned or shared.
     // In the former case initialized runtime is required to throw exceptions
     // in the latter case -- to provide proper execution context for caller.
-    // TODO: Can this be called in uninitialized state in the new MM?
+    // TODO: this probably can't be called in uninitialized state in the new MM.
     Kotlin_initRuntimeIfNeeded();
 
     return IsForeignRefAccessible(object, context);
@@ -85,6 +85,7 @@ void KRefSharedHolder::init(ObjHeader* obj) {
 
 template <ErrorPolicy errorPolicy>
 ObjHeader* KRefSharedHolder::ref() const {
+  kotlin::AssertThreadState(kotlin::ThreadState::kRunnable);
   if (!ensureRefAccessible<errorPolicy>(obj_, context_)) {
     return nullptr;
   }
@@ -129,11 +130,15 @@ void BackRefFromAssociatedObject::addRef() {
 
     // There are no references to the associated object itself, so Kotlin object is being passed from Kotlin,
     // and it is owned therefore.
+    kotlin::AssertThreadState(kotlin::ThreadState::kRunnable);
     ensureRefAccessible<errorPolicy>(obj_, context_); // TODO: consider removing explicit verification.
 
     // Foreign reference has already been deinitialized (see [releaseRef]).
     // Create a new one:
     context_ = InitForeignRef(obj_);
+  } else {
+    // Can be called both from Native state (if ObjC or Swift code adds RC)
+    // and from Runnable state (Kotlin_ObjCExport_refToObjC).
   }
 }
 
@@ -143,6 +148,7 @@ template void BackRefFromAssociatedObject::addRef<ErrorPolicy::kTerminate>();
 template <ErrorPolicy errorPolicy>
 bool BackRefFromAssociatedObject::tryAddRef() {
   static_assert(errorPolicy != ErrorPolicy::kDefaultValue, "Cannot use default return value here");
+  kotlin::CalledFromNativeGuard guard;
 
   if (obj_ == nullptr) return false; // E.g. after [detach].
 
@@ -180,6 +186,8 @@ void BackRefFromAssociatedObject::releaseRef() {
   if (atomicAdd(&refCount, -1) == 0) {
     if (obj_ == nullptr) return; // E.g. after [detach].
 
+    kotlin::CalledFromNativeGuard guard;
+
     // Note: by this moment "subsequent" addRef may have already happened and patched context_.
     // So use the value loaded before refCount update:
     DeinitForeignRef(obj_, context);
@@ -200,6 +208,7 @@ ALWAYS_INLINE void BackRefFromAssociatedObject::assertDetached() {
 
 template <ErrorPolicy errorPolicy>
 ObjHeader* BackRefFromAssociatedObject::ref() const {
+  kotlin::AssertThreadState(kotlin::ThreadState::kRunnable);
   RuntimeAssert(obj_ != nullptr, "no valid Kotlin object found");
 
   if (!ensureRefAccessible<errorPolicy>(obj_, context_)) {

kotlin-native/runtime/src/main/cpp/MemorySharedRefs.hpp

@@ -30,6 +30,11 @@ class KRefSharedHolder {
 
   void dispose() const;
 
+  void disposeFromNative() const {
+    kotlin::CalledFromNativeGuard guard;
+    dispose();
+  }
+
   OBJ_GETTER0(describe) const;
 
  private:

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

@@ -126,24 +126,32 @@ extern "C" id objc_retainAutoreleaseReturnValue(id self);
 namespace {
 
 ALWAYS_INLINE void send_releaseAsAssociatedObject(void* associatedObject, ReleaseMode mode) {
-  if (associatedObject != nullptr) {
+  auto msgSend = reinterpret_cast<void (*)(void* self, SEL cmd, ReleaseMode mode)>(&objc_msgSend);
-    auto msgSend = reinterpret_cast<void (*)(void* self, SEL cmd, ReleaseMode mode)>(&objc_msgSend);
+  msgSend(associatedObject, Kotlin_ObjCExport_releaseAsAssociatedObjectSelector, mode);
-    msgSend(associatedObject, Kotlin_ObjCExport_releaseAsAssociatedObjectSelector, mode);
-  }
 }
 
 } // namespace
 
 extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_releaseAssociatedObject(void* associatedObject) {
-  send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kRelease);
+  if (associatedObject != nullptr) {
+    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
+    send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kRelease);
+  }
 }
 
 extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_detachAndReleaseAssociatedObject(void* associatedObject) {
-  send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kDetachAndRelease);
+  if (associatedObject != nullptr) {
+    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
+    send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kDetachAndRelease);
+  }
 }
 
 extern "C" ALWAYS_INLINE void Kotlin_ObjCExport_detachAssociatedObject(void* associatedObject) {
-  send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kDetach);
+  if (associatedObject != nullptr) {
+    // Switching to Native state is not required, because detach is fast and can't call user code.
+    // Also switching is not possible, because this is called from GC.
+    send_releaseAsAssociatedObject(associatedObject, ReleaseMode::kDetach);
+  }
 }
 
 extern "C" id Kotlin_ObjCExport_convertUnit(ObjHeader* unitInstance) {
@@ -498,6 +506,8 @@ static id Kotlin_ObjCExport_refToObjC_slowpath(ObjHeader* obj);
 
 template <bool retainAutorelease>
 static ALWAYS_INLINE id Kotlin_ObjCExport_refToObjCImpl(ObjHeader* obj) {
+  kotlin::AssertThreadState(kotlin::ThreadState::kRunnable);
+
   if (obj == nullptr) return nullptr;
 
   id associatedObject = GetAssociatedObject(obj);
@@ -536,6 +546,8 @@ extern "C" OBJ_GETTER(Kotlin_Interop_CreateObjCObjectHolder, id obj) {
 }
 
 extern "C" OBJ_GETTER(Kotlin_ObjCExport_refFromObjC, id obj) {
+  kotlin::AssertThreadState(kotlin::ThreadState::kRunnable);
+
   if (obj == nullptr) RETURN_OBJ(nullptr);
   auto msgSend = reinterpret_cast<ObjHeader* (*)(id self, SEL cmd, ObjHeader** slot)>(&objc_msgSend);
   RETURN_RESULT_OF(msgSend, obj, Kotlin_ObjCExport_toKotlinSelector);

kotlin-native/runtime/src/main/cpp/ObjCExportCollections.h

@@ -46,8 +46,8 @@ static inline OBJ_GETTER(refFromObjCOrNSNull, id obj) {
 }
 
 static inline OBJ_GETTER(invokeAndAssociate, KRef (*func)(KRef* result), id obj) {
+  // TODO: this probably can't be called in uninitialized state in the new MM.
   Kotlin_initRuntimeIfNeeded();
-  // TODO: Does this need a switch to runnable state?
 
   KRef kotlinObj = func(OBJ_RESULT);
 

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

@@ -358,6 +358,7 @@ konan::AutoreleasePool::AutoreleasePool()
   : handle(objc_autoreleasePoolPush()) {}
 
 konan::AutoreleasePool::~AutoreleasePool() {
+  kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
   objc_autoreleasePoolPop(handle);
 }
 
@@ -366,6 +367,7 @@ void* Kotlin_objc_autoreleasePoolPush() {
 }
 
 void Kotlin_objc_autoreleasePoolPop(void* ptr) {
+  kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
   objc_autoreleasePoolPop(ptr);
 }
 

kotlin-native/runtime/src/main/cpp/Runtime.cpp

@@ -149,6 +149,7 @@ RuntimeState* initRuntime() {
 }
 
 void deinitRuntime(RuntimeState* state, bool destroyRuntime) {
+  AssertThreadState(state->memoryState, kotlin::ThreadState::kRunnable);
   RuntimeAssert(state->status == RuntimeStatus::kRunning, "Runtime must be in the running state");
   state->status = RuntimeStatus::kDestroying;
   // This may be called after TLS is zeroed out, so ::runtimeState and ::memoryState in Memory cannot be trusted.

kotlin-native/runtime/src/mm/cpp/Memory.cpp

@@ -465,18 +465,21 @@ extern "C" void EnsureNeverFrozen(ObjHeader* obj) {
 }
 
 extern "C" ForeignRefContext InitLocalForeignRef(ObjHeader* object) {
+    AssertThreadState(ThreadState::kRunnable);
     // TODO: Remove when legacy MM is gone.
     // Nothing to do.
     return nullptr;
 }
 
 extern "C" ForeignRefContext InitForeignRef(ObjHeader* object) {
+    AssertThreadState(ThreadState::kRunnable);
     auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
     auto* node = mm::StableRefRegistry::Instance().RegisterStableRef(threadData, object);
     return ToForeignRefManager(node);
 }
 
 extern "C" void DeinitForeignRef(ObjHeader* object, ForeignRefContext context) {
+    AssertThreadState(ThreadState::kRunnable);
     RuntimeAssert(context != nullptr, "DeinitForeignRef must not be called for InitLocalForeignRef");
     auto* threadData = mm::ThreadRegistry::Instance().CurrentThreadData();
     auto* node = FromForeignRefManager(context);
@@ -529,4 +532,10 @@ MemoryState* kotlin::mm::GetMemoryState() {
     return ToMemoryState(ThreadRegistry::Instance().CurrentThreadDataNode());
 }
 
+ALWAYS_INLINE kotlin::CalledFromNativeGuard::CalledFromNativeGuard() noexcept {
+    Kotlin_initRuntimeIfNeeded();
+    thread_ = mm::GetMemoryState();
+    SwitchThreadState(thread_, ThreadState::kRunnable);
+}
+
 const bool kotlin::kSupportsMultipleMutators = kotlin::gc::kSupportsMultipleMutators;

kotlin-native/runtime/src/objc/cpp/ObjCExportClasses.mm

@@ -61,7 +61,7 @@ static void injectToRuntime();
 
 +(instancetype)allocWithZone:(NSZone*)zone {
   Kotlin_initRuntimeIfNeeded();
-  // TODO: Does this need a switch to runnable state?
+  kotlin::ThreadStateGuard guard(kotlin::ThreadState::kRunnable);
 
   KotlinBase* result = [super allocWithZone:zone];
 
@@ -86,6 +86,8 @@ static void injectToRuntime();
 }
 
 +(instancetype)createWrapper:(ObjHeader*)obj {
+  kotlin::AssertThreadState(kotlin::ThreadState::kRunnable);
+
   KotlinBase* candidate = [super allocWithZone:nil];
   // TODO: should we call NSObject.init ?
   candidate->refHolder.initAndAddRef(obj);
@@ -97,8 +99,11 @@ static void injectToRuntime();
     } else {
       id old = AtomicCompareAndSwapAssociatedObject(obj, nullptr, candidate);
       if (old != nullptr) {
-        candidate->refHolder.releaseRef();
+        {
-        [candidate releaseAsAssociatedObject:ReleaseMode::kDetachAndRelease];
+          kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
+          candidate->refHolder.releaseRef();
+          [candidate releaseAsAssociatedObject:ReleaseMode::kDetachAndRelease];
+        }
         return objc_retainAutoreleaseReturnValue(old);
       }
     }

kotlin-native/runtime/src/objc/cpp/ObjCExportCollections.mm

@@ -156,7 +156,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }
 
 -(void)dealloc {
-  iteratorHolder.dispose();
+  iteratorHolder.disposeFromNative();
   [super dealloc];
 }
 
@@ -167,6 +167,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }
 
 - (id)nextObject {
+  kotlin::CalledFromNativeGuard guard;
   KRef iterator = iteratorHolder.ref<ErrorPolicy::kTerminate>();
   if (Kotlin_Iterator_hasNext(iterator)) {
     ObjHolder holder;
@@ -185,7 +186,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }
 
 -(void)dealloc {
-  listHolder.dispose();
+  listHolder.disposeFromNative();
   [super dealloc];
 }
 
@@ -200,12 +201,14 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }
 
 -(id)objectAtIndex:(NSUInteger)index {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder kotlinValueHolder;
   KRef kotlinValue = Kotlin_List_get(listHolder.ref<ErrorPolicy::kTerminate>(), index, kotlinValueHolder.slot());
   return refToObjCOrNSNull(kotlinValue);
 }
 
 -(NSUInteger)count {
+  kotlin::CalledFromNativeGuard guard;
   return Kotlin_Collection_getSize(listHolder.ref<ErrorPolicy::kTerminate>());
 }
 
@@ -219,7 +222,7 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }
 
 -(void)dealloc {
-  listHolder.dispose();
+  listHolder.disposeFromNative();
   [super dealloc];
 }
 
@@ -234,35 +237,42 @@ static inline KInt objCIndexToKotlinOrThrow(NSUInteger index) {
 }
 
 -(id)objectAtIndex:(NSUInteger)index {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder kotlinValueHolder;
   KRef kotlinValue = Kotlin_List_get(listHolder.ref<ErrorPolicy::kTerminate>(), index, kotlinValueHolder.slot());
   return refToObjCOrNSNull(kotlinValue);
 }
 
 -(NSUInteger)count {
+  kotlin::CalledFromNativeGuard guard;
   return Kotlin_Collection_getSize(listHolder.ref<ErrorPolicy::kTerminate>());
 }
 
 - (void)insertObject:(id)anObject atIndex:(NSUInteger)index {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   KRef kotlinObject = refFromObjCOrNSNull(anObject, holder.slot());
   Kotlin_MutableList_addObjectAtIndex(listHolder.ref<ErrorPolicy::kTerminate>(), objCIndexToKotlinOrThrow(index), kotlinObject);
 }
 
 - (void)removeObjectAtIndex:(NSUInteger)index {
+  kotlin::CalledFromNativeGuard guard;
   Kotlin_MutableList_removeObjectAtIndex(listHolder.ref<ErrorPolicy::kTerminate>(), objCIndexToKotlinOrThrow(index));
 }
 
 - (void)addObject:(id)anObject {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   Kotlin_MutableCollection_addObject(listHolder.ref<ErrorPolicy::kTerminate>(), refFromObjCOrNSNull(anObject, holder.slot()));
 }
 
 - (void)removeLastObject {
+  kotlin::CalledFromNativeGuard guard;
   Kotlin_MutableList_removeLastObject(listHolder.ref<ErrorPolicy::kTerminate>());
 }
 
 - (void)replaceObjectAtIndex:(NSUInteger)index withObject:(id)anObject {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   KRef kotlinObject = refFromObjCOrNSNull(anObject, holder.slot());
   Kotlin_MutableList_setObject(listHolder.ref<ErrorPolicy::kTerminate>(), objCIndexToKotlinOrThrow(index), kotlinObject);
@@ -292,7 +302,7 @@ static inline id KSet_getElement(KRef set, id object) {
 }
 
 -(void)dealloc {
-  setHolder.dispose();
+  setHolder.disposeFromNative();
   [super dealloc];
 }
 
@@ -307,20 +317,24 @@ static inline id KSet_getElement(KRef set, id object) {
 }
 
 -(NSUInteger) count {
+  kotlin::CalledFromNativeGuard guard;
   return Kotlin_Collection_getSize(setHolder.ref<ErrorPolicy::kTerminate>());
 }
 
 - (id)member:(id)object {
+  kotlin::CalledFromNativeGuard guard;
   return KSet_getElement(setHolder.ref<ErrorPolicy::kTerminate>(), object);
 }
 
 // Not mandatory, just an optimization:
 - (BOOL)containsObject:(id)anObject {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   return Kotlin_Set_contains(setHolder.ref<ErrorPolicy::kTerminate>(), refFromObjCOrNSNull(anObject, holder.slot()));
 }
 
 - (NSEnumerator*)objectEnumerator {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   return [KIteratorAsNSEnumerator createWithKIterator:Kotlin_Set_iterator(setHolder.ref<ErrorPolicy::kTerminate>(), holder.slot())];
 }
@@ -336,7 +350,7 @@ static inline id KSet_getElement(KRef set, id object) {
 -(instancetype)init {
   if (self = [super init]) {
     Kotlin_initRuntimeIfNeeded();
-    // TODO: Does this need a switch to runnable state?
+    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kRunnable);
     ObjHolder holder;
     KRef set = Kotlin_MutableSet_createWithCapacity(8, holder.slot());
     self->setHolder.init(set);
@@ -348,7 +362,7 @@ static inline id KSet_getElement(KRef set, id object) {
 - (instancetype)initWithCapacity:(NSUInteger)numItems {
   if (self = [super init]) {
     Kotlin_initRuntimeIfNeeded();
-    // TODO: Does this need a switch to runnable state?
+    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kRunnable);
     ObjHolder holder;
     KRef set = Kotlin_MutableSet_createWithCapacity(objCCapacityToKotlin(numItems), holder.slot());
     self->setHolder.init(set);
@@ -376,7 +390,7 @@ static inline id KSet_getElement(KRef set, id object) {
   // 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();
+  setHolder.disposeFromNative();
   [super dealloc];
 }
 
@@ -393,30 +407,36 @@ static inline id KSet_getElement(KRef set, id object) {
 }
 
 -(NSUInteger) count {
+  kotlin::CalledFromNativeGuard guard;
   return Kotlin_Collection_getSize(setHolder.ref<ErrorPolicy::kTerminate>());
 }
 
 - (id)member:(id)object {
+  kotlin::CalledFromNativeGuard guard;
   return KSet_getElement(setHolder.ref<ErrorPolicy::kTerminate>(), object);
 }
 
 // Not mandatory, just an optimization:
 - (BOOL)containsObject:(id)anObject {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   return Kotlin_Set_contains(setHolder.ref<ErrorPolicy::kTerminate>(), refFromObjCOrNSNull(anObject, holder.slot()));
 }
 
 - (NSEnumerator*)objectEnumerator {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   return [KIteratorAsNSEnumerator createWithKIterator:Kotlin_Set_iterator(setHolder.ref<ErrorPolicy::kTerminate>(), holder.slot())];
 }
 
 - (void)addObject:(id)object {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   Kotlin_MutableCollection_addObject(setHolder.ref<ErrorPolicy::kTerminate>(), refFromObjCOrNSNull(object, holder.slot()));
 }
 
 - (void)removeObject:(id)object {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   Kotlin_MutableCollection_removeObject(setHolder.ref<ErrorPolicy::kTerminate>(), refFromObjCOrNSNull(object, holder.slot()));
 }
@@ -444,7 +464,7 @@ static inline id KMap_get(KRef map, id aKey) {
 }
 
 -(void)dealloc {
-  mapHolder.dispose();
+  mapHolder.disposeFromNative();
   [super dealloc];
 }
 
@@ -462,14 +482,17 @@ static inline id KMap_get(KRef map, id aKey) {
 // But that doesn't make any sense, since this class can't be arbitrary initialized.
 
 -(NSUInteger) count {
+  kotlin::CalledFromNativeGuard guard;
   return Kotlin_Map_getSize(mapHolder.ref<ErrorPolicy::kTerminate>());
 }
 
 - (id)objectForKey:(id)aKey {
+  kotlin::CalledFromNativeGuard guard;
   return KMap_get(mapHolder.ref<ErrorPolicy::kTerminate>(), aKey);
 }
 
 - (NSEnumerator *)keyEnumerator {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   return [KIteratorAsNSEnumerator createWithKIterator:Kotlin_Map_keyIterator(mapHolder.ref<ErrorPolicy::kTerminate>(), holder.slot())];
 }
@@ -487,14 +510,14 @@ static inline id KMap_get(KRef map, id aKey) {
   // 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();
+  mapHolder.disposeFromNative();
   [super dealloc];
 }
 
 -(instancetype)init {
   if (self = [super init]) {
     Kotlin_initRuntimeIfNeeded();
-    // TODO: Does this need a switch to runnable state?
+    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kRunnable);
     ObjHolder holder;
     KRef map = Kotlin_MutableMap_createWithCapacity(8, holder.slot());
     self->mapHolder.init(map);
@@ -511,7 +534,7 @@ static inline id KMap_get(KRef map, id aKey) {
 - (instancetype)initWithCapacity:(NSUInteger)numItems {
   if (self = [super init]) {
     Kotlin_initRuntimeIfNeeded();
-    // TODO: Does this need a switch to runnable state?
+    kotlin::ThreadStateGuard guard(kotlin::ThreadState::kRunnable);
     ObjHolder holder;
     KRef map = Kotlin_MutableMap_createWithCapacity(objCCapacityToKotlin(numItems), holder.slot());
     self->mapHolder.init(map);
@@ -532,31 +555,38 @@ static inline id KMap_get(KRef map, id aKey) {
 }
 
 -(NSUInteger) count {
+  kotlin::CalledFromNativeGuard guard;
   return Kotlin_Map_getSize(mapHolder.ref<ErrorPolicy::kTerminate>());
 }
 
 - (id)objectForKey:(id)aKey {
+  kotlin::CalledFromNativeGuard guard;
   return KMap_get(mapHolder.ref<ErrorPolicy::kTerminate>(), aKey);
 }
 
 - (NSEnumerator *)keyEnumerator {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   return [KIteratorAsNSEnumerator createWithKIterator:Kotlin_Map_keyIterator(mapHolder.ref<ErrorPolicy::kTerminate>(), holder.slot())];
 }
 
 - (void)setObject:(id)anObject forKey:(id<NSCopying>)aKey {
-  ObjHolder keyHolder, valueHolder;
-
   id keyCopy = [aKey copyWithZone:nullptr]; // Correspond to the expected NSMutableDictionary behaviour.
-  KRef kotlinKey = refFromObjCOrNSNull(keyCopy, keyHolder.slot());
+  {
+    kotlin::CalledFromNativeGuard guard;
+    ObjHolder keyHolder, valueHolder;
+
+    KRef kotlinKey = refFromObjCOrNSNull(keyCopy, keyHolder.slot());
+
+    KRef kotlinValue = refFromObjCOrNSNull(anObject, valueHolder.slot());
+
+    Kotlin_MutableMap_set(mapHolder.ref<ErrorPolicy::kTerminate>(), kotlinKey, kotlinValue);
+  }
   objc_release(keyCopy);
-
-  KRef kotlinValue = refFromObjCOrNSNull(anObject, valueHolder.slot());
-
-  Kotlin_MutableMap_set(mapHolder.ref<ErrorPolicy::kTerminate>(), kotlinKey, kotlinValue);
 }
 
 - (void)removeObjectForKey:(id)aKey {
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   KRef kotlinKey = refFromObjCOrNSNull(aKey, holder.slot());