Make Kotlin_initRuntimeIfNeeded switch to native state

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

@@ -366,6 +366,7 @@ private class ExportedElement(val kind: ElementKind,
             |extern "C" KObjHeader* ${cname}_instance(KObjHeader**);
             |static $objectClassC ${cname}_instance_impl(void) {
             |  Kotlin_initRuntimeIfNeeded();
+            |  ScopedRunnableState stateGuard;
             |  KObjHolder result_holder;
             |  KObjHeader* result = ${cname}_instance(result_holder.slot());
             |  return $objectClassC { .pinned = CreateStablePointer(result)};
@@ -384,6 +385,7 @@ private class ExportedElement(val kind: ElementKind,
               |extern "C" KObjHeader* $cname(KObjHeader**);
               |static $enumClassC ${cname}_impl(void) {
               |  Kotlin_initRuntimeIfNeeded();
+              |  ScopedRunnableState stateGuard;
               |  KObjHolder result_holder;
               |  KObjHeader* result = $cname(result_holder.slot());
               |  return $enumClassC { .pinned = CreateStablePointer(result)};
@@ -428,6 +430,9 @@ private class ExportedElement(val kind: ElementKind,
         val builder = StringBuilder()
         builder.append("$visibility ${owner.translateType(cfunction[0])} ${cnameImpl}(${cfunction.drop(1).
                 mapIndexed { index, it -> "${owner.translateType(it)} arg${index}" }.joinToString(", ")}) {\n")
+        // TODO: do we really need that in every function?
+        builder.append("  Kotlin_initRuntimeIfNeeded();\n")
+        builder.append("  ScopedRunnableState stateGuard;\n");
         val args = ArrayList(cfunction.drop(1).mapIndexed { index, pair ->
             translateArgument("arg$index", pair, Direction.C_TO_KOTLIN, builder)
         })
@@ -435,8 +440,6 @@ private class ExportedElement(val kind: ElementKind,
         val isConstructor = declaration is ConstructorDescriptor
         val isObjectReturned = !isConstructor && owner.isMappedToReference(cfunction[0].type)
         val isStringReturned = owner.isMappedToString(cfunction[0].type)
-        // TODO: do we really need that in every function?
-        builder.append("  Kotlin_initRuntimeIfNeeded();\n")
         builder.append("   try {\n")
         if (isObjectReturned || isStringReturned) {
             builder.append("  KObjHolder result_holder;\n")
@@ -918,6 +921,8 @@ internal class CAdapterGenerator(val context: Context) : DeclarationDescriptorVi
         |void EnterFrame(KObjHeader** start, int parameters, int count) RUNTIME_NOTHROW;
         |void LeaveFrame(KObjHeader** start, int parameters, int count) RUNTIME_NOTHROW;
         |void Kotlin_initRuntimeIfNeeded();
+        |void Kotlin_mm_switchThreadStateRunnable() RUNTIME_NOTHROW;
+        |void Kotlin_mm_switchThreadStateNative() RUNTIME_NOTHROW;
         |void TerminateWithUnhandledException(KObjHeader*) RUNTIME_NORETURN;
         |
         |KObjHeader* CreateStringFromCString(const char*, KObjHeader**);
@@ -960,6 +965,16 @@ internal class CAdapterGenerator(val context: Context) : DeclarationDescriptorVi
         |    KObjHeader* GetExceptionObject() noexcept;
         |};
         |
+        |class ScopedRunnableState {
+        |public:
+        |   ScopedRunnableState() noexcept { Kotlin_mm_switchThreadStateRunnable(); }
+        |   ~ScopedRunnableState() { Kotlin_mm_switchThreadStateNative(); }
+        |   ScopedRunnableState(const ScopedRunnableState&) = delete;
+        |   ScopedRunnableState(ScopedRunnableState&&) = delete;
+        |   ScopedRunnableState& operator=(const ScopedRunnableState&) = delete;
+        |   ScopedRunnableState& operator=(ScopedRunnableState&&) = delete;
+        |};
+        |
         |static void DisposeStablePointerImpl(${prefix}_KNativePtr ptr) {
         |  DisposeStablePointer(ptr);
         |}
@@ -967,6 +982,7 @@ internal class CAdapterGenerator(val context: Context) : DeclarationDescriptorVi
         |  DisposeCString((char*)ptr);
         |}
         |static ${prefix}_KBoolean IsInstanceImpl(${prefix}_KNativePtr ref, const ${prefix}_KType* type) {
+        |  ScopedRunnableState stateGuard;
         |  KObjHolder holder;
         |  return IsInstance(DerefStablePointer(ref, holder.slot()), (const KTypeInfo*)type);
         |}
@@ -981,6 +997,7 @@ internal class CAdapterGenerator(val context: Context) : DeclarationDescriptorVi
             output("extern \"C\" KObjHeader* Kotlin_box${it.shortNameForPredefinedType}($parameter$maybeComma KObjHeader**);")
             output("static ${translateType(nullableIt)} ${it.createNullableNameForPredefinedType}Impl($parameter) {")
             output("Kotlin_initRuntimeIfNeeded();", 1)
+            output("ScopedRunnableState stateGuard;", 1)
             output("KObjHolder result_holder;", 1)
             output("KObjHeader* result = Kotlin_box${it.shortNameForPredefinedType}($argument result_holder.slot());", 1)
             output("return ${translateType(nullableIt)} { .pinned = CreateStablePointer(result) };", 1)

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

@@ -1276,6 +1276,8 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
                 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?
+            call(context.llvm.initRuntimeIfNeeded, emptyList())
             switchThreadState(Runnable)
         }
 
@@ -1292,6 +1294,7 @@ internal class FunctionGenerationContext(val function: LLVMValueRef,
             if (needsRuntimeInit) {
                 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, "")!!

kotlin-native/backend.native/tests/build.gradle

@@ -4899,6 +4899,7 @@ if (isAppleTarget(project)) {
     }
 
     frameworkTest('testValuesGenericsFramework') {
+        enabled = !isExperimentalMM // Experimental MM does not support thread state switching for ObjC interop yet.
         framework('ValuesGenerics') {
             sources = ['objcexport/values.kt', 'framework/values_generics']
         }
@@ -4906,6 +4907,7 @@ if (isAppleTarget(project)) {
     }
 
     frameworkTest("testStdlibFramework") {
+        enabled = !isExperimentalMM // Experimental MM does not support thread state switching for ObjC interop yet.
         framework('Stdlib') {
             sources = ['framework/stdlib']
             bitcode = true
@@ -4917,6 +4919,7 @@ if (isAppleTarget(project)) {
     if (cacheTesting != null && cacheTesting.isDynamic) {
         // testMultipleFrameworks disabled until https://youtrack.jetbrains.com/issue/KT-34262 is fixed.
     } else frameworkTest("testMultipleFrameworks") {
+        enabled = !isExperimentalMM // Experimental MM does not support thread state switching for ObjC interop yet.
         framework('First') {
             sources = ['framework/multiple/framework1', 'framework/multiple/shared']
             bitcode = true
@@ -4929,6 +4932,7 @@ if (isAppleTarget(project)) {
     }
 
     frameworkTest("testMultipleFrameworksStatic") {
+        enabled = !isExperimentalMM // Experimental MM does not support thread state switching for ObjC interop yet.
         if (cacheTesting != null) {
             // See https://youtrack.jetbrains.com/issue/KT-34261.
             expectedExitStatus = 134
@@ -4962,7 +4966,8 @@ if (isAppleTarget(project)) {
     }
 
     frameworkTest("testKt42397Framework") {
-        enabled = !project.globalTestArgs.contains('-opt')
+        enabled = !project.globalTestArgs.contains('-opt') &&
+            !isExperimentalMM // Experimental MM does not support thread state switching for ObjC interop yet.
         framework("Kt42397") {
             sources = ['framework/kt42397']
         }
@@ -4970,6 +4975,7 @@ if (isAppleTarget(project)) {
     }
 
     frameworkTest("testKt43517Framework") {
+        enabled = !isExperimentalMM // Experimental MM does not support thread state switching for ObjC interop yet.
         framework('Kt43517') {
             sources = ['framework/kt43517']
             library = 'objcKt43517'

kotlin-native/backend.native/tests/interop/leakMemoryWithRunningThread/checked.kt

@@ -7,6 +7,7 @@ import kotlinx.cinterop.*
 val global = AtomicInt(0)
 
 fun ensureInititalized() {
+    // Only needed with the legacy MM. TODO: Remove when legacy MM is gone.
     kotlin.native.initRuntimeIfNeeded()
     // Leak memory
     StableRef.create(Any())

kotlin-native/backend.native/tests/interop/leakMemoryWithRunningThread/unchecked.kt

@@ -7,6 +7,7 @@ import kotlinx.cinterop.*
 val global = AtomicInt(0)
 
 fun ensureInititalized() {
+    // Only needed with the legacy MM. TODO: Remove when legacy MM is gone.
     kotlin.native.initRuntimeIfNeeded()
     // Leak memory
     StableRef.create(Any())

kotlin-native/runtime/src/launcher/cpp/androidLauncher.cpp

@@ -74,6 +74,7 @@ namespace {
 
 void launchMain() {
   Kotlin_initRuntimeIfNeeded();
+  Kotlin_mm_switchThreadStateRunnable();
 
   {
     ObjHolder args;

kotlin-native/runtime/src/launcher/cpp/launcher.cpp

@@ -53,6 +53,7 @@ extern "C" RUNTIME_USED int Init_and_run_start(int argc, const char** argv, int
 #endif
 
   Kotlin_initRuntimeIfNeeded();
+  Kotlin_mm_switchThreadStateRunnable();
 
   KInt exitStatus = Konan_run_start(argc, argv);
 

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

@@ -19,6 +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?
     Kotlin_initRuntimeIfNeeded();
 
     return IsForeignRefAccessible(object, context);

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

@@ -47,6 +47,7 @@ static inline OBJ_GETTER(refFromObjCOrNSNull, id obj) {
 
 static inline OBJ_GETTER(invokeAndAssociate, KRef (*func)(KRef* result), id obj) {
   Kotlin_initRuntimeIfNeeded();
+  // TODO: Does this need a switch to runnable state?
 
   KRef kotlinObj = func(OBJ_RESULT);
 
@@ -56,4 +57,4 @@ static inline OBJ_GETTER(invokeAndAssociate, KRef (*func)(KRef* result), id obj)
 }
 
 #endif // KONAN_OBJC_INTEROP
-#endif // RUNTIME_OBJCEXPORTCOLLECTIONS_H
+#endif // RUNTIME_OBJCEXPORTCOLLECTIONS_H

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

@@ -140,6 +140,9 @@ RuntimeState* initRuntime() {
   InitOrDeinitGlobalVariables(INIT_THREAD_LOCAL_GLOBALS, result->memoryState);
   RuntimeAssert(result->status == RuntimeStatus::kUninitialized, "Runtime must still be in the uninitialized state");
   result->status = RuntimeStatus::kRunning;
+
+  kotlin::SwitchThreadState(result->memoryState, kotlin::ThreadState::kNative);
+
   return result;
 }
 
@@ -172,6 +175,8 @@ void deinitRuntime(RuntimeState* state, bool destroyRuntime) {
 
 void Kotlin_deinitRuntimeCallback(void* argument) {
   auto* state = reinterpret_cast<RuntimeState*>(argument);
+  // This callback may be called from any state, make sure it runs in the runnable state.
+  kotlin::SwitchThreadState(state->memoryState, kotlin::ThreadState::kRunnable, /* reentrant = */ true);
   deinitRuntime(state, false);
 }
 
@@ -378,9 +383,18 @@ KBoolean Kotlin_Debugging_isThreadStateNative() {
     return kotlin::GetThreadState() == kotlin::ThreadState::kNative;
 }
 
-
 KBoolean Kotlin_Debugging_isPermanent(KRef obj) {
     return obj->permanent();
 }
 
+RUNTIME_NOTHROW void Kotlin_initRuntimeIfNeededFromKotlin() {
+    switch (CurrentMemoryModel) {
+        case MemoryModel::kExperimental:
+            return;
+        case MemoryModel::kStrict:
+        case MemoryModel::kRelaxed:
+            Kotlin_initRuntimeIfNeeded();
+    }
+}
+
 }  // extern "C"

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

@@ -33,6 +33,7 @@ enum DestroyRuntimeMode {
 
 DestroyRuntimeMode Kotlin_getDestroyRuntimeMode();
 
+// For experimental MM, if runtime gets initialized, it will be in the native state after this.
 RUNTIME_NOTHROW void Kotlin_initRuntimeIfNeeded();
 void Kotlin_deinitRuntimeIfNeeded();
 

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

@@ -881,6 +881,7 @@ void* workerRoutine(void* argument) {
   // to see there's already a worker created for this thread.
   ::g_worker = worker;
   Kotlin_initRuntimeIfNeeded();
+  SwitchThreadState(worker->memoryState(), ThreadState::kRunnable);
 
   do {
     if (worker->processQueueElement(true) == JOB_TERMINATE) break;

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

@@ -11,7 +11,7 @@ import kotlin.native.internal.UnhandledExceptionHookHolder
 /**
  * Initializes Kotlin runtime for the current thread, if not inited already.
  */
-@GCUnsafeCall("Kotlin_initRuntimeIfNeeded")
+@GCUnsafeCall("Kotlin_initRuntimeIfNeededFromKotlin")
 external public fun initRuntimeIfNeeded(): Unit
 
 /**

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

@@ -13,8 +13,9 @@
 #include "InitializationScheme.hpp"
 #include "KAssert.h"
 #include "Natives.h"
-#include "Porting.h"
 #include "ObjectOps.hpp"
+#include "Porting.h"
+#include "Runtime.h"
 #include "StableRefRegistry.hpp"
 #include "ThreadData.hpp"
 #include "ThreadRegistry.hpp"

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

@@ -61,6 +61,7 @@ static void injectToRuntime();
 
 +(instancetype)allocWithZone:(NSZone*)zone {
   Kotlin_initRuntimeIfNeeded();
+  // TODO: Does this need a switch to runnable state?
 
   KotlinBase* result = [super allocWithZone:zone];
 

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

@@ -327,6 +327,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?
     ObjHolder holder;
     KRef set = Kotlin_MutableSet_createWithCapacity(8, holder.slot());
     self->setHolder.init(set);
@@ -338,6 +339,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?
     ObjHolder holder;
     KRef set = Kotlin_MutableSet_createWithCapacity(objCCapacityToKotlin(numItems), holder.slot());
     self->setHolder.init(set);
@@ -483,6 +485,7 @@ static inline id KMap_get(KRef map, id aKey) {
 -(instancetype)init {
   if (self = [super init]) {
     Kotlin_initRuntimeIfNeeded();
+    // TODO: Does this need a switch to runnable state?
     ObjHolder holder;
     KRef map = Kotlin_MutableMap_createWithCapacity(8, holder.slot());
     self->mapHolder.init(map);
@@ -499,6 +502,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?
     ObjHolder holder;
     KRef map = Kotlin_MutableMap_createWithCapacity(objCCapacityToKotlin(numItems), holder.slot());
     self->mapHolder.init(map);