Make ReturnableBlockLowering common

and remove special handling in JVM_IR codegen.

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

@@ -3,14 +3,15 @@
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
-package org.jetbrains.kotlin.ir.backend.js.lower.inline
+package org.jetbrains.kotlin.backend.common.lower
 
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
-import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
+import org.jetbrains.kotlin.ir.builders.irBoolean
-import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.builders.irBreak
-import org.jetbrains.kotlin.ir.declarations.IrDeclarationParent
+import org.jetbrains.kotlin.ir.builders.irGet
+import org.jetbrains.kotlin.ir.builders.irSetVar
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.expressions.IrContainerExpression
 import org.jetbrains.kotlin.ir.expressions.IrExpression
@@ -20,7 +21,6 @@ import org.jetbrains.kotlin.ir.expressions.impl.IrBlockImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrDoWhileLoopImpl
 import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
-import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 
 /**
  * Replaces returnable blocks and `return`'s with loops and `break`'s correspondingly.
@@ -71,38 +71,25 @@ import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
  */
 class ReturnableBlockLowering(val context: CommonBackendContext) : FileLoweringPass {
     override fun lower(irFile: IrFile) {
-        irFile.transform(ReturnableBlockTransformer(context), ReturnableBlockLoweringContext(irFile))
+        irFile.transform(ReturnableBlockTransformer(context), null)
     }
 }
 
-private class ReturnableBlockLoweringContext(val containingDeclaration: IrDeclarationParent) {
+private class ReturnableBlockTransformer(val context: CommonBackendContext) : IrElementTransformerVoidWithContext() {
-    var labelCnt = 0
+    private var labelCnt = 0
-    val returnMap = mutableMapOf<IrReturnableBlockSymbol, (IrReturn) -> IrExpression>()
+    private val returnMap = mutableMapOf<IrReturnableBlockSymbol, (IrReturn) -> IrExpression>()
-}
 
-private class ReturnableBlockTransformer(
+    override fun visitReturn(expression: IrReturn): IrExpression {
-    val context: CommonBackendContext
+        expression.transformChildrenVoid()
-) : IrElementTransformer<ReturnableBlockLoweringContext> {
+        return returnMap[expression.returnTargetSymbol]?.invoke(expression) ?: expression
-
-    override fun visitReturn(expression: IrReturn, data: ReturnableBlockLoweringContext): IrExpression {
-        expression.transformChildren(this, data)
-        return data.returnMap[expression.returnTargetSymbol]?.invoke(expression) ?: expression
     }
 
-    override fun visitDeclaration(declaration: IrDeclaration, data: ReturnableBlockLoweringContext): IrStatement {
+    override fun visitContainerExpression(expression: IrContainerExpression): IrExpression {
-        if (declaration is IrDeclarationParent) {
+        if (expression !is IrReturnableBlock) return super.visitContainerExpression(expression)
-            declaration.transformChildren(this, ReturnableBlockLoweringContext(declaration))
-        }
-        return super.visitDeclaration(declaration, data)
-    }
-
-    private val constFalse get() = JsIrBuilder.buildBoolean(context.irBuiltIns.booleanType, false)
-
-    override fun visitContainerExpression(expression: IrContainerExpression, data: ReturnableBlockLoweringContext): IrExpression {
-        if (expression !is IrReturnableBlock) return super.visitContainerExpression(expression, data)
 
+        val builder = context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol)
         val variable by lazy {
-            JsIrBuilder.buildVar(expression.type, data.containingDeclaration, "tmp\$ret\$${data.labelCnt++}", true)
+            builder.scope.createTemporaryVariableDeclaration(expression.type, "tmp\$ret\$${labelCnt++}", true)
         }
 
         val loop by lazy {
@@ -112,38 +99,33 @@ private class ReturnableBlockTransformer(
                 context.irBuiltIns.unitType,
                 expression.origin
             ).apply {
-                label = "l\$ret\$${data.labelCnt++}"
+                label = "l\$ret\$${labelCnt++}"
-                condition = constFalse
+                condition = builder.irBoolean(false)
             }
         }
 
         var hasReturned = false
 
-        data.returnMap[expression.symbol] = { returnExpression ->
+        returnMap[expression.symbol] = { returnExpression ->
             hasReturned = true
-
+            builder.irComposite(returnExpression) {
-            IrCompositeImpl(
+                +irSetVar(variable.symbol, returnExpression.value)
-                returnExpression.startOffset,
+                +irBreak(loop)
-                returnExpression.endOffset,
-                context.irBuiltIns.unitType
-            ).apply {
-                statements += JsIrBuilder.buildSetVariable(variable.symbol, returnExpression.value, context.irBuiltIns.unitType)
-                statements += JsIrBuilder.buildBreak(context.irBuiltIns.unitType, loop)
             }
         }
 
         val newStatements = expression.statements.mapIndexed { i, s ->
             if (i == expression.statements.lastIndex && s is IrReturn && s.returnTargetSymbol == expression.symbol) {
-                s.transformChildren(this, data)
+                s.transformChildrenVoid()
                 if (!hasReturned) s.value else {
-                    JsIrBuilder.buildSetVariable(variable.symbol, s.value, context.irBuiltIns.unitType)
+                    builder.irSetVar(variable.symbol, s.value)
                 }
             } else {
-                s.transform(this, data)
+                s.transform(this, null)
             }
         }
 
-        data.returnMap.remove(expression.symbol)
+        returnMap.remove(expression.symbol)
 
         if (!hasReturned) {
             return IrCompositeImpl(
@@ -162,15 +144,10 @@ private class ReturnableBlockTransformer(
                 newStatements
             )
 
-            return IrCompositeImpl(
+            return builder.irComposite(expression, expression.origin) {
-                expression.startOffset,
+                +variable
-                expression.endOffset,
+                +loop
-                expression.type,
+                +irGet(variable)
-                expression.origin
-            ).apply {
-                statements += variable
-                statements += loop
-                statements += JsIrBuilder.buildGetValue(variable.symbol)
             }
         }
     }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsLoweringPhases.kt

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.ir.backend.js.lower.*
 import org.jetbrains.kotlin.ir.backend.js.lower.calls.CallsLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.coroutines.JsSuspendFunctionsLowering
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.RemoveInlineFunctionsWithReifiedTypeParametersLowering
-import org.jetbrains.kotlin.ir.backend.js.lower.inline.ReturnableBlockLowering
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmLower.kt

@@ -121,6 +121,13 @@ private val innerClassesPhase = makeIrFilePhase(
     prerequisite = setOf(localDeclarationsPhase)
 )
 
+private val returnableBlocksPhase = makeIrFilePhase(
+    ::ReturnableBlockLowering,
+    name = "ReturnableBlock",
+    description = "Replace returnable blocks with do-while(false) loops",
+    prerequisite = setOf(arrayConstructorPhase, assertionPhase)
+)
+
 private val jvmFilePhases =
         stripTypeAliasDeclarationsPhase then
         provisionalFunctionExpressionPhase then
@@ -141,6 +148,7 @@ private val jvmFilePhases =
         renameFieldsPhase then
         assertionPhase then
         tailrecPhase then
+        returnableBlocksPhase then
 
         jvmInlineClassPhase then
 

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

@@ -62,21 +62,12 @@ class TryInfo(val onExit: IrExpression) : ExpressionInfo() {
     val gaps = mutableListOf<Pair<Label, Label>>()
 }
 
-class ReturnableBlockInfo(
-    val returnLabel: Label,
-    val returnSymbol: IrSymbol,
-    val returnTemporary: Int? = null
-) : ExpressionInfo()
-
 class BlockInfo(val parent: BlockInfo? = null) {
     val variables = mutableListOf<VariableInfo>()
     private val infos: Stack<ExpressionInfo> = parent?.infos ?: Stack()
 
     fun hasFinallyBlocks(): Boolean = infos.firstIsInstanceOrNull<TryInfo>() != null
 
-    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 {
@@ -257,34 +248,13 @@ class ExpressionCodegen(
     }
 
     override fun visitBlock(expression: IrBlock, data: BlockInfo): PromisedValue {
+        assert(expression !is IrReturnableBlock) { "unlowered returnable block: ${expression.dump()}" }
         if (expression.isTransparentScope)
             return super.visitBlock(expression, data)
         val info = BlockInfo(data)
-        return if (expression is IrReturnableBlock) {
+        // Force materialization to avoid reading from out-of-scope variables.
-            val returnType = expression.asmType
+        return super.visitBlock(expression, info).materialized.also {
-            val returnLabel = Label()
+            writeLocalVariablesInTable(info, markNewLabel())
-            // 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())
-            }
         }
     }
 
@@ -571,21 +541,15 @@ class ExpressionCodegen(
             return immaterialUnitValue
         }
 
-        val target = data.findBlock<ReturnableBlockInfo> { it.returnSymbol == expression.returnTargetSymbol }
         val returnType = methodSignatureMapper.mapReturnType(owner)
         val afterReturnLabel = Label()
         expression.value.accept(this, data).coerce(returnType, owner.returnType).materialize()
-        generateFinallyBlocksIfNeeded(returnType, afterReturnLabel, data, target)
+        generateFinallyBlocksIfNeeded(returnType, afterReturnLabel, data)
         expression.markLineNumber(startOffset = true)
-        if (target != null) {
+        if (isNonLocalReturn) {
-            target.returnTemporary?.let { mv.store(it, returnType) }
+            generateGlobalReturnFlag(mv, owner.name.asString())
-            mv.fixStackAndJump(target.returnLabel)
-        } else {
-            if (isNonLocalReturn) {
-                generateGlobalReturnFlag(mv, owner.name.asString())
-            }
-            mv.areturn(returnType)
         }
+        mv.areturn(returnType)
         mv.mark(afterReturnLabel)
         mv.nop()/*TODO check RESTORE_STACK_IN_TRY_CATCH processor*/
         return immaterialUnitValue
@@ -757,7 +721,7 @@ class ExpressionCodegen(
         return immaterialUnitValue
     }
 
-    private fun unwindBlockStack(endLabel: Label, data: BlockInfo, stop: (ExpressionInfo) -> Boolean): ExpressionInfo? {
+    private fun unwindBlockStack(endLabel: Label, data: BlockInfo, stop: (ExpressionInfo) -> Boolean = { false }): ExpressionInfo? {
         return data.handleBlock {
             if (it is TryInfo)
                 genFinallyBlock(it, null, endLabel, data)
@@ -900,16 +864,16 @@ class ExpressionCodegen(
         tryInfo.gaps.add(gapStart to (afterJumpLabel ?: markNewLabel()))
     }
 
-    fun generateFinallyBlocksIfNeeded(returnType: Type, afterReturnLabel: Label, data: BlockInfo, target: ReturnableBlockInfo? = null) {
+    fun generateFinallyBlocksIfNeeded(returnType: Type, afterReturnLabel: Label, data: BlockInfo) {
         if (data.hasFinallyBlocks()) {
             if (Type.VOID_TYPE != returnType) {
                 val returnValIndex = frameMap.enterTemp(returnType)
                 mv.store(returnValIndex, returnType)
-                unwindBlockStack(afterReturnLabel, data) { it == target }
+                unwindBlockStack(afterReturnLabel, data)
                 mv.load(returnValIndex, returnType)
                 frameMap.leaveTemp(returnType)
             } else {
-                unwindBlockStack(afterReturnLabel, data) { it == target }
+                unwindBlockStack(afterReturnLabel, data)
             }
         }
     }

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmLoweringPhases.kt

@@ -15,7 +15,6 @@ import org.jetbrains.kotlin.backend.wasm.lower.excludeDeclarationsFromCodegen
 import org.jetbrains.kotlin.ir.backend.js.JsLoweredDeclarationOrigin
 import org.jetbrains.kotlin.ir.backend.js.lower.*
 import org.jetbrains.kotlin.ir.backend.js.lower.inline.RemoveInlineFunctionsWithReifiedTypeParametersLowering
-import org.jetbrains.kotlin.ir.backend.js.lower.inline.ReturnableBlockLowering
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents