JVM_IR: implement IrReturnableBlock codegen

(cherry picked from commit 530ad368fe)
2019-04-17 12:53:55 +02:00
parent 683f8417d3
commit 08da078842
4 changed files with 119 additions and 115 deletions
@@ -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) {
-                    putTypeArgument(0, expression.getTypeArgument(0))
-                }
+                copyTypeArgumentsFrom(expression)
                putValueArgument(0, irGet(sizeVar))
            })

            val lambda = invokable.asSingleArgumentLambda()
-            if (lambda == null) {
-                val invokableVar = irTemporary(invokable)
-                val invoke = invokable.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
-                fromZeroTo(sizeVar) { _, index ->
-                    +setItem(result, index, irCall(invoke).apply {
-                        dispatchReceiver = irGet(invokableVar)
-                        putValueArgument(0, irGet(index))
-                    })
+            val invoke = invokable.type.getClass()!!.functions.single { it.name == OperatorNameConventions.INVOKE }
+            val invokableVar = if (lambda == null) irTemporary(invokable) else null
+            +irWhile().apply {
+                condition = irCall(context.irBuiltIns.lessFunByOperandType[index.type.toKotlinType()]!!).apply {
+                    putValueArgument(0, irGet(index))
+                    putValueArgument(1, irGet(sizeVar))
                }
-            } else {
-                // Inline `invokable` by replacing the argument with `i` and `return x` with `result[i] = x; continue`.
-                fromZeroTo(sizeVar) { loop, index ->
-                    val body = lambda.body!!.transform(object : IrElementTransformerVoidWithContext() {
-                        override fun visitGetValue(expression: IrGetValue) =
-                            if (expression.symbol == lambda.valueParameters[0].symbol)
-                                IrGetValueImpl(expression.startOffset, expression.endOffset, index.symbol)
-                            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")
+                body = irBlock {
+                    val value =
+                        lambda?.inline(listOf(index)) ?: irCallOp(invoke.symbol, invoke.returnType, irGet(invokableVar!!), irGet(index))
+                    +irCall(result.type.getClass()!!.functions.single { it.name == OperatorNameConventions.SET }).apply {
+                        dispatchReceiver = irGet(result)
+                        putValueArgument(0, irGet(index))
+                        putValueArgument(1, value)
                    }
+                    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
+    }
 }
@@ -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]
 }

@@ -52,17 +52,22 @@ class TryInfo(val onExit: IrExpression) : ExpressionInfo() {
    val gaps = mutableListOf<Pair<Label, Label>>()
 }

-class BlockInfo private constructor(val parent: BlockInfo?) {
-    val variables = mutableListOf<VariableInfo>()
-    val infos = Stack<ExpressionInfo>()
+class ReturnableBlockInfo(
+    val returnLabel: Label,
+    val returnSymbol: IrSymbol,
+    val returnTemporary: Int? = null
+) : ExpressionInfo()

-    fun create() = BlockInfo(this).apply {
-        this@apply.infos.addAll(this@BlockInfo.infos)
-    }
+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

-    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)
-        writeParameterInLocalVariableTable(startLabel)
+        val endLabel = markNewLabel()
+        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()
-        // Force materialization to avoid reading from out-of-scope variables.
-        return super.visitBlock(expression, info).materialized.apply {
-            writeLocalVariablesInTable(info)
+        val info = BlockInfo(data)
+        return if (expression is IrReturnableBlock) {
+            val returnType = expression.asmType
+            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) {
-        val endLabel = markNewLabel()
+    private fun writeLocalVariablesInTable(info: BlockInfo, endLabel: Label) {
        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) {
-            generateGlobalReturnFlag(mv, (owner as IrFunction).name.asString())
+        if (target != null) {
+            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 {
-                it is TryInfo -> genFinallyBlock(it, null, endLabel, data)
-                it is LoopInfo && it.loop == loop -> return it
-            }
-            unwindBlockStack(endLabel, data, loop)
+            if (it is TryInfo)
+                genFinallyBlock(it, null, endLabel, data)
+            return if (stop(it)) it else unwindBlockStack(endLabel, data, stop)
        }
    }

    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 }
            }
        }
    }
@@ -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) {
-        val onStack = codegen.gen(argumentExpression, valueType, BlockInfo.create())
+    private fun putValueOnStack(argumentExpression: IrExpression, valueType: Type, paramIndex: Int, blockInfo: BlockInfo) {
+        val onStack = codegen.gen(argumentExpression, valueType, blockInfo)
        putArgumentOrCapturedToLocalVal(JvmKotlinType(onStack.type, onStack.kotlinType), onStack, -1, paramIndex, ValueKind.CAPTURED)
    }