Native: improve thread state switching for Obj-C interop

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

@@ -785,6 +785,9 @@ private class InteropTransformer(val context: Context, override val irFile: IrFi
         // Calls to other ObjC class constructors must be lowered.
         require(constructedClass.isKotlinObjCClass()) { renderCompilerError(expression) }
         return builder.at(expression).irBlock {
+            // Note: using [interopAllocObjCObject] and [interopObjCRelease] here is suboptimal: they switch the thread to Native state
+            // and then back to Runnable.
+            // TODO: consider calling specialized versions of allocWithZoneImp and releaseImp directly.
             val rawPtr = irTemporary(irCall(symbols.interopAllocObjCObject.owner).apply {
                 putValueArgument(0, getObjCClass(symbols, constructedClass.symbol))
             })

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

@@ -125,20 +125,21 @@ template <ErrorPolicy errorPolicy>
 void BackRefFromAssociatedObject::addRef() {
   static_assert(errorPolicy != ErrorPolicy::kDefaultValue, "Cannot use default return value here");
 
+  // Can be called both from Native state (if ObjC or Swift code adds RC)
+  // and from Runnable state (Kotlin_ObjCExport_refToObjC).
+
   if (atomicAdd(&refCount, 1) == 1) {
     if (obj_ == nullptr) return; // E.g. after [detach].
 
+    kotlin::CalledFromNativeGuard guard(/* reentrant */ true);
+
     // 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).
   }
 }
 

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

@@ -89,6 +89,8 @@ id allocWithZoneImp(Class self, SEL _cmd, void* zone) {
   id result = messenger(&s, _cmd, zone);
 
   auto* typeInfo = classData->typeInfo;
+
+  kotlin::CalledFromNativeGuard guard;
   ObjHolder holder;
   auto kotlinObj = AllocInstanceWithAssociatedObject(typeInfo, result, holder.slot());
 
@@ -372,6 +374,7 @@ void Kotlin_objc_autoreleasePoolPop(void* ptr) {
 }
 
 id Kotlin_objc_allocWithZone(Class clazz) {
+  kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
   return objc_allocWithZone(clazz);
 }
 
@@ -380,6 +383,7 @@ id Kotlin_objc_retain(id ptr) {
 }
 
 void Kotlin_objc_release(id ptr) {
+  kotlin::ThreadStateGuard guard(kotlin::ThreadState::kNative);
   objc_release(ptr);
 }
 

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

@@ -35,7 +35,7 @@
 }
 
 -(void)dealloc {
-  refHolder.dispose();
+  refHolder.disposeFromNative();
   [super dealloc];
 }
 
@@ -88,9 +88,23 @@ extern "C" OBJ_GETTER(Kotlin_Interop_refFromObjC, id obj);
 static OBJ_GETTER(Konan_ObjCInterop_getWeakReference, KRef ref) {
   KotlinObjCWeakReference* objcRef = (KotlinObjCWeakReference*)ref->GetAssociatedObject();
 
-  id objcReferred = objc_loadWeakRetained(&objcRef->referred);
+  // objc_loadWeakRetained can call arbitrary user code, so it needs Native state:
+  id objcReferred = kotlin::CallWithThreadState<kotlin::ThreadState::kNative>(objc_loadWeakRetained, &objcRef->referred);
+
+  // Kotlin_Interop_refFromObjC creates Kotlin objects, so it needs Runnable state:
   KRef result = Kotlin_Interop_refFromObjC(objcReferred, OBJ_RESULT);
-  objc_release(objcReferred);
+
+  // objc_release can call arbitrary user code, so it needs Native state:
+  kotlin::CallWithThreadState<kotlin::ThreadState::kNative>(objc_release, objcReferred);
+
+  // Possible optimizations for thread state switching above:
+  //  1. `Kotlin_Interop_refFromObjC` typically does `objc_retain` under the hood. We could coalesce it with `objc_release` above.
+  //  2. `objc_loadWeakRetained` and `objc_release` typically shouldn't call arbitrary user code here:
+  //     the latter can't deallocate the object because it is still alive by this point, so the only possibility for user code here is
+  //     when a user overrides `_tryRetain` or `release`.
+  //     We do this for Kotlin subclasses of Obj-C classes, which also use this implementation.
+  //     But we could use the other weak implementation for such classes, and assume no one else overrides these methods
+  //     (in a dangerous way).
 
   return result;
 }