Add basic support for suspend functions in produced Obj-C header

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

@@ -207,6 +207,11 @@ internal class KonanSymbols(
                     NoLookupLocation.FROM_BACKEND
             ).single())
 
+    val objCExportResumeContinuation = internalFunction("resumeContinuation")
+    val objCExportResumeContinuationWithException = internalFunction("resumeContinuationWithException")
+    val objCExportGetCoroutineSuspended = internalFunction("getCoroutineSuspended")
+    val objCExportInterceptedContinuation = internalFunction("interceptedContinuation")
+
     val getNativeNullPtr = symbolTable.referenceSimpleFunction(context.getNativeNullPtr)
 
     val boxCachePredicates = BoxCache.values().associate {

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

@@ -488,6 +488,8 @@ internal class Llvm(val context: Context, val llvmModule: LLVMModuleRef) {
     val Kotlin_ObjCExport_RethrowExceptionAsNSError by lazyRtFunction
     val Kotlin_ObjCExport_RethrowNSErrorAsException by lazyRtFunction
     val Kotlin_ObjCExport_AllocInstanceWithAssociatedObject by lazyRtFunction
+    val Kotlin_ObjCExport_createContinuationArgument by lazyRtFunction
+    val Kotlin_ObjCExport_resumeContinuation by lazyRtFunction
 
     val tlsMode by lazy {
         when (target) {

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

@@ -86,7 +86,7 @@ internal open class ObjCExportCodeGeneratorBase(codegen: CodeGenerator) : ObjCCo
                 generateBlockToKotlinFunctionConverter(bridge)
             }
 
-    private val blockGenerator = BlockGenerator(this.codegen)
+    protected val blockGenerator = BlockGenerator(this.codegen)
     private val functionToBlockConverterCache = mutableMapOf<BlockPointerBridge, LLVMValueRef>()
 
     internal fun kotlinFunctionToBlockConverter(bridge: BlockPointerBridge): LLVMValueRef =
@@ -123,6 +123,10 @@ internal class ObjCExportCodeGenerator(
 
     val externalGlobalInitializers = mutableMapOf<LLVMValueRef, ConstValue>()
 
+    internal val continuationToCompletionConverter: LLVMValueRef by lazy {
+        generateContinuationToCompletionConverter(blockGenerator)
+    }
+
     fun FunctionGenerationContext.genSendMessage(
             returnType: LLVMTypeRef,
             receiver: LLVMValueRef,
@@ -584,6 +588,20 @@ private fun ObjCExportCodeGenerator.emitBoxConverter(
     setObjCExportTypeInfo(boxClass, constPointer(converter))
 }
 
+private fun ObjCExportCodeGenerator.generateContinuationToCompletionConverter(
+        blockGenerator: BlockGenerator
+): LLVMValueRef = with(blockGenerator) {
+    generateWrapKotlinObjectToBlock(
+            BlockType(numberOfParameters = 2, returnsVoid = true),
+            convertName = "convertContinuation",
+            invokeName = "invokeCompletion"
+    ) { continuation, arguments ->
+        check(arguments.size == 2)
+        callFromBridge(context.llvm.Kotlin_ObjCExport_resumeContinuation, listOf(continuation) + arguments)
+        ret(null)
+    }
+}
+
 private const val maxConvertorsInCache = 33
 
 private fun ObjCExportBlockCodeGenerator.emitFunctionConverters() {
@@ -760,6 +778,7 @@ private fun ObjCExportCodeGenerator.generateObjCImp(
     }
 
     var errorOutPtr: LLVMValueRef? = null
+    var continuation: LLVMValueRef? = null
 
     val kotlinArgs = methodBridge.paramBridges.mapIndexedNotNull { index, paramBridge ->
         val parameter = param(index)
@@ -777,17 +796,22 @@ private fun ObjCExportCodeGenerator.generateObjCImp(
                 errorOutPtr = parameter
                 null
             }
+
+            MethodBridgeValueParameter.SuspendCompletion -> {
+                callFromBridge(
+                        context.llvm.Kotlin_ObjCExport_createContinuationArgument,
+                        listOf(parameter, generateExceptionTypeInfoArray(baseMethod!!)),
+                        Lifetime.ARGUMENT
+                ).also {
+                    continuation = it
+                }
+            }
         }
     }
 
     // TODO: consider merging this handler with function cleanup.
-    val exceptionHandler = if (errorOutPtr == null) {
-        kotlinExceptionHandler { exception ->
-            callFromBridge(symbols.objCExportTrapOnUndeclaredException.owner.llvmFunction, listOf(exception))
-            unreachable()
-        }
-    } else {
-        kotlinExceptionHandler { exception ->
+    val exceptionHandler = when {
+        errorOutPtr != null -> kotlinExceptionHandler { exception ->
             callFromBridge(
                     context.llvm.Kotlin_ObjCExport_RethrowExceptionAsNSError,
                     listOf(exception, errorOutPtr!!, generateExceptionTypeInfoArray(baseMethod!!))
@@ -810,6 +834,22 @@ private fun ObjCExportCodeGenerator.generateObjCImp(
 
             ret(returnValue)
         }
+
+        continuation != null -> kotlinExceptionHandler { exception ->
+            // Callee haven't suspended, so it isn't going to call the completion. Call it here:
+            callFromBridge(
+                    context.ir.symbols.objCExportResumeContinuationWithException.owner.llvmFunction,
+                    listOf(continuation!!, exception)
+            )
+            // Note: completion block could be called directly instead, but this implementation is
+            // simpler and avoids duplication.
+            ret(null)
+        }
+
+        else -> kotlinExceptionHandler { exception ->
+            callFromBridge(symbols.objCExportTrapOnUndeclaredException.owner.llvmFunction, listOf(exception))
+            unreachable()
+        }
     }
 
     val targetResult = callKotlin(kotlinArgs, Lifetime.ARGUMENT, exceptionHandler)
@@ -825,6 +865,23 @@ private fun ObjCExportCodeGenerator.generateObjCImp(
         is MethodBridge.ReturnValue.WithError.ZeroForError -> genReturnValueOnSuccess(returnBridge.successBridge)
         MethodBridge.ReturnValue.Instance.InitResult -> param(0)
         MethodBridge.ReturnValue.Instance.FactoryResult -> kotlinReferenceToObjC(targetResult!!) // provided by [callKotlin]
+        MethodBridge.ReturnValue.Suspend -> {
+            val coroutineSuspended = callFromBridge(
+                    codegen.llvmFunction(context.ir.symbols.objCExportGetCoroutineSuspended.owner),
+                    emptyList(),
+                    Lifetime.LOCAL
+            )
+            ifThen(icmpNe(targetResult!!, coroutineSuspended)) {
+                // Callee haven't suspended, so it isn't going to call the completion. Call it here:
+                callFromBridge(
+                        context.ir.symbols.objCExportResumeContinuation.owner.llvmFunction,
+                        listOf(continuation!!, targetResult)
+                )
+                // Note: completion block could be called directly instead, but this implementation is
+                // simpler and avoids duplication.
+            }
+            null
+        }
     }
 
     ret(genReturnValueOnSuccess(returnType))
@@ -928,6 +985,17 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
                         store(kNullInt8Ptr, it)
                         errorOutPtr = it
                     }
+
+                MethodBridgeValueParameter.SuspendCompletion -> {
+                    val continuation = param(irFunction.allParameters.size) // The last argument.
+                    // TODO: consider placing interception into the converter to reduce code size.
+                    val intercepted = callFromBridge(
+                            context.ir.symbols.objCExportInterceptedContinuation.owner.llvmFunction,
+                            listOf(continuation),
+                            Lifetime.ARGUMENT
+                    )
+                    callFromBridge(continuationToCompletionConverter, listOf(intercepted))
+                }
             }
         }
 
@@ -985,12 +1053,25 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
             MethodBridge.ReturnValue.Instance.InitResult,
             MethodBridge.ReturnValue.Instance.FactoryResult ->
                 error("init or factory method can't have bridge for overriding: $baseMethod")
+
+            MethodBridge.ReturnValue.Suspend -> {
+                // Objective-C implementation of Kotlin suspend function is always responsible
+                // for calling the completion, so in Kotlin coroutines machinery terms it suspends,
+                // which is indicated by the return value:
+                callFromBridge(
+                        context.ir.symbols.objCExportGetCoroutineSuspended.owner.llvmFunction,
+                        emptyList(),
+                        Lifetime.RETURN_VALUE
+                )
+            }
         }
 
         val baseReturnType = baseIrFunction.returnType
         val actualReturnType = irFunction.returnType
 
         val retVal = when {
+            baseIrFunction.isSuspend -> genKotlinBaseMethodResult(Lifetime.RETURN_VALUE, methodBridge.returnBridge)
+
             actualReturnType.isUnit() || actualReturnType.isNothing() -> {
                 genKotlinBaseMethodResult(Lifetime.ARGUMENT, methodBridge.returnBridge)
                 null
@@ -1433,10 +1514,12 @@ private val MethodBridgeParameter.objCType: LLVMTypeRef get() = when (this) {
     is MethodBridgeReceiver -> ReferenceBridge.objCType
     MethodBridgeSelector -> int8TypePtr
     MethodBridgeValueParameter.ErrorOutParameter -> pointerType(ReferenceBridge.objCType)
+    MethodBridgeValueParameter.SuspendCompletion -> int8TypePtr
 }
 
 private fun MethodBridge.ReturnValue.objCType(context: Context): LLVMTypeRef {
     return when (this) {
+        MethodBridge.ReturnValue.Suspend,
         MethodBridge.ReturnValue.Void -> voidType
         MethodBridge.ReturnValue.HashCode -> if (context.is64BitNSInteger()) int64Type else int32Type
         is MethodBridge.ReturnValue.Mapped -> this.bridge.objCType
@@ -1483,6 +1566,7 @@ private val Family.nsUIntegerEncoding: String get() = when (this) {
 }
 
 private fun MethodBridge.ReturnValue.getObjCEncoding(targetFamily: Family): String = when (this) {
+    MethodBridge.ReturnValue.Suspend,
     MethodBridge.ReturnValue.Void -> "v"
     MethodBridge.ReturnValue.HashCode -> targetFamily.nsUIntegerEncoding
     is MethodBridge.ReturnValue.Mapped -> this.bridge.objCEncoding
@@ -1498,6 +1582,7 @@ private val MethodBridgeParameter.objCEncoding: String get() = when (this) {
     is MethodBridgeReceiver -> ReferenceBridge.objCEncoding
     MethodBridgeSelector -> ":"
     MethodBridgeValueParameter.ErrorOutParameter -> "^${ReferenceBridge.objCEncoding}"
+    MethodBridgeValueParameter.SuspendCompletion -> "@"
 }
 
 private val TypeBridge.objCEncoding: String get() = when (this) {

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

@@ -35,6 +35,7 @@ internal object MethodBridgeSelector : MethodBridgeParameter()
 internal sealed class MethodBridgeValueParameter : MethodBridgeParameter() {
     data class Mapped(val bridge: TypeBridge) : MethodBridgeValueParameter()
     object ErrorOutParameter : MethodBridgeValueParameter()
+    object SuspendCompletion : MethodBridgeValueParameter()
 }
 
 internal data class MethodBridge(
@@ -56,6 +57,8 @@ internal data class MethodBridge(
             object Success : WithError()
             data class ZeroForError(val successBridge: ReturnValue, val successMayBeZero: Boolean) : WithError()
         }
+
+        object Suspend : ReturnValue()
     }
 
     val paramBridges: List<MethodBridgeParameter> =
@@ -86,6 +89,7 @@ internal fun MethodBridge.valueParametersAssociated(
         when (it) {
             is MethodBridgeValueParameter.Mapped -> it to kotlinParameters.next()
 
+            MethodBridgeValueParameter.SuspendCompletion,
             is MethodBridgeValueParameter.ErrorOutParameter -> it to null
         }
     }.also { assert(!kotlinParameters.hasNext()) }
@@ -101,6 +105,7 @@ internal fun MethodBridge.parametersAssociated(
             is MethodBridgeValueParameter.Mapped, MethodBridgeReceiver.Instance ->
                 it to kotlinParameters.next()
 
+            MethodBridgeValueParameter.SuspendCompletion,
             MethodBridgeReceiver.Static, MethodBridgeSelector, MethodBridgeValueParameter.ErrorOutParameter ->
                 it to null
 

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

@@ -581,6 +581,7 @@ internal class ObjCExportTranslatorImpl(
                         }
                     }
                     MethodBridgeValueParameter.ErrorOutParameter -> "error"
+                    MethodBridgeValueParameter.SuspendCompletion -> "completionHandler"
                 }
 
                 val uniqueName = unifyName(candidateName, usedNames)
@@ -590,6 +591,16 @@ internal class ObjCExportTranslatorImpl(
                     is MethodBridgeValueParameter.Mapped -> mapType(p!!.type, bridge.bridge, objCExportScope)
                     MethodBridgeValueParameter.ErrorOutParameter ->
                         ObjCPointerType(ObjCNullableReferenceType(ObjCClassType("NSError")), nullable = true)
+
+                    MethodBridgeValueParameter.SuspendCompletion -> {
+                        ObjCBlockPointerType(
+                                returnType = ObjCVoidType,
+                                parameterTypes = listOf(
+                                        mapReferenceType(method.returnType!!, objCExportScope).makeNullable(),
+                                        ObjCNullableReferenceType(ObjCClassType("NSError"))
+                                )
+                        )
+                    }
                 }
 
                 parameters += ObjCParameter(uniqueName, p, type)
@@ -652,6 +663,7 @@ internal class ObjCExportTranslatorImpl(
     }
 
     private fun mapReturnType(returnBridge: MethodBridge.ReturnValue, method: FunctionDescriptor, objCExportScope: ObjCExportScope): ObjCType = when (returnBridge) {
+        MethodBridge.ReturnValue.Suspend,
         MethodBridge.ReturnValue.Void -> ObjCVoidType
         MethodBridge.ReturnValue.HashCode -> ObjCPrimitiveType.NSUInteger
         is MethodBridge.ReturnValue.Mapped -> mapType(method.returnType!!, returnBridge.bridge, objCExportScope)

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

@@ -160,7 +160,7 @@ internal class ObjCExportLazyImpl(
 
         file.children.filterIsInstance<KtCallableDeclaration>().forEach {
             // Supposed to be similar to ObjCExportMapper.shouldBeVisible.
-            if ((it is KtFunction || it is KtProperty) && it.isPublic && !it.isSuspend && !it.hasExpectModifier()) {
+            if ((it is KtFunction || it is KtProperty) && it.isPublic && !it.hasExpectModifier()) {
                 val classDescriptor = getClassIfExtension(it)
                 if (classDescriptor != null) {
                     extensions.getOrPut(classDescriptor, { mutableListOf() }) += it
@@ -438,9 +438,6 @@ private fun ModuleDescriptor.isCommonStdlib() =
 private val KtModifierListOwner.isPublic: Boolean
     get() = this.visibilityModifierTypeOrDefault() == KtTokens.PUBLIC_KEYWORD
 
-private val KtCallableDeclaration.isSuspend: Boolean
-    get() = this.hasModifier(KtTokens.SUSPEND_KEYWORD)
-
 internal val KtPureClassOrObject.isInterface: Boolean
     get() = this is KtClass && this.isInterface()
 

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

@@ -72,7 +72,7 @@ internal fun ObjCExportMapper.getClassIfCategory(extensionReceiverType: KotlinTy
 
 // Note: partially duplicated in ObjCExportLazyImpl.translateTopLevels.
 internal fun ObjCExportMapper.shouldBeExposed(descriptor: CallableMemberDescriptor): Boolean =
-        descriptor.isEffectivelyPublicApi && !descriptor.isSuspend && !descriptor.isExpect &&
+        descriptor.isEffectivelyPublicApi && !descriptor.isExpect &&
                 !isHiddenByDeprecation(descriptor)
 
 internal fun ObjCExportMapper.shouldBeExposed(descriptor: ClassDescriptor): Boolean =
@@ -241,6 +241,8 @@ private fun ObjCExportMapper.bridgeReturnType(
 ): MethodBridge.ReturnValue {
     val returnType = descriptor.returnType!!
     return when {
+        descriptor.isSuspend -> MethodBridge.ReturnValue.Suspend
+
         descriptor is ConstructorDescriptor -> if (descriptor.constructedClass.isArray) {
             MethodBridge.ReturnValue.Instance.FactoryResult
         } else {
@@ -315,7 +317,10 @@ private fun ObjCExportMapper.bridgeMethodImpl(descriptor: FunctionDescriptor): M
     }
 
     val returnBridge = bridgeReturnType(descriptor, convertExceptionsToErrors)
-    if (convertExceptionsToErrors) {
+
+    if (descriptor.isSuspend) {
+        valueParameters += MethodBridgeValueParameter.SuspendCompletion
+    } else if (convertExceptionsToErrors) {
         // Add error out parameter before tail block parameters. The convention allows this.
         // Placing it after would trigger https://bugs.swift.org/browse/SR-12201
         // (see also https://github.com/JetBrains/kotlin-native/issues/3825).
@@ -332,6 +337,7 @@ private fun MethodBridgeValueParameter.isBlockPointer(): Boolean = when (this) {
         is BlockPointerBridge -> true
     }
     MethodBridgeValueParameter.ErrorOutParameter -> false
+    MethodBridgeValueParameter.SuspendCompletion -> true
 }
 
 internal fun ObjCExportMapper.bridgePropertyType(descriptor: PropertyDescriptor): TypeBridge {

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

@@ -453,13 +453,17 @@ internal class ObjCExportNamerImpl(
                         else -> it!!.name.asString().toIdentifier()
                     }
                     MethodBridgeValueParameter.ErrorOutParameter -> "error"
+                    MethodBridgeValueParameter.SuspendCompletion -> "completionHandler"
                 }
 
                 if (index == 0) {
                     append(when {
                         bridge is MethodBridgeValueParameter.ErrorOutParameter -> "AndReturn"
 
+                        bridge is MethodBridgeValueParameter.SuspendCompletion -> "With"
+
                         method is ConstructorDescriptor -> "With"
+
                         else -> ""
                     })
                     append(name.capitalize())
@@ -502,6 +506,7 @@ internal class ObjCExportNamerImpl(
                         else -> it!!.name.asString().toIdentifier()
                     }
                     MethodBridgeValueParameter.ErrorOutParameter -> continue@parameters
+                    MethodBridgeValueParameter.SuspendCompletion -> "completionHandler"
                 }
 
                 append(label)

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

@@ -163,3 +163,8 @@ internal fun ObjCType.makeNullableIfReferenceOrPointer(): ObjCType = when (this)
 
     is ObjCNullableReferenceType, is ObjCRawType, is ObjCPrimitiveType, ObjCVoidType -> this
 }
+
+internal fun ObjCReferenceType.makeNullable(): ObjCNullableReferenceType = when (this) {
+    is ObjCNonNullReferenceType -> ObjCNullableReferenceType(this)
+    is ObjCNullableReferenceType -> this
+}

backend.native/tests/framework/main.swift

@@ -37,6 +37,15 @@ func assertEquals<T: Equatable>(actual: T, expected: T,
     }
 }
 
+func assertSame(actual: AnyObject?, expected: AnyObject?,
+                _ message: String = "Assertion failed:",
+                file: String = #file, line: Int = #line) throws {
+    if (actual !== expected) {
+        try throwAssertFailed(message: message + " Expected value: \(expected), but got: \(actual)",
+                file: file, line: line)
+    }
+}
+
 func assertEquals<T: Equatable>(actual: [T], expected: [T],
                                 _ message: String = "Assertion failed: arrays not equal",
                                 file: String = #file, line: Int = #line) throws {