JVM_IR: implement IrReturnableBlock codegen

(cherry picked from commit 530ad368fe)

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ArrayConstructorLowering.kt

@@ -11,8 +11,10 @@ import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.ir.builders.*
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
+import org.jetbrains.kotlin.ir.expressions.impl.IrReturnableBlockImpl
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.impl.IrReturnableBlockSymbolImpl
 import org.jetbrains.kotlin.ir.types.getClass
 import org.jetbrains.kotlin.ir.types.toKotlinType
 import org.jetbrains.kotlin.ir.util.*
@@ -33,31 +35,6 @@ class ArrayConstructorLowering(val context: CommonBackendContext) : IrElementTra
             fromInit to fromSize
         }
 
-    // Generate `array[index] = value`.
-    private fun IrBuilderWithScope.setItem(array: IrVariable, index: IrVariable, value: IrExpression) =
-        irCall(array.type.getClass()!!.functions.single { it.name.toString() == "set" }).apply {
-            dispatchReceiver = irGet(array)
-            putValueArgument(0, irGet(index))
-            putValueArgument(1, value)
-        }
-
-    // Generate `for (index in 0 until end) { element }`.
-    private fun IrBlockBuilder.fromZeroTo(end: IrVariable, element: IrBlockBuilder.(IrLoop, IrVariable) -> Unit) {
-        val index = irTemporaryVar(irInt(0))
-        +irWhile().apply {
-            condition = irCall(context.irBuiltIns.lessFunByOperandType[index.type.toKotlinType()]!!).apply {
-                putValueArgument(0, irGet(index))
-                putValueArgument(1, irGet(end))
-            }
-            body = irBlock {
-                val currentIndex = irTemporary(irGet(index))
-                val inc = index.type.getClass()!!.functions.single { it.name.asString() == "inc" }
-                +irSetVar(index.symbol, irCall(inc).apply { dispatchReceiver = irGet(index) })
-                element(this@apply, currentIndex)
-            }
-        }
-    }
-
     private fun IrExpression.asSingleArgumentLambda(): IrSimpleFunction? {
         // A lambda is represented as a block with a function declaration and a reference to it.
         if (this !is IrBlock || statements.size != 2)
@@ -75,8 +52,6 @@ class ArrayConstructorLowering(val context: CommonBackendContext) : IrElementTra
     override fun visitConstructorCall(expression: IrConstructorCall): IrExpression {
         val sizeConstructor = arrayInlineToSizeCtor[expression.symbol]
             ?: return super.visitConstructorCall(expression)
-        val arrayOfReferences = sizeConstructor == context.ir.symbols.arrayOfNulls
-
         // inline fun <reified T> Array(size: Int, invokable: (Int) -> T): Array<T> {
         //     val result = arrayOfNulls<T>(size)
         //     for (i in 0 until size) {
@@ -85,59 +60,61 @@ class ArrayConstructorLowering(val context: CommonBackendContext) : IrElementTra
         //     return result as Array<T>
         // }
         // (and similar for primitive arrays)
-        val loweringContext = context
         val size = expression.getValueArgument(0)!!.transform(this, null)
         val invokable = expression.getValueArgument(1)!!.transform(this, null)
         return context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol).irBlock(expression.startOffset, expression.endOffset) {
+            val index = irTemporaryVar(irInt(0))
             val sizeVar = irTemporary(size)
             val result = irTemporary(irCall(sizeConstructor, expression.type).apply {
-                if (arrayOfReferences) {
+                copyTypeArgumentsFrom(expression)
-                    putTypeArgument(0, expression.getTypeArgument(0))
-                }
                 putValueArgument(0, irGet(sizeVar))
             })
 
             val lambda = invokable.asSingleArgumentLambda()
-            if (lambda == null) {
+            val invoke = invokable.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
-                val invokableVar = irTemporary(invokable)
+            val invokableVar = if (lambda == null) irTemporary(invokable) else null
-                val invoke = invokable.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
+            +irWhile().apply {
-                fromZeroTo(sizeVar) { _, index ->
+                condition = irCall(context.irBuiltIns.lessFunByOperandType[index.type.toKotlinType()]!!).apply {
-                    +setItem(result, index, irCall(invoke).apply {
+                    putValueArgument(0, irGet(index))
-                        dispatchReceiver = irGet(invokableVar)
+                    putValueArgument(1, irGet(sizeVar))
-                        putValueArgument(0, irGet(index))
-                    })
                 }
-            } else {
+                body = irBlock {
-                // Inline `invokable` by replacing the argument with `i` and `return x` with `result[i] = x; continue`.
+                    val value =
-                fromZeroTo(sizeVar) { loop, index ->
+                        lambda?.inline(listOf(index)) ?: irCallOp(invoke.symbol, invoke.returnType, irGet(invokableVar!!), irGet(index))
-                    val body = lambda.body!!.transform(object : IrElementTransformerVoidWithContext() {
+                    +irCall(result.type.getClass()!!.functions.single { it.name == OperatorNameConventions.SET }).apply {
-                        override fun visitGetValue(expression: IrGetValue) =
+                        dispatchReceiver = irGet(result)
-                            if (expression.symbol == lambda.valueParameters[0].symbol)
+                        putValueArgument(0, irGet(index))
-                                IrGetValueImpl(expression.startOffset, expression.endOffset, index.symbol)
+                        putValueArgument(1, value)
-                            else
-                                super.visitGetValue(expression)
-
-                        override fun visitReturn(expression: IrReturn) =
-                            if (expression.returnTargetSymbol == lambda.symbol) {
-                                val value = expression.value.transform(this, null)
-                                val scope = currentScope?.scope?.scopeOwnerSymbol ?: lambda.symbol
-                                loweringContext.createIrBuilder(scope).irBlock(expression.startOffset, expression.endOffset) {
-                                    +setItem(result, index, value)
-                                    +irContinue(loop)
-                                }
-                            } else {
-                                super.visitReturn(expression)
-                            }
-                    }, null)
-
-                    when (body) {
-                        is IrExpressionBody -> +setItem(result, index, body.expression)
-                        is IrBlockBody -> body.statements.forEach { +it }
-                        else -> throw AssertionError("unexpected function body type: $body")
                     }
+                    val inc = index.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INC }
+                    +irSetVar(index.symbol, irCallOp(inc.symbol, index.type, irGet(index)))
                 }
             }
             +irGet(result)
         }
     }
+
+    // TODO use a generic inliner (e.g. JS/Native's FunctionInlining.Inliner)
+    private fun IrFunction.inline(arguments: List<IrValueDeclaration>): IrReturnableBlock {
+        val argumentMap = valueParameters.zip(arguments).toMap()
+        val blockSymbol = IrReturnableBlockSymbolImpl(descriptor)
+        val block = IrReturnableBlockImpl(startOffset, endOffset, returnType, blockSymbol, null, symbol)
+        val remapper = object : AbstractVariableRemapper() {
+            override fun remapVariable(value: IrValueDeclaration): IrValueDeclaration? =
+                argumentMap[value]
+
+            override fun visitReturn(expression: IrReturn): IrExpression = super.visitReturn(
+                if (expression.returnTargetSymbol == symbol)
+                    IrReturnImpl(expression.startOffset, expression.endOffset, expression.type, blockSymbol, expression.value)
+                else
+                    expression
+            )
+        }
+        when (val transformed = body?.transform(remapper, null)) {
+            is IrBlockBody -> block.statements += transformed.statements
+            is IrExpressionBody -> block.statements += transformed.expression
+            else -> throw AssertionError("unexpected body type: $this")
+        }
+        return block
+    }
 }

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/LowerUtils.kt

@@ -59,7 +59,7 @@ class DeclarationIrBuilder(
 )
 
 abstract class AbstractVariableRemapper : IrElementTransformerVoid() {
-    protected abstract fun remapVariable(value: IrValueDeclaration): IrValueParameter?
+    protected abstract fun remapVariable(value: IrValueDeclaration): IrValueDeclaration?
 
     override fun visitGetValue(expression: IrGetValue): IrExpression =
         remapVariable(expression.symbol.owner)?.let {
@@ -68,12 +68,12 @@ abstract class AbstractVariableRemapper : IrElementTransformerVoid() {
 }
 
 class VariableRemapper(val mapping: Map<IrValueParameter, IrValueParameter>) : AbstractVariableRemapper() {
-    override fun remapVariable(value: IrValueDeclaration): IrValueParameter? =
+    override fun remapVariable(value: IrValueDeclaration): IrValueDeclaration? =
         mapping[value]
 }
 
 class VariableRemapperDesc(val mapping: Map<ValueDescriptor, IrValueParameter>) : AbstractVariableRemapper() {
-    override fun remapVariable(value: IrValueDeclaration): IrValueParameter? =
+    override fun remapVariable(value: IrValueDeclaration): IrValueDeclaration? =
         mapping[value.descriptor]
 }
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/ExpressionCodegen.kt

@@ -52,17 +52,22 @@ class TryInfo(val onExit: IrExpression) : ExpressionInfo() {
     val gaps = mutableListOf<Pair<Label, Label>>()
 }
 
-class BlockInfo private constructor(val parent: BlockInfo?) {
+class ReturnableBlockInfo(
-    val variables = mutableListOf<VariableInfo>()
+    val returnLabel: Label,
-    val infos = Stack<ExpressionInfo>()
+    val returnSymbol: IrSymbol,
+    val returnTemporary: Int? = null
+) : ExpressionInfo()
 
-    fun create() = BlockInfo(this).apply {
+class BlockInfo(val parent: BlockInfo? = null) {
-        this@apply.infos.addAll(this@BlockInfo.infos)
+    val variables = mutableListOf<VariableInfo>()
-    }
+    private val infos: Stack<ExpressionInfo> = parent?.infos ?: Stack()
 
     fun hasFinallyBlocks(): Boolean = infos.firstIsInstanceOrNull<TryInfo>() != null
 
-    inline fun <T : ExpressionInfo, R> withBlock(info: T, f: (T) -> R): R {
+    internal inline fun <reified T : ExpressionInfo> findBlock(predicate: (T) -> Boolean): T? =
+        infos.find { it is T && predicate(it) } as? T
+
+    internal inline fun <T : ExpressionInfo, R> withBlock(info: T, f: (T) -> R): R {
         infos.add(info)
         try {
             return f(info)
@@ -71,7 +76,7 @@ class BlockInfo private constructor(val parent: BlockInfo?) {
         }
     }
 
-    inline fun <R> handleBlock(f: (ExpressionInfo) -> R): R? {
+    internal inline fun <R> handleBlock(f: (ExpressionInfo) -> R): R? {
         if (infos.isEmpty()) {
             return null
         }
@@ -82,10 +87,6 @@ class BlockInfo private constructor(val parent: BlockInfo?) {
             infos.add(top)
         }
     }
-
-    companion object {
-        fun create() = BlockInfo(null)
-    }
 }
 
 class VariableInfo(val declaration: IrVariable, val index: Int, val type: Type, val startLabel: Label)
@@ -153,7 +154,7 @@ class ExpressionCodegen(
     fun generate() {
         mv.visitCode()
         val startLabel = markNewLabel()
-        val info = BlockInfo.create()
+        val info = BlockInfo()
         val body = irFunction.body!!
         val result = body.accept(this, info)
         // If this function has an expression body, return the result of that expression.
@@ -169,25 +170,26 @@ class ExpressionCodegen(
             result.coerce(returnType).materialize()
             mv.areturn(returnType)
         }
-        writeLocalVariablesInTable(info)
+        val endLabel = markNewLabel()
-        writeParameterInLocalVariableTable(startLabel)
+        writeLocalVariablesInTable(info, endLabel)
+        writeParameterInLocalVariableTable(startLabel, endLabel)
         mv.visitEnd()
     }
 
-    private fun writeParameterInLocalVariableTable(startLabel: Label) {
+    private fun writeParameterInLocalVariableTable(startLabel: Label, endLabel: Label) {
         if (!irFunction.isStatic) {
-            mv.visitLocalVariable("this", classCodegen.type.descriptor, null, startLabel, markNewLabel(), 0)
+            mv.visitLocalVariable("this", classCodegen.type.descriptor, null, startLabel, endLabel, 0)
         }
         val extensionReceiverParameter = irFunction.extensionReceiverParameter
         if (extensionReceiverParameter != null) {
-            writeValueParameterInLocalVariableTable(extensionReceiverParameter, startLabel)
+            writeValueParameterInLocalVariableTable(extensionReceiverParameter, startLabel, endLabel)
         }
         for (param in irFunction.valueParameters) {
-            writeValueParameterInLocalVariableTable(param, startLabel)
+            writeValueParameterInLocalVariableTable(param, startLabel, endLabel)
         }
     }
 
-    private fun writeValueParameterInLocalVariableTable(param: IrValueParameter, startLabel: Label) {
+    private fun writeValueParameterInLocalVariableTable(param: IrValueParameter, startLabel: Label, endLabel: Label) {
         // TODO: old code has a special treatment for destructuring lambda parameters.
         // There is no (easy) way to reproduce it with IR structures.
         // Does not show up in tests, but might come to bite us at some point.
@@ -196,22 +198,43 @@ class ExpressionCodegen(
         val type = typeMapper.mapType(param)
         // NOTE: we expect all value parameters to be present in the frame.
         mv.visitLocalVariable(
-            name, type.descriptor, null, startLabel, markNewLabel(), findLocalIndex(param.symbol)
+            name, type.descriptor, null, startLabel, endLabel, findLocalIndex(param.symbol)
         )
     }
 
     override fun visitBlock(expression: IrBlock, data: BlockInfo): PromisedValue {
         if (expression.isTransparentScope)
             return super.visitBlock(expression, data)
-        val info = data.create()
+        val info = BlockInfo(data)
-        // Force materialization to avoid reading from out-of-scope variables.
+        return if (expression is IrReturnableBlock) {
-        return super.visitBlock(expression, info).materialized.apply {
+            val returnType = expression.asmType
-            writeLocalVariablesInTable(info)
+            val returnLabel = Label()
+            // Because the return might be inside an expression, it will need to pop excess items
+            // before jumping and store the result in a temporary variable.
+            val returnTemporary = if (returnType != Type.VOID_TYPE) frameMap.enterTemp(returnType) else null
+            info.withBlock(ReturnableBlockInfo(returnLabel, expression.symbol, returnTemporary)) {
+                // Remember current stack depth.
+                mv.fakeAlwaysFalseIfeq(returnLabel)
+                super.visitBlock(expression, info).materialized.also {
+                    returnTemporary?.let { mv.store(it, returnType) }
+                    // Variables leave the scope in reverse order, so must write locals first.
+                    mv.mark(returnLabel)
+                    writeLocalVariablesInTable(info, returnLabel)
+                    returnTemporary?.let {
+                        mv.load(it, returnType)
+                        frameMap.leaveTemp(returnType)
+                    }
+                }
+            }
+        } else {
+            // Force materialization to avoid reading from out-of-scope variables.
+            super.visitBlock(expression, info).materialized.also {
+                writeLocalVariablesInTable(info, markNewLabel())
+            }
         }
     }
 
-    private fun writeLocalVariablesInTable(info: BlockInfo) {
+    private fun writeLocalVariablesInTable(info: BlockInfo, endLabel: Label) {
-        val endLabel = markNewLabel()
         info.variables.forEach {
             when (it.declaration.origin) {
                 IrDeclarationOrigin.IR_TEMPORARY_VARIABLE,
@@ -601,15 +624,21 @@ class ExpressionCodegen(
             return voidValue
         }
 
+        val target = data.findBlock<ReturnableBlockInfo> { it.returnSymbol == expression.returnTargetSymbol }
         val returnType = typeMapper.mapReturnType(owner)
         val afterReturnLabel = Label()
         expression.value.accept(this, data).coerce(returnType).materialize()
-        generateFinallyBlocksIfNeeded(returnType, afterReturnLabel, data)
+        generateFinallyBlocksIfNeeded(returnType, afterReturnLabel, data, target)
         expression.markLineNumber(startOffset = true)
-        if (isNonLocalReturn) {
+        if (target != null) {
-            generateGlobalReturnFlag(mv, (owner as IrFunction).name.asString())
+            target.returnTemporary?.let { mv.store(it, returnType) }
+            mv.fixStackAndJump(target.returnLabel)
+        } else {
+            if (isNonLocalReturn) {
+                generateGlobalReturnFlag(mv, (owner as IrFunction).name.asString())
+            }
+            mv.areturn(returnType)
         }
-        mv.areturn(returnType)
         mv.mark(afterReturnLabel)
         mv.nop()/*TODO check RESTORE_STACK_IN_TRY_CATCH processor*/
         return voidValue
@@ -797,20 +826,18 @@ class ExpressionCodegen(
         return voidValue
     }
 
-    private fun unwindBlockStack(endLabel: Label, data: BlockInfo, loop: IrLoop? = null): LoopInfo? {
+    private fun unwindBlockStack(endLabel: Label, data: BlockInfo, stop: (ExpressionInfo) -> Boolean): ExpressionInfo? {
         return data.handleBlock {
-            when {
+            if (it is TryInfo)
-                it is TryInfo -> genFinallyBlock(it, null, endLabel, data)
+                genFinallyBlock(it, null, endLabel, data)
-                it is LoopInfo && it.loop == loop -> return it
+            return if (stop(it)) it else unwindBlockStack(endLabel, data, stop)
-            }
-            unwindBlockStack(endLabel, data, loop)
         }
     }
 
     override fun visitBreakContinue(jump: IrBreakContinue, data: BlockInfo): PromisedValue {
         jump.markLineNumber(startOffset = true)
         val endLabel = Label()
-        val stackElement = unwindBlockStack(endLabel, data, jump.loop)
+        val stackElement = unwindBlockStack(endLabel, data) { it is LoopInfo && it.loop == jump.loop } as LoopInfo?
             ?: throw AssertionError("Target label for break/continue not found")
         mv.fixStackAndJump(if (jump is IrBreak) stackElement.breakLabel else stackElement.continueLabel)
         mv.mark(endLabel)
@@ -942,16 +969,16 @@ class ExpressionCodegen(
         tryInfo.gaps.add(gapStart to (afterJumpLabel ?: markNewLabel()))
     }
 
-    fun generateFinallyBlocksIfNeeded(returnType: Type, afterReturnLabel: Label, data: BlockInfo) {
+    fun generateFinallyBlocksIfNeeded(returnType: Type, afterReturnLabel: Label, data: BlockInfo, target: ReturnableBlockInfo? = null) {
         if (data.hasFinallyBlocks()) {
             if (Type.VOID_TYPE != returnType) {
                 val returnValIndex = frameMap.enterTemp(returnType)
                 mv.store(returnValIndex, returnType)
-                unwindBlockStack(afterReturnLabel, data, null)
+                unwindBlockStack(afterReturnLabel, data) { it == target }
                 mv.load(returnValIndex, returnType)
                 frameMap.leaveTemp(returnType)
             } else {
-                unwindBlockStack(afterReturnLabel, data, null)
+                unwindBlockStack(afterReturnLabel, data) { it == target }
             }
         }
     }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrInlineCodegen.kt

@@ -54,7 +54,7 @@ class IrInlineCodegen(
                 (argumentExpression as IrBlock).statements.filterIsInstance<IrFunctionReference>().single()
             rememberClosure(irReference, parameterType, irValueParameter) as IrExpressionLambdaImpl
         } else {
-            putValueOnStack(argumentExpression, parameterType, irValueParameter?.index ?: -1)
+            putValueOnStack(argumentExpression, parameterType, irValueParameter?.index ?: -1, blockInfo)
         }
     }
 
@@ -71,14 +71,14 @@ class IrInlineCodegen(
     }
 
     private fun putCapturedValueOnStack(argumentExpression: IrExpression, valueType: Type, capturedParamIndex: Int) {
-        val onStack = codegen.gen(argumentExpression, valueType, BlockInfo.create())
+        val onStack = codegen.gen(argumentExpression, valueType, BlockInfo())
         putArgumentOrCapturedToLocalVal(
             JvmKotlinType(onStack.type, onStack.kotlinType), onStack, capturedParamIndex, capturedParamIndex, ValueKind.CAPTURED
         )
     }
 
-    private fun putValueOnStack(argumentExpression: IrExpression, valueType: Type, paramIndex: Int) {
+    private fun putValueOnStack(argumentExpression: IrExpression, valueType: Type, paramIndex: Int, blockInfo: BlockInfo) {
-        val onStack = codegen.gen(argumentExpression, valueType, BlockInfo.create())
+        val onStack = codegen.gen(argumentExpression, valueType, blockInfo)
         putArgumentOrCapturedToLocalVal(JvmKotlinType(onStack.type, onStack.kotlinType), onStack, -1, paramIndex, ValueKind.CAPTURED)
     }