[JS IR] Generate correct source locations in sourcemap for inline declarations

* Keep returnable blocks.
* Add a new lowering which simplifies returnable blocks by introducing
  temporary variable for result and changing returnable block's type to Unit.
* Use information from returnable blocks in codegen to generate
  the right source locations in sourcemap.
* Support in namer (LocalNameGenerator).
* Fix some lowerings to work correctly with returnable blocks.

#KT-46551 In Progress

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSymbol
 import org.jetbrains.kotlin.ir.transformStatement
 
+// TODO migrate other usages and move this file to backend.jvm
 /**
  * Replaces returnable blocks and `return`'s with loops and `break`'s correspondingly.
  *
@@ -149,4 +150,4 @@ class ReturnableBlockTransformer(val context: CommonBackendContext, val containe
             }
         }
     }
-}
+}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -348,9 +348,9 @@ private val callableReferenceLowering = makeBodyLoweringPhase(
 )
 
 private val returnableBlockLoweringPhase = makeBodyLoweringPhase(
-    ::ReturnableBlockLowering,
-    name = "ReturnableBlockLowering",
-    description = "Replace returnable block with do-while loop",
+    ::JsReturnableBlockLowering,
+    name = "JsReturnableBlockLowering",
+    description = "Introduce temporary variable for result and change returnable block's type to Unit",
     prerequisite = setOf(functionInliningPhase)
 )
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/JsReturnableBlockLowering.kt

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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
+
+import org.jetbrains.kotlin.backend.common.BodyLoweringPass
+import org.jetbrains.kotlin.backend.common.CommonBackendContext
+import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.lower.at
+import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
+import org.jetbrains.kotlin.backend.common.lower.irComposite
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.*
+import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.expressions.impl.IrGetObjectValueImpl
+import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
+
+/**
+ * Wraps returnable blocks with returns to composite and replaces returns with assignment to temporary variable + `return Unit`,
+ * also, it changes type of returnable block to Unit.
+ *
+ * ```
+ * returnable_block {
+ *   ...
+ *   return@returnable_block e
+ *   ...
+ * }: T
+ * ```
+ *
+ * is transformed into
+ *
+ * ```
+ * composite {
+ *   val result
+ *   returnable_block {
+ *     ...
+ *     result = e
+ *     return@returnable_block Unit
+ *     ...
+ *   }: Unit
+ *   result
+ * }: T
+ * ```
+ */
+class JsReturnableBlockLowering(val context: CommonBackendContext) : BodyLoweringPass {
+    override fun lower(irBody: IrBody, container: IrDeclaration) {
+        container.transform(JsReturnableBlockTransformer(context), null)
+    }
+}
+
+class JsReturnableBlockTransformer(val context: CommonBackendContext) : IrElementTransformerVoidWithContext() {
+    private var labelCnt = 0
+    private var map = mutableMapOf<IrReturnableBlockSymbol, IrVariable>()
+
+    private val unitType = context.irBuiltIns.unitType
+    private val unitValue get() = IrGetObjectValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, unitType, context.irBuiltIns.unitClass)
+
+    override fun visitBlock(expression: IrBlock): IrExpression {
+        if (expression !is IrReturnableBlock) return super.visitBlock(expression)
+
+        expression.transformChildrenVoid()
+
+        val variable = map.remove(expression.symbol) ?: return expression
+
+        val builder = context.createIrBuilder(expression.symbol)
+
+        expression.type = context.irBuiltIns.unitType
+        return builder.irComposite(expression, expression.origin, variable.type) {
+            +variable
+            +expression
+            +irGet(variable)
+        }
+    }
+
+    override fun visitReturn(expression: IrReturn): IrExpression {
+        expression.transformChildrenVoid()
+
+        val targetSymbol = expression.returnTargetSymbol
+        if (targetSymbol !is IrReturnableBlockSymbol) return expression
+
+        val variable = map.getOrPut(targetSymbol) {
+            currentScope!!.scope.createTmpVariable(targetSymbol.owner.type, "tmp\$ret\$${labelCnt++}", true)
+        }
+
+        val builder = context.createIrBuilder(targetSymbol)
+        return builder.at(UNDEFINED_OFFSET, UNDEFINED_OFFSET).irComposite {
+            +at(expression).irSet(variable.symbol, expression.value)
+            +at(UNDEFINED_OFFSET, UNDEFINED_OFFSET).irReturn(unitValue)
+        }
+    }
+}

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/SecondaryCtorLowering.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2018 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -10,8 +10,8 @@ import org.jetbrains.kotlin.backend.common.DeclarationTransformer
 import org.jetbrains.kotlin.backend.common.getOrPut
 import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom
-import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.symbols.IrConstructorSymbol
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
@@ -130,25 +131,30 @@ class SecondaryConstructorLowering(val context: JsIrBackendContext) : Declaratio
         delegate.body = context.irFactory.createBlockBody(UNDEFINED_OFFSET, UNDEFINED_OFFSET) {
             statements += (constructorBody.deepCopyWithSymbols(delegate) as IrStatementContainer).statements
             statements += JsIrBuilder.buildReturn(delegate.symbol, JsIrBuilder.buildGetValue(thisParam.symbol), context.irBuiltIns.nothingType)
-            transformChildrenVoid(ThisUsageReplaceTransformer(delegate.symbol, oldValueParameters.zip(delegate.valueParameters).toMap()))
+            transformChildrenVoid(ThisUsageReplaceTransformer(constructor.symbol, delegate.symbol, oldValueParameters.zip(delegate.valueParameters).toMap()))
         }
     }
 
 
     private class ThisUsageReplaceTransformer(
+        val constructor: IrConstructorSymbol,
         val function: IrFunctionSymbol,
         val symbolMapping: Map<IrValueParameter, IrValueParameter>
     ) : IrElementTransformerVoid() {
 
         val newThisSymbol = symbolMapping.values.last().symbol
 
-        override fun visitReturn(expression: IrReturn): IrExpression = IrReturnImpl(
-            expression.startOffset,
-            expression.endOffset,
-            expression.type,
-            function,
-            IrGetValueImpl(expression.startOffset, expression.endOffset, newThisSymbol.owner.type, newThisSymbol)
-        )
+        override fun visitReturn(expression: IrReturn): IrExpression =
+            if (expression.returnTargetSymbol != constructor)
+                expression
+            else
+                IrReturnImpl(
+                    expression.startOffset,
+                    expression.endOffset,
+                    expression.type,
+                    function,
+                    IrGetValueImpl(expression.startOffset, expression.endOffset, newThisSymbol.owner.type, newThisSymbol)
+                )
 
         override fun visitGetValue(expression: IrGetValue) = symbolMapping[expression.symbol.owner]?.let {
             expression.run { IrGetValueImpl(startOffset, endOffset, type, it.symbol, origin) }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/cleanup/CleanupLowering.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -36,7 +36,12 @@ private class BlockRemover : IrElementVisitorVoid {
 
     private fun process(container: IrStatementContainer) {
         container.statements.transformFlat { statement ->
-            (statement as? IrStatementContainer)?.statements
+            when (statement) {
+                // returnable blocks required for sourcemaps generation, so keep them
+                is IrReturnableBlock -> null
+                is IrStatementContainer -> statement.statements
+                else -> null
+            }
         }
     }
 
@@ -110,4 +115,4 @@ private class CodeCleaner : IrElementVisitorVoid {
         super.visitContainerExpression(expression)
         expression.cleanUpStatements()
     }
-}
+}

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/transformers/irToJs/IrElementToJsStatementTransformer.kt

@@ -11,8 +11,10 @@ import org.jetbrains.kotlin.ir.backend.js.utils.emptyScope
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.types.isAny
 import org.jetbrains.kotlin.ir.util.constructedClassType
+import org.jetbrains.kotlin.ir.util.file
 import org.jetbrains.kotlin.js.backend.ast.*
 
 @Suppress("PARAMETER_NAME_CHANGED_ON_OVERRIDE")
@@ -28,8 +30,19 @@ class IrElementToJsStatementTransformer : BaseIrElementToJsNodeTransformer<JsSta
         return JsBlock(body.statements.map { it.accept(this, context) })
     }
 
-    override fun visitBlock(expression: IrBlock, context: JsGenerationContext): JsBlock {
-        return JsBlock(expression.statements.map { it.accept(this, context) })
+    override fun visitBlock(expression: IrBlock, context: JsGenerationContext): JsStatement {
+        val newContext = (expression as? IrReturnableBlock)?.inlineFunctionSymbol?.let {
+            context.newFile(it.owner.file, context.currentFunction, context.localNames)
+        } ?: context
+
+        val block = JsBlock(expression.statements.map { it.accept(this, newContext) })
+
+        if (expression is IrReturnableBlock) {
+            val label = context.getNameForReturnableBlock(expression)
+            if (label != null) return JsLabel(label, block)
+        }
+
+        return block
     }
 
     override fun visitComposite(expression: IrComposite, context: JsGenerationContext): JsStatement {
@@ -72,7 +85,16 @@ class IrElementToJsStatementTransformer : BaseIrElementToJsNodeTransformer<JsSta
     }
 
     override fun visitReturn(expression: IrReturn, context: JsGenerationContext): JsStatement {
-        return expression.value.maybeOptimizeIntoSwitch(context) { JsReturn(it) }.withSource(expression, context)
+        val targetSymbol = expression.returnTargetSymbol
+        val lastStatementTransformer: (JsExpression) -> JsStatement =
+            if (targetSymbol is IrReturnableBlockSymbol) {
+                // TODO assert that value is Unit?
+                { JsBreak(context.getNameForReturnableBlock(targetSymbol.owner)!!.makeRef()) }
+            } else {
+                { JsReturn(it) }
+            }
+
+        return expression.value.maybeOptimizeIntoSwitch(context, lastStatementTransformer).withSource(expression, context)
     }
 
     override fun visitThrow(expression: IrThrow, context: JsGenerationContext): JsStatement {

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/utils/JsGenerationContext.kt

@@ -71,6 +71,11 @@ class JsGenerationContext(
         return JsName(name)
     }
 
+    fun getNameForReturnableBlock(block: IrReturnableBlock): JsName? {
+        val name = localNames!!.localReturnableBlockNames.names[block] ?: return null
+        return JsName(name)
+    }
+
     private fun isCoroutineDoResume(): Boolean {
         val overriddenSymbols = (currentFunction as? IrSimpleFunction)?.overriddenSymbols ?: return false
         return overriddenSymbols.any {

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/utils/NameTables.kt

@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
  * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
  */
 
@@ -10,10 +10,8 @@ import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.lower.serialization.ir.JsManglerIr
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrBreak
-import org.jetbrains.kotlin.ir.expressions.IrExpression
-import org.jetbrains.kotlin.ir.expressions.IrLoop
-import org.jetbrains.kotlin.ir.expressions.IrWhen
+import org.jetbrains.kotlin.ir.expressions.*
+import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.types.isUnit
 import org.jetbrains.kotlin.ir.util.file
 import org.jetbrains.kotlin.ir.util.fqNameWhenAvailable
@@ -304,6 +302,7 @@ class NameTables(
 class LocalNameGenerator(parentScope: NameScope) : IrElementVisitorVoid {
     val variableNames = NameTable<IrDeclarationWithName>(parentScope)
     val localLoopNames = NameTable<IrLoop>()
+    val localReturnableBlockNames = NameTable<IrReturnableBlock>()
 
     private val breakableDeque: Deque<IrExpression> = LinkedList()
 
@@ -327,6 +326,15 @@ class LocalNameGenerator(parentScope: NameScope) : IrElementVisitorVoid {
         super.visitBreak(jump)
     }
 
+    override fun visitReturn(expression: IrReturn) {
+        val targetSymbol = expression.returnTargetSymbol
+        if (targetSymbol is IrReturnableBlockSymbol) {
+            persistReturnableBlockName(SYNTHETIC_BLOCK_LABEL, targetSymbol.owner)
+        }
+
+        super.visitReturn(expression)
+    }
+
     override fun visitWhen(expression: IrWhen) {
         breakableDeque.push(expression)
 
@@ -352,6 +360,10 @@ class LocalNameGenerator(parentScope: NameScope) : IrElementVisitorVoid {
     private fun persistLoopName(label: String, loop: IrLoop) {
         localLoopNames.declareFreshName(loop, label)
     }
+
+    private fun persistReturnableBlockName(label: String, loop: IrReturnableBlock) {
+        localReturnableBlockNames.declareFreshName(loop, label)
+    }
 }
 
 
@@ -373,3 +385,4 @@ fun sanitizeName(name: String): String {
 }
 
 private const val SYNTHETIC_LOOP_LABEL = "\$l\$break"
+private const val SYNTHETIC_BLOCK_LABEL = "\$l\$block"