Native: implement autoreleaseAndRet for ObjCExportCodeGenerator

This operation is equivalent to "autorelease and then return",
but it supports Obj-C runtime autorelease optimization for return value:
if the caller is optimized (which is usually true), then no autorelease
will happen at runtime.

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

@@ -132,7 +132,7 @@ internal inline fun generateFunction(
 
     val isCToKotlinBridge = function.origin == CBridgeOrigin.C_TO_KOTLIN_BRIDGE
 
-    val functionGenerationContext = FunctionGenerationContext(
+    val functionGenerationContext = DefaultFunctionGenerationContext(
             llvmFunction,
             codegen,
             startLocation,
@@ -175,7 +175,7 @@ internal inline fun generateFunction(
         switchToRunnable: Boolean = false,
         code: FunctionGenerationContext.() -> Unit
 ) {
-    val functionGenerationContext = FunctionGenerationContext(
+    val functionGenerationContext = DefaultFunctionGenerationContext(
             function,
             codegen,
             startLocation,
@@ -212,7 +212,7 @@ internal inline fun generateFunctionNoRuntime(
         function: LLVMValueRef,
         code: FunctionGenerationContext.() -> Unit,
 ) {
-    val functionGenerationContext = FunctionGenerationContext(function, codegen, null, null, switchToRunnable = false)
+    val functionGenerationContext = DefaultFunctionGenerationContext(function, codegen, null, null, switchToRunnable = false)
     try {
         functionGenerationContext.forbidRuntime = true
         require(!functionGenerationContext.isObjectType(functionGenerationContext.returnType!!)) {
@@ -425,7 +425,7 @@ internal abstract class FunctionGenerationContextBuilder<T : FunctionGenerationC
     abstract fun build(): T
 }
 
-internal open class FunctionGenerationContext(
+internal abstract class FunctionGenerationContext(
         val function: LLVMValueRef,
         val codegen: CodeGenerator,
         private val startLocation: LocationInfo?,
@@ -453,7 +453,6 @@ internal open class FunctionGenerationContext(
     }
 
     var returnType: LLVMTypeRef? = LLVMGetReturnType(getFunctionType(function))
-    private val returns: MutableMap<LLVMBasicBlockRef, LLVMValueRef> = mutableMapOf()
     val constructedClass: IrClass?
         get() = (irFunction as? IrConstructor)?.constructedClass
     private var returnSlot: LLVMValueRef? = null
@@ -470,8 +469,7 @@ internal open class FunctionGenerationContext(
     private val localsInitBb = basicBlockInFunction("locals_init", null)
     private val stackLocalsInitBb = basicBlockInFunction("stack_locals_init", null)
     private val entryBb = basicBlockInFunction("entry", startLocation)
-    private val epilogueBb = basicBlockInFunction("epilogue", endLocation)
-    private val cleanupLandingpad = basicBlockInFunction("cleanup_landingpad", endLocation)
+    protected val cleanupLandingpad = basicBlockInFunction("cleanup_landingpad", endLocation)
 
     val stackLocalsManager = StackLocalsManagerImpl(this, stackLocalsInitBb)
 
@@ -509,7 +507,7 @@ internal open class FunctionGenerationContext(
         currentPositionHolder.dispose()
     }
 
-    private fun basicBlockInFunction(name: String, locationInfo: LocationInfo?): LLVMBasicBlockRef {
+    protected fun basicBlockInFunction(name: String, locationInfo: LocationInfo?): LLVMBasicBlockRef {
         val bb = LLVMAppendBasicBlockInContext(llvmContext, function, name)!!
         update(bb, locationInfo)
         return bb
@@ -551,19 +549,7 @@ internal open class FunctionGenerationContext(
     }
 
 
-    fun ret(value: LLVMValueRef?): LLVMValueRef {
-        val res = LLVMBuildBr(builder, epilogueBb)!!
-        if (returns.containsKey(currentBlock)) {
-            // TODO: enable error throwing.
-            throw Error("ret() in the same basic block twice! in ${function.name}")
-        }
-
-        if (value != null)
-            returns[currentBlock] = value
-
-        currentPositionHolder.setAfterTerminator()
-        return res
-    }
+    abstract fun ret(value: LLVMValueRef?): LLVMValueRef
 
     fun param(index: Int): LLVMValueRef = LLVMGetParam(this.function, index)!!
 
@@ -1377,7 +1363,6 @@ internal open class FunctionGenerationContext(
     }
 
     internal fun prologue() {
-        assert(returns.isEmpty())
         if (isObjectType(returnType!!)) {
             returnSlot = LLVMGetParam(function, numParameters(function.type) - 1)
         }
@@ -1447,28 +1432,7 @@ internal open class FunctionGenerationContext(
             br(entryBb)
         }
 
-        appendingTo(epilogueBb) {
-            when {
-                returnType == voidType -> {
-                    releaseVars()
-                    assert(returnSlot == null)
-                    handleEpilogueForExperimentalMM(context.llvm.Kotlin_mm_safePointFunctionEpilogue)
-                    LLVMBuildRetVoid(builder)
-                }
-                returns.isNotEmpty() -> {
-                    val returnPhi = phi(returnType!!)
-                    addPhiIncoming(returnPhi, *returns.toList().toTypedArray())
-                    if (returnSlot != null) {
-                        updateReturnRef(returnPhi, returnSlot!!)
-                    }
-                    releaseVars()
-                    handleEpilogueForExperimentalMM(context.llvm.Kotlin_mm_safePointFunctionEpilogue)
-                    LLVMBuildRet(builder, returnPhi)
-                }
-                // Do nothing, all paths throw.
-                else -> LLVMBuildUnreachable(builder)
-            }
-        }
+        processReturns()
 
         val shouldHaveCleanupLandingpad = forwardingForeignExceptionsTerminatedWith != null ||
                 needSlots ||
@@ -1526,12 +1490,38 @@ internal open class FunctionGenerationContext(
             LLVMDeleteBasicBlock(cleanupLandingpad)
         }
 
-        returns.clear()
         vars.clear()
         returnSlot = null
         slotsPhi = null
     }
 
+    protected abstract fun processReturns()
+
+    protected fun retValue(value: LLVMValueRef): LLVMValueRef {
+        if (returnSlot != null) {
+            updateReturnRef(value, returnSlot!!)
+        }
+        onReturn()
+        return rawRet(value)
+    }
+
+    protected fun rawRet(value: LLVMValueRef): LLVMValueRef = LLVMBuildRet(builder, value)!!.also {
+        currentPositionHolder.setAfterTerminator()
+    }
+
+    protected fun retVoid(): LLVMValueRef {
+        check(returnSlot == null)
+        onReturn()
+        return LLVMBuildRetVoid(builder)!!.also {
+            currentPositionHolder.setAfterTerminator()
+        }
+    }
+
+    protected fun onReturn() {
+        releaseVars()
+        handleEpilogueForExperimentalMM(context.llvm.Kotlin_mm_safePointFunctionEpilogue)
+    }
+
     private fun handleEpilogueForExperimentalMM(safePointFunction: LLVMValueRef) {
         if (context.memoryModel == MemoryModel.EXPERIMENTAL) {
             if (!forbidRuntime) {
@@ -1681,5 +1671,58 @@ internal open class FunctionGenerationContext(
     }
 }
 
+internal class DefaultFunctionGenerationContext(
+        function: LLVMValueRef,
+        codegen: CodeGenerator,
+        startLocation: LocationInfo?,
+        endLocation: LocationInfo?,
+        switchToRunnable: Boolean,
+        irFunction: IrFunction? = null
+) : FunctionGenerationContext(
+        function,
+        codegen,
+        startLocation,
+        endLocation,
+        switchToRunnable,
+        irFunction
+) {
+    // Note: return handling can be extracted to a separate class.
 
+    private val returns: MutableMap<LLVMBasicBlockRef, LLVMValueRef> = mutableMapOf()
 
+    private val epilogueBb = basicBlockInFunction("epilogue", endLocation).also {
+        LLVMMoveBasicBlockBefore(it, cleanupLandingpad) // Just to make the produced code a bit more readable.
+    }
+
+    override fun ret(value: LLVMValueRef?): LLVMValueRef {
+        val res = br(epilogueBb)
+
+        if (returns.containsKey(currentBlock)) {
+            // TODO: enable error throwing.
+            throw Error("ret() in the same basic block twice! in ${function.name}")
+        }
+
+        if (value != null)
+            returns[currentBlock] = value
+
+        return res
+    }
+
+    override fun processReturns() {
+        appendingTo(epilogueBb) {
+            when {
+                returnType == voidType -> {
+                    retVoid()
+                }
+                returns.isNotEmpty() -> {
+                    val returnPhi = phi(returnType!!)
+                    addPhiIncoming(returnPhi, *returns.toList().toTypedArray())
+                    retValue(returnPhi)
+                }
+                // Do nothing, all paths throw.
+                else -> unreachable()
+            }
+        }
+        returns.clear()
+    }
+}

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

@@ -35,6 +35,14 @@ internal open class ObjCCodeGenerator(val codegen: CodeGenerator) {
             context.stdlibModule.llvmSymbolOrigin
     )
 
+    val objcAutoreleaseReturnValue = context.llvm.externalFunction(
+            "llvm.objc.autoreleaseReturnValue",
+            functionType(int8TypePtr, false, int8TypePtr),
+            context.stdlibModule.llvmSymbolOrigin
+    ).also {
+        setFunctionNoUnwind(it)
+    }
+
     // TODO: this doesn't support stret.
     fun msgSender(functionType: LLVMTypeRef): LLVMValueRef =
             objcMsgSend.bitcast(pointerType(functionType)).llvm

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

@@ -52,6 +52,37 @@ internal class ObjCExportFunctionGenerationContext(
 ) : FunctionGenerationContext(builder) {
     private 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].
+    // Remaining usages can be optimized ad hoc.
+
+    override fun ret(value: LLVMValueRef?): LLVMValueRef = if (value == null) {
+        retVoid()
+    } else {
+        retValue(value)
+    }
+
+    /**
+     * autoreleases and returns [value].
+     * It is equivalent to `ret(autorelease(value))`, but optimizes the autorelease out if the caller is prepared for it.
+     *
+     * See the Clang documentation and the Obj-C runtime source code for more details:
+     * https://clang.llvm.org/docs/AutomaticReferenceCounting.html#arc-runtime-objc-autoreleasereturnvalue
+     * https://github.com/opensource-apple/objc4/blob/cd5e62a5597ea7a31dccef089317abb3a661c154/runtime/objc-object.h#L930
+     */
+    fun autoreleaseAndRet(value: LLVMValueRef) {
+        onReturn()
+        // Note: it is important to make this call tail (otherwise the elimination magic won't work),
+        // so it should go after other "epilogue" instructions, and that's why we couldn't just use
+        // ret(autorelease(value))
+        val result = call(objCExportCodegen.objcAutoreleaseReturnValue, listOf(value))
+        LLVMSetTailCall(result, 1)
+        rawRet(result)
+    }
+
+    override fun processReturns() {
+        // Do nothing.
+    }
 }
 
 internal class ObjCExportFunctionGenerationContextBuilder(