ObjCExport: optimize receiving autoreleased references from Obj-C/Swift

when the callee supports this.

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

@@ -9,6 +9,8 @@ import org.jetbrains.kotlin.backend.konan.llvm.LlvmParameterAttribute
 import org.jetbrains.kotlin.builtins.UnsignedType
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.util.classId
+import org.jetbrains.kotlin.konan.target.Architecture
+import org.jetbrains.kotlin.konan.target.KonanTarget
 
 // TODO: We have [org.jetbrains.kotlin.native.interop.gen.ObjCAbiInfo] which does similar thing.
 //  Consider unifying these classes.
@@ -96,4 +98,36 @@ class WindowsX64TargetAbiInfo : ExtendOnCalleeSideTargetAbiInfo(shouldZeroExtBoo
  * "Generic" ABI info that is applicable for the most of our current targets.
  * In time will be replaced by more specific [TargetAbiInfo] inheritors.
  */
-typealias DefaultTargetAbiInfo = ExtendOnCallerSideTargetAbiInfo
+typealias DefaultTargetAbiInfo = ExtendOnCallerSideTargetAbiInfo
+
+/// Equivalent to TargetCodeGenInfo.markARCOptimizedReturnCallsAsNoTail in Clang.
+///
+/// Determine whether a call to objc_retainAutoreleasedReturnValue or
+/// objc_unsafeClaimAutoreleasedReturnValue should be marked as 'notail'.
+fun KonanTarget.markARCOptimizedReturnCallsAsNoTail(): Boolean = when (this.architecture) {
+    Architecture.X64 -> {
+        /// Disable tail call on x86-64. The epilogue code before the tail jump blocks
+        /// autoreleaseRV/retainRV and autoreleaseRV/unsafeClaimRV optimizations.
+        true
+    }
+    else -> false
+}
+
+/// Equivalent to TargetCodeGenInfo.getARCRetainAutoreleasedReturnValueMarker in Clang.
+///
+/// Retrieve the address of a function to call immediately before
+/// calling objc_retainAutoreleasedReturnValue.  The
+/// implementation of objc_autoreleaseReturnValue sniffs the
+/// instruction stream following its return address to decide
+/// whether it's a call to objc_retainAutoreleasedReturnValue.
+/// This can be prohibitively expensive, depending on the
+/// relocation model, and so on some targets it instead sniffs for
+/// a particular instruction sequence.  This functions returns
+/// that instruction sequence in inline assembly, which will be
+/// empty if none is required.
+fun KonanTarget.getARCRetainAutoreleasedReturnValueMarker(): String? = when (this.architecture) {
+    Architecture.X86 -> "movl\t%ebp, %ebp\t\t// marker for objc_retainAutoreleaseReturnValue"
+    Architecture.ARM64 -> "mov\tfp, fp\t\t// marker for objc_retainAutoreleaseReturnValue"
+    Architecture.ARM32 -> "mov\tr7, r7\t\t// marker for objc_retainAutoreleaseReturnValue"
+    else -> null
+}

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

@@ -5,7 +5,12 @@
 
 package org.jetbrains.kotlin.backend.konan.llvm.objc
 
+import kotlinx.cinterop.signExtend
+import kotlinx.cinterop.toCValues
+import llvm.LLVMGetInlineAsm
+import llvm.LLVMInlineAsmDialect
 import llvm.LLVMValueRef
+import org.jetbrains.kotlin.backend.konan.getARCRetainAutoreleasedReturnValueMarker
 import org.jetbrains.kotlin.backend.konan.llvm.*
 
 internal open class ObjCCodeGenerator(val codegen: CodeGenerator) {
@@ -56,6 +61,30 @@ internal open class ObjCCodeGenerator(val codegen: CodeGenerator) {
             origin = context.stdlibModule.llvmSymbolOrigin
     ))
 
+    val objcRetainAutoreleasedReturnValue = context.llvm.externalFunction(LlvmFunctionProto(
+            "llvm.objc.retainAutoreleasedReturnValue",
+            LlvmRetType(int8TypePtr),
+            listOf(LlvmParamType(int8TypePtr)),
+            listOf(LlvmFunctionAttribute.NoUnwind),
+            origin = context.stdlibModule.llvmSymbolOrigin
+    ))
+
+    val objcRetainAutoreleasedReturnValueMarker: LLVMValueRef? by lazy {
+        // See emitAutoreleasedReturnValueMarker in Clang.
+        val asmString = codegen.context.config.target.getARCRetainAutoreleasedReturnValueMarker() ?: return@lazy null
+        val asmStringBytes = asmString.toByteArray()
+        LLVMGetInlineAsm(
+                Ty = functionType(voidType, false),
+                AsmString = asmStringBytes.toCValues(),
+                AsmStringSize = asmStringBytes.size.signExtend(),
+                Constraints = null,
+                ConstraintsSize = 0,
+                HasSideEffects = 1,
+                IsAlignStack = 0,
+                Dialect = LLVMInlineAsmDialect.LLVMInlineAsmDialectATT
+        )
+    }
+
     // TODO: this doesn't support stret.
     fun msgSender(functionType: LlvmFunctionSignature): LlvmCallable =
             LlvmCallable(

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

@@ -55,8 +55,13 @@ internal fun ObjCExportCodeGeneratorBase.generateBlockToKotlinFunctionConverter(
         switchThreadStateIfExperimentalMM(ThreadState.Native)
         // Using terminatingExceptionHandler, so any exception thrown by `invoke` will lead to the termination,
         // and switching the thread state back to `Runnable` on exceptional path is not required.
-        val result = call(invoke, listOf(blockPtr) + args, exceptionHandler = terminatingExceptionHandler)
-
+        val result = callAndMaybeRetainAutoreleased(
+                invoke,
+                bridge.blockType.blockInvokeLlvmType,
+                listOf(blockPtr) + args,
+                exceptionHandler = terminatingExceptionHandler,
+                doRetain = !bridge.returnsVoid
+        )
         args.forEach {
             objcReleaseFromNativeThreadState(it)
         }
@@ -66,7 +71,11 @@ internal fun ObjCExportCodeGeneratorBase.generateBlockToKotlinFunctionConverter(
         val kotlinResult = if (bridge.returnsVoid) {
             theUnitInstanceRef.llvm
         } else {
+            // TODO: in some cases the sequence below will have redundant retain-release pair.
+            // We could implement an optimized objCRetainedReferenceToKotlin, which takes ownership
+            // of its argument (i.e. consumes retained reference).
             objCReferenceToKotlin(result, Lifetime.RETURN_VALUE)
+                    .also { objcReleaseFromRunnableThreadState(result) }
         }
         ret(kotlinResult)
     }.also {

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.backend.konan.llvm.objcexport
 
+import kotlinx.cinterop.toCValues
 import kotlinx.cinterop.toKString
 import llvm.*
 import org.jetbrains.kotlin.backend.common.ir.allParameters
@@ -48,14 +49,10 @@ internal fun TypeBridge.makeNothing() = when (this) {
 }
 
 internal class ObjCExportFunctionGenerationContext(
-        builder: ObjCExportFunctionGenerationContextBuilder
+        builder: ObjCExportFunctionGenerationContextBuilder,
+        override val needCleanupLandingpadAndLeaveFrame: Boolean
 ) : FunctionGenerationContext(builder) {
-    private val objCExportCodegen = builder.objCExportCodegen
-
-    // All generated bridges by ObjCExport should have `LeaveFrame`
-    // because there is no guarantee of catching Kotlin exception in Kotlin code.
-    override val needCleanupLandingpadAndLeaveFrame: Boolean
-        get() = true
+    val objCExportCodegen = builder.objCExportCodegen
 
     // Note: we could generate single "epilogue" and make all [ret]s just branch to it (like [DefaultFunctionGenerationContext]),
     // but this would be useless for most of the usages, which have only one [ret].
@@ -127,7 +124,11 @@ internal class ObjCExportFunctionGenerationContextBuilder(
         functionType.addFunctionAttributes(function)
     }
 
-    override fun build() = ObjCExportFunctionGenerationContext(this)
+    // Unless specified otherwise, all generated bridges by ObjCExport should have `LeaveFrame`
+    // because there is no guarantee of catching Kotlin exception in Kotlin code.
+    var needCleanupLandingpadAndLeaveFrame = true
+
+    override fun build() = ObjCExportFunctionGenerationContext(this, needCleanupLandingpadAndLeaveFrame)
 }
 
 internal inline fun ObjCExportCodeGeneratorBase.functionGenerator(
@@ -140,6 +141,78 @@ internal inline fun ObjCExportCodeGeneratorBase.functionGenerator(
         this
 ).apply(configure)
 
+internal fun ObjCExportFunctionGenerationContext.callAndMaybeRetainAutoreleased(
+        function: LLVMValueRef,
+        signature: LlvmFunctionSignature,
+        args: List<LLVMValueRef>,
+        resultLifetime: Lifetime = Lifetime.IRRELEVANT,
+        exceptionHandler: ExceptionHandler,
+        doRetain: Boolean
+): LLVMValueRef {
+    if (!doRetain) return call(function, args, resultLifetime, exceptionHandler, attributeProvider = signature)
+
+    // Objective-C runtime provides "optimizable" return for autoreleased references:
+    // the caller (this code) handles the return value with objc_retainAutoreleasedReturnValue,
+    // and the callee tries to detect this by looking at the code location at the return address.
+    // The latter is implemented as tail calls to objc_retainAutoreleaseReturnValue or objc_autoreleaseReturnValue.
+    //
+    // These functions look for a specific pattern immediately following the call site.
+    // Depending on the platform, this pattern is either
+    // * move from the return value register to the argument register and call to objcRetainAutoreleasedReturnValue, or
+    // * special instruction (like `mov fp, fp`) that is not supposed to be generated in any other case.
+    //
+    // Unfortunately, we can't just generate this straightforwardly in LLVM,
+    // because we have to catch exceptions thrown by `function`.
+    // So we have to use `invoke` LLVM instructions. In this case LLVM sometimes inserts
+    // a redundant jump after the call when generating the machine code, ruining the expected code pattern.
+    //
+    // To workaround this, we generate the call and the pattern after it as a separate noinline function ("outlined"),
+    // and catch the exceptions in the caller of this function.
+    // So, no exception handler in "outlined" => no redundant jump => the optimized return works properly.
+
+    val functionIsPassedAsLastParameter = LLVMIsConstant(function) == 0
+
+    val valuesToPass = args + if (functionIsPassedAsLastParameter) listOf(function) else emptyList()
+    val outlinedType = LlvmFunctionSignature(
+            signature.returnType,
+            signature.parameterTypes + if (functionIsPassedAsLastParameter) listOf(LlvmParamType(function.type)) else emptyList()
+    )
+
+    val outlined = objCExportCodegen.functionGenerator(outlinedType, this.function.name.orEmpty() + "_outlined") {
+        setupBridgeDebugInfo()
+        // Don't generate redundant cleanup landingpad (the generation would fail due to forbidRuntime below):
+        needCleanupLandingpadAndLeaveFrame = false
+    }.generate {
+        forbidRuntime = true // Don't emit safe points, frame management etc.
+
+        val actualArgs = signature.parameterTypes.indices.map { param(it) }
+        val actualFunction = if (functionIsPassedAsLastParameter) param(signature.parameterTypes.size) else function
+
+        // Use LLVMBuildCall instead of call, because the latter enforces using exception handler, which is exactly what we have to avoid.
+        val result = LLVMBuildCall(builder, actualFunction, actualArgs.toCValues(), actualArgs.size, "")!!.also {
+            signature.addCallSiteAttributes(it)
+        }.let { callResult ->
+            // Simplified version of emitAutoreleasedReturnValueMarker in Clang:
+            objCExportCodegen.objcRetainAutoreleasedReturnValueMarker?.let {
+                LLVMBuildCall(arg0 = builder, Fn = it, Args = null, NumArgs = 0, Name = "")
+            }
+
+            call(objCExportCodegen.objcRetainAutoreleasedReturnValue, listOf(callResult)).also {
+                if (context.config.target.markARCOptimizedReturnCallsAsNoTail())
+                    LLVMSetNoTailCall(it)
+            }
+        }
+
+        ret(result)
+    }
+
+    outlinedType.addFunctionAttributes(outlined)
+    LLVMSetLinkage(outlined, LLVMLinkage.LLVMPrivateLinkage)
+    setFunctionNoInline(outlined)
+
+    return call(LlvmCallable(outlined, outlinedType), valuesToPass, resultLifetime, exceptionHandler)
+}
+
 internal open class ObjCExportCodeGeneratorBase(codegen: CodeGenerator) : ObjCCodeGenerator(codegen) {
     val symbols get() = context.ir.symbols
     val runtime get() = codegen.runtime
@@ -1159,8 +1232,6 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
 
     val parameterToBase = irFunction.allParameters.zip(baseIrFunction.allParameters).toMap()
 
-    val objcMsgSend = msgSender(objCFunctionType(context, methodBridge))
-
     val functionType = LlvmFunctionSignature(irFunction, codegen)
 
     val result = functionGenerator(functionType, "kotlin2objc").generate {
@@ -1241,9 +1312,21 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
         }
 
         switchThreadStateIfExperimentalMM(ThreadState.Native)
+
+        val retainAutoreleasedTargetResult = methodBridge.returnBridge.isAutoreleasedObjCReference()
+
+        val objCFunctionType = objCFunctionType(context, methodBridge)
+        val objcMsgSend = msgSender(objCFunctionType)
+
         // Using terminatingExceptionHandler, so any exception thrown by the method will lead to the termination,
         // and switching the thread state back to `Runnable` on exceptional path is not required.
-        val targetResult = call(objcMsgSend, objCArgs, exceptionHandler = terminatingExceptionHandler)
+        val targetResult = callAndMaybeRetainAutoreleased(
+                objcMsgSend.llvmValue,
+                objCFunctionType,
+                objCArgs,
+                exceptionHandler = terminatingExceptionHandler,
+                doRetain = retainAutoreleasedTargetResult
+        )
 
         objCReferenceArgsToRelease.forEach {
             objcReleaseFromNativeThreadState(it)
@@ -1285,6 +1368,7 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
 
             MethodBridge.ReturnValue.WithError.Success -> {
                 ifThen(icmpEq(targetResult, Int8(0).llvm)) {
+                    check(!retainAutoreleasedTargetResult)
                     rethrow()
                 }
                 null
@@ -1294,10 +1378,12 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
                 if (returnBridge.successMayBeZero) {
                     val error = load(errorOutPtr!!)
                     ifThen(icmpNe(error, kNullInt8Ptr)) {
+                        // error is not null, so targetResult should be null => no need for objc_release on it.
                         rethrow()
                     }
                 } else {
                     ifThen(icmpEq(targetResult, kNullInt8Ptr)) {
+                        // targetResult is null => no need for objc_release on it.
                         rethrow()
                     }
                 }
@@ -1343,6 +1429,14 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
                 )
         }
 
+        if (retainAutoreleasedTargetResult) {
+            // TODO: in some cases the return sequence will have redundant retain-release pair:
+            //  retain in the return value conversion and release here.
+            // We could implement an optimized objCRetainedReferenceToKotlin, which takes ownership
+            // of its argument (i.e. consumes retained reference).
+            objcReleaseFromRunnableThreadState(targetResult)
+        }
+
         ret(retVal)
     }
 
@@ -1351,6 +1445,20 @@ private fun ObjCExportCodeGenerator.generateKotlinToObjCBridge(
     return constPointer(result)
 }
 
+private fun MethodBridge.ReturnValue.isAutoreleasedObjCReference(): Boolean = when (this) {
+    MethodBridge.ReturnValue.HashCode, // integer
+    MethodBridge.ReturnValue.Instance.FactoryResult, // retained
+    MethodBridge.ReturnValue.Instance.InitResult, // retained
+    MethodBridge.ReturnValue.Suspend, // void
+    MethodBridge.ReturnValue.WithError.Success, // boolean
+    MethodBridge.ReturnValue.Void -> false
+    is MethodBridge.ReturnValue.Mapped -> when (this.bridge) {
+        is BlockPointerBridge, ReferenceBridge -> true
+        is ValueTypeBridge -> false
+    }
+    is MethodBridge.ReturnValue.WithError.ZeroForError -> this.successBridge.isAutoreleasedObjCReference()
+}
+
 private fun ObjCExportCodeGenerator.createReverseAdapter(
         irFunction: IrFunction,
         baseMethod: ObjCMethodSpec.BaseMethod<IrSimpleFunctionSymbol>,

kotlin-native/backend.native/tests/objcexport/noAutorelease.swift

@@ -351,10 +351,7 @@ private func testSendKotlinObjectToSwift() throws {
 }
 
 private func testSendKotlinObjectToSwiftBlock() throws {
-    // Getting Swift block in Kotlin still adds it to the autoreleasepool.
-    // The block is not our target, so just use `checkAutorelease: false` flag
-    // to disable the relevant check for now.
-    try testCallToSwift(flags: TestFlags(checkAutorelease: false)) {
+    try testCallToSwift {
         NoAutoreleaseKt.sendKotlinObjectToBlock(helper: $0, tracker: $1)
     }
 }
@@ -506,7 +503,6 @@ class NoAutoreleaseTests : SimpleTestProvider {
         test("testSendBlockToSwift", testSendBlockToSwift)
         test("testSendCompletionToSwift", testSendCompletionToSwift)
 
-#if false
         test("testReceiveKotlinObjectFromSwift", testReceiveKotlinObjectFromSwift)
         test("testReceiveSwiftObjectFromSwift", testReceiveSwiftObjectFromSwift)
         test("testReceiveListFromSwift", testReceiveListFromSwift)
@@ -514,6 +510,5 @@ class NoAutoreleaseTests : SimpleTestProvider {
         test("testReceiveNumberFromSwift", testReceiveNumberFromSwift)
         test("testReceiveBlockFromSwift", testReceiveBlockFromSwift)
         test("testReceiveBlockFromSwiftAndCall", testReceiveBlockFromSwiftAndCall)
-#endif
     }
 }

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

@@ -245,6 +245,13 @@ extern "C" const char* Kotlin_callsCheckerGoodFunctionNames[] = {
         "llvm.objc.autorelease",
         "llvm.objc.autoreleaseReturnValue",
         "llvm.objc.retain",
+
+        // Not used in Runnable state, but this would be ok.
+        // If we don't include it to the good functions list, the code generator will emit redundant state check at the callsite,
+        // and this would ruin the code: the state check would be inserted between retainAutoreleasedReturnValue and the actual call
+        // producing "autoreleased return value", so the latter won't be able to detect the former, and the autorelease elimination
+        // won't work.
+        "llvm.objc.retainAutoreleasedReturnValue",
         "llvm.objectsize.*",
         "llvm.pow.*",
         "llvm.rint.*",