IR: unify 3 copies of function body remapping

Also, 1. remove some redundant copies; 2. fix remapping of non-local returns in lambdas if the body is moved after LocalDeclarationsLowering (the lambda is no longer inside the body, but must still be visited)
2019-12-30 11:46:03 +01:00
parent 4d1677be06
commit 99eab5a058
15 changed files with 120 additions and 131 deletions
@@ -6,13 +6,14 @@
 package org.jetbrains.kotlin.backend.common.ir

 import org.jetbrains.kotlin.backend.common.lower.VariableRemapper
-import org.jetbrains.kotlin.ir.declarations.IrFunction
-import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
-import org.jetbrains.kotlin.ir.declarations.IrValueDeclaration
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.IrReturnImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrReturnableBlockImpl
+import org.jetbrains.kotlin.ir.symbols.IrReturnTargetSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrReturnableBlockSymbolImpl
+import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.util.isVararg
 import org.jetbrains.kotlin.ir.util.statements

@@ -36,21 +37,56 @@ fun IrExpression.asSimpleLambda(): IrSimpleFunction? {
    return function
 }

+private fun createParameterMapping(source: IrFunction, target: IrFunction): Map<IrValueParameter, IrValueParameter> {
+    val sourceParameters = listOfNotNull(source.dispatchReceiverParameter, source.extensionReceiverParameter) + source.valueParameters
+    val targetParameters = listOfNotNull(target.dispatchReceiverParameter, target.extensionReceiverParameter) + target.valueParameters
+    assert(sourceParameters.size == targetParameters.size)
+    return sourceParameters.zip(targetParameters).toMap()
+}
+
+fun IrFunction.copyBodyTo(target: IrFunction): IrBody? =
+    copyBodyTo(target, createParameterMapping(this, target))
+
+fun IrFunction.copyBodyTo(target: IrFunction, arguments: Map<IrValueParameter, IrValueDeclaration>): IrBody? =
+    body?.deepCopyWithSymbols(target)?.move(this, target, target.symbol, arguments)
+
+fun IrFunction.moveBodyTo(target: IrFunction): IrBody? =
+    moveBodyTo(target, createParameterMapping(this, target))
+
+fun IrFunction.moveBodyTo(target: IrFunction, arguments: Map<IrValueParameter, IrValueDeclaration>): IrBody? =
+    body?.move(this, target, target.symbol, arguments)
+
+private fun IrBody.move(
+    source: IrFunction,
+    target: IrDeclarationParent,
+    targetSymbol: IrReturnTargetSymbol,
+    arguments: Map<IrValueParameter, IrValueDeclaration>
+): IrBody = transform(object : VariableRemapper(arguments) {
+    override fun visitReturn(expression: IrReturn): IrExpression = super.visitReturn(
+        if (expression.returnTargetSymbol == source.symbol)
+            IrReturnImpl(expression.startOffset, expression.endOffset, expression.type, targetSymbol, expression.value)
+        else
+            expression
+    )
+
+    override fun visitBlock(expression: IrBlock): IrExpression {
+        // Might be an inline lambda argument; if the function has already been moved out, visit it explicitly.
+        if (expression.origin == IrStatementOrigin.LAMBDA || expression.origin == IrStatementOrigin.ANONYMOUS_FUNCTION)
+            if (expression.statements[0] !is IrFunction && expression.statements[1] is IrFunctionReference)
+                (expression.statements[1] as IrFunctionReference).symbol.owner.transformChildrenVoid()
+        return super.visitBlock(expression)
+    }
+
+    override fun visitDeclaration(declaration: IrDeclaration): IrStatement {
+        if (declaration.parent == source)
+            declaration.parent = target
+        return super.visitDeclaration(declaration)
+    }
+}, null)
+
 // TODO use a generic inliner (e.g. JS/Native's FunctionInlining.Inliner)
 // Inline simple function calls without type parameters, default parameters, or varargs.
-fun IrFunction.inline(arguments: List<IrValueDeclaration> = listOf()): IrReturnableBlock {
-    require(body != null)
-    val argumentMap = valueParameters.zip(arguments).toMap()
-    val blockSymbol = IrReturnableBlockSymbolImpl(descriptor)
-    val block = IrReturnableBlockImpl(startOffset, endOffset, returnType, blockSymbol, null, symbol)
-    val remapper = object : VariableRemapper(argumentMap) {
-        override fun visitReturn(expression: IrReturn): IrExpression = super.visitReturn(
-            if (expression.returnTargetSymbol == symbol)
-                IrReturnImpl(expression.startOffset, expression.endOffset, expression.type, blockSymbol, expression.value)
-            else
-                expression
-        )
+fun IrFunction.inline(target: IrDeclarationParent, arguments: List<IrValueDeclaration> = listOf()): IrReturnableBlock =
+    IrReturnableBlockImpl(startOffset, endOffset, returnType, IrReturnableBlockSymbolImpl(descriptor), null, symbol).apply {
+        statements += body!!.move(this@inline, target, symbol, valueParameters.zip(arguments).toMap()).statements
    }
-    body!!.statements.mapTo(block.statements) { it.transform(remapper, null) }
-    return block
-}
@@ -598,46 +598,5 @@ fun createStaticFunctionWithReceivers(
    }
 }

-fun copyBodyToStatic(oldFunction: IrFunction, staticFunction: IrFunction) {
-    val mapping: Map<IrValueParameter, IrValueParameter> =
-        (listOfNotNull(oldFunction.dispatchReceiverParameter, oldFunction.extensionReceiverParameter) + oldFunction.valueParameters)
-            .zip(staticFunction.valueParameters).toMap()
-    copyBodyWithParametersMapping(staticFunction, oldFunction, mapping)
-}
-
-fun copyBodyWithParametersMapping(
-    newFunction: IrFunction,
-    oldFunction: IrFunction,
-    mapping: Map<IrValueParameter, IrValueParameter>
-) {
-    newFunction.body = oldFunction.body?.deepCopyWithSymbols(oldFunction)
-        ?.transform(
-            object : IrElementTransformerVoid() {
-                // Remap return targets to the static method so they do not appear to be
-                // non-local returns.
-                override fun visitReturn(expression: IrReturn): IrExpression {
-                    expression.transformChildrenVoid(this);
-                    return if (expression.returnTargetSymbol == oldFunction.symbol) {
-                        IrReturnImpl(
-                            expression.startOffset,
-                            expression.endOffset,
-                            expression.type,
-                            newFunction.symbol,
-                            expression.value
-                        )
-                    } else expression
-                }
-
-                // Remap argument values.
-                override fun visitGetValue(expression: IrGetValue): IrExpression =
-                    mapping[expression.symbol.owner]?.let {
-                        IrGetValueImpl(expression.startOffset, expression.endOffset, it.type, it.symbol, expression.origin)
-                    } ?: expression
-
-            }, null
-        )
-        ?.patchDeclarationParents(newFunction)
-}
-
 val IrSymbol.isSuspend: Boolean
    get() = this is IrSimpleFunctionSymbol && owner.isSuspend
@@ -23,7 +23,6 @@ import org.jetbrains.kotlin.ir.types.getClass
 import org.jetbrains.kotlin.ir.util.constructedClass
 import org.jetbrains.kotlin.ir.util.constructors
 import org.jetbrains.kotlin.ir.util.functions
-import org.jetbrains.kotlin.ir.util.patchDeclarationParents
 import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.util.OperatorNameConventions

@@ -84,7 +83,7 @@ class ArrayConstructorLowering(val context: CommonBackendContext) : IrElementTra
                }
                body = irBlock {
                    val tempIndex = irTemporary(irGet(index))
-                    val value = lambda?.inline(listOf(tempIndex)) ?: irCallOp(
+                    val value = lambda?.inline(parent, listOf(tempIndex)) ?: irCallOp(
                        invoke.symbol,
                        invoke.returnType,
                        irGet(invokableVar!!),
@@ -100,10 +99,6 @@ class ArrayConstructorLowering(val context: CommonBackendContext) : IrElementTra
                }
            }
            +irGet(result)
-        }.also {
-            // Some parents of local declarations are not updated during ad-hoc inlining
-            // TODO: Remove when generic inliner is used
-            it.patchDeclarationParents(scope.getLocalDeclarationParent())
        }
    }
 }