From 5aa2deb449cecfba6389f410e832d64cd2c340a5 Mon Sep 17 00:00:00 2001
From: Anton Bannykh <Anton.Bannykh@jetbrains.com>
Date: Thu, 19 Dec 2019 17:57:53 +0300
Subject: [PATCH] Make ReturnableBlockLowering a BodyLoweringPass

---
 .../common/lower/ReturnableBlockLowering.kt   | 27 +++++++++----------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt
index cf607aa7d19..b91966477bd 100644
--- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt
+++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt
@@ -5,22 +5,18 @@
 
 package org.jetbrains.kotlin.backend.common.lower
 
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
 import org.jetbrains.kotlin.backend.common.CommonBackendContext
-import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
-import org.jetbrains.kotlin.ir.builders.irBoolean
-import org.jetbrains.kotlin.ir.builders.irBreak
-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
-import org.jetbrains.kotlin.ir.expressions.IrReturn
-import org.jetbrains.kotlin.ir.expressions.IrReturnableBlock
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrSymbolOwner
+import org.jetbrains.kotlin.ir.expressions.*
 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.symbols.IrSymbol
 
 /**
  * Replaces returnable blocks and `return`'s with loops and `break`'s correspondingly.
@@ -69,13 +65,13 @@ import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
  * }
  *
  */
-class ReturnableBlockLowering(val context: CommonBackendContext) : FileLoweringPass {
-    override fun lower(irFile: IrFile) {
-        irFile.transform(ReturnableBlockTransformer(context), null)
+class ReturnableBlockLowering(val context: CommonBackendContext) : BodyLoweringPass {
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        container.transform(ReturnableBlockTransformer(context, (container as IrSymbolOwner).symbol), null)
     }
 }
 
-class ReturnableBlockTransformer(val context: CommonBackendContext) : IrElementTransformerVoidWithContext() {
+class ReturnableBlockTransformer(val context: CommonBackendContext, val containerSymbol: IrSymbol? = null) : IrElementTransformerVoidWithContext() {
     private var labelCnt = 0
     private val returnMap = mutableMapOf<IrReturnableBlockSymbol, (IrReturn) -> IrExpression>()
 
@@ -87,7 +83,8 @@ class ReturnableBlockTransformer(val context: CommonBackendContext) : IrElementT
     override fun visitContainerExpression(expression: IrContainerExpression): IrExpression {
         if (expression !is IrReturnableBlock) return super.visitContainerExpression(expression)
 
-        val builder = context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol)
+        val scopeSymbol = currentScope?.scope?.scopeOwnerSymbol ?: containerSymbol
+        val builder = context.createIrBuilder(scopeSymbol!!)
         val variable by lazy {
             builder.scope.createTemporaryVariableDeclaration(expression.type, "tmp\$ret\$${labelCnt++}", true)
         }