JVM IR: replace function accesses in multi-file parts

Any access to a function from a multi-file part needs to be replaced with the access to the corresponding public method (if it exists) from the facade class. Note that this has no immediate effect because we use KotlinTypeMapper for mapping calls, and it understands that a call to a function from the part must actually be generated into a call to the function from the facade in the bytecode. This commit merely changes the IR to better reflect what's generated in the final bytecode, and to be able to use simplified IR-based method signature mapping instead of the legacy KotlinTypeMapper in the future.
2019-08-21 15:10:14 +02:00
parent 21b2487290
commit 2b4424b564
7 changed files with 87 additions and 8 deletions
@@ -68,6 +68,7 @@ class JvmBackendContext(

    internal val multifileFacadesToAdd = mutableMapOf<JvmClassName, MutableList<IrClass>>()
    internal val multifileFacadeForPart = mutableMapOf<IrClass, JvmClassName>()
+    internal val multifileFacadeMemberToPartMember = mutableMapOf<IrFunctionSymbol, IrFunctionSymbol>()

    override var inVerbosePhase: Boolean = false

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.backend.jvm.codegen
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.backend.common.ir.ir2string
+import org.jetbrains.kotlin.backend.jvm.lower.MultifileFacadeFileEntry
 import org.jetbrains.kotlin.codegen.BaseExpressionCodegen
 import org.jetbrains.kotlin.codegen.ClassBuilder
 import org.jetbrains.kotlin.codegen.OwnerKind
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrMemberAccessExpression
 import org.jetbrains.kotlin.ir.util.parentAsClass
+import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature
@@ -102,9 +104,10 @@ class IrSourceCompilerForInline(

    private fun getFunctionToInline(call: IrCall, jvmSignature: JvmMethodSignature, callDefault: Boolean): IrFunction {
        val callee = call.symbol.owner
+        val parent = callee.parentAsClass
        if (callDefault) {
            /*TODO: get rid of hack*/
-            return callee.parentAsClass.declarations.filterIsInstance<IrFunction>().single {
+            return parent.declarations.filterIsInstance<IrFunction>().single {
                it.descriptor.name.asString() == jvmSignature.asmMethod.name + JvmAbi.DEFAULT_PARAMS_IMPL_SUFFIX &&
                        codegen.context.methodSignatureMapper.mapSignatureSkipGeneric(callee).asmMethod.descriptor.startsWith(
                            jvmSignature.asmMethod.descriptor.substringBeforeLast(')')
@@ -112,6 +115,11 @@ class IrSourceCompilerForInline(
            }
        }

+        if (parent.fileParent.fileEntry is MultifileFacadeFileEntry) {
+            return (codegen.context.multifileFacadeMemberToPartMember[callee.symbol]
+                ?: error("Function from a multi-file facade without the link to the function in the part: ${callee.render()}")).owner
+        }
+
        return callee
    }

@@ -32,10 +32,14 @@ import org.jetbrains.kotlin.ir.builders.irReturn
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.IrFieldImpl
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
+import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrGetField
+import org.jetbrains.kotlin.ir.expressions.copyTypeArgumentsFrom
+import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
 import org.jetbrains.kotlin.ir.symbols.IrFieldSymbol
+import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.ir.util.transformFlat
@@ -55,9 +59,17 @@ internal val generateMultifileFacadesPhase = namedIrModulePhase(
            input: IrModuleFragment
        ): IrModuleFragment {
            val movedFields = mutableMapOf<IrFieldSymbol, IrFieldSymbol>()
-            input.files.addAll(generateMultifileFacades(input.descriptor, context, movedFields))
+            val functionDelegates = mutableMapOf<IrFunctionSymbol, IrFunctionSymbol>()
+
+            input.files.addAll(generateMultifileFacades(input.descriptor, context, movedFields, functionDelegates))
+
            UpdateFieldCallSites(movedFields).lower(input)
+            UpdateFunctionCallSites(functionDelegates).lower(input)
+
            context.multifileFacadesToAdd.clear()
+
+            functionDelegates.entries.associateTo(context.multifileFacadeMemberToPartMember) { (member, newMember) -> newMember to member }
+
            return input
        }
    }
@@ -86,7 +98,8 @@ internal class MultifileFacadeFileEntry(
 private fun generateMultifileFacades(
    module: ModuleDescriptor,
    context: JvmBackendContext,
-    movedFields: MutableMap<IrFieldSymbol, IrFieldSymbol>
+    movedFields: MutableMap<IrFieldSymbol, IrFieldSymbol>,
+    functionDelegates: MutableMap<IrFunctionSymbol, IrFunctionSymbol>
 ): List<IrFile> =
    context.multifileFacadesToAdd.map { (jvmClassName, partClasses) ->
        val fileEntry = MultifileFacadeFileEntry(jvmClassName, partClasses.map(IrClass::fileParent))
@@ -106,7 +119,12 @@ private fun generateMultifileFacades(
            moveFieldsOfConstProperties(partClass, facadeClass, movedFields)

            for (member in partClass.declarations) {
-                member.createMultifileDelegateIfNeeded(context, facadeClass)
+                if (member is IrFunction) {
+                    val newMember = member.createMultifileDelegateIfNeeded(context, facadeClass)
+                    if (newMember != null) {
+                        functionDelegates[member.symbol] = newMember.symbol
+                    }
+                }
            }
        }

@@ -135,12 +153,11 @@ private fun IrField.shouldMoveToFacade(): Boolean {
    return property != null && property.isConst && !Visibilities.isPrivate(visibility)
 }

-private fun IrDeclaration.createMultifileDelegateIfNeeded(context: JvmBackendContext, facadeClass: IrClass) {
-    if (this !is IrFunction ||
-        Visibilities.isPrivate(visibility) ||
+private fun IrFunction.createMultifileDelegateIfNeeded(context: JvmBackendContext, facadeClass: IrClass): IrFunction? {
+    if (Visibilities.isPrivate(visibility) ||
        name == InitializersLowering.clinitName ||
        origin == JvmLoweredDeclarationOrigin.SYNTHETIC_ACCESSOR
-    ) return
+    ) return null

    // TODO: perform copy of the signature only, without body
    val function = deepCopyWithSymbols(facadeClass)
@@ -158,6 +175,8 @@ private fun IrDeclaration.createMultifileDelegateIfNeeded(context: JvmBackendCon
    function.origin = JvmLoweredDeclarationOrigin.MULTIFILE_BRIDGE

    facadeClass.declarations.add(function)
+
+    return function
 }

 // This deep copy is needed while we still use KotlinTypeMapper to map signatures in method calls. Without it, KotlinTypeMapper takes
@@ -190,3 +209,24 @@ private class UpdateFieldCallSites(
        }
    }
 }
+
+private class UpdateFunctionCallSites(
+    private val functionDelegates: MutableMap<IrFunctionSymbol, IrFunctionSymbol>
+) : FileLoweringPass, IrElementTransformerVoid() {
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        val newFunction = functionDelegates[expression.symbol] ?: return super.visitCall(expression)
+        return expression.run {
+            IrCallImpl(startOffset, endOffset, type, newFunction).apply {
+                copyTypeArgumentsFrom(expression)
+                extensionReceiver = expression.extensionReceiver?.transform(this@UpdateFunctionCallSites, null)
+                for (i in 0 until valueArgumentsCount) {
+                    putValueArgument(i, expression.getValueArgument(i)?.transform(this@UpdateFunctionCallSites, null))
+                }
+            }
+        }
+    }
+}
@@ -0,0 +1,15 @@
+// TARGET_BACKEND: JVM
+// WITH_RUNTIME
+// FILE: 1.kt
+
+@file:JvmName("Facade")
+@file:JvmMultifileClass
+
+inline fun foo(o: String, k: String = "K", body: (String) -> String): String =
+    o + bar(body(k))
+
+fun bar(x: String): String = x
+
+// FILE: 2.kt
+
+fun box(): String = foo("O") { it }
@@ -1924,6 +1924,11 @@ public class BlackBoxInlineCodegenTestGenerated extends AbstractBlackBoxInlineCo
            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/multifileClasses"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
        }

+        @TestMetadata("defaultArguments.kt")
+        public void testDefaultArguments() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/multifileClasses/defaultArguments.kt");
+        }
+
        @TestMetadata("inlineFromOptimizedMultifileClass.kt")
        public void testInlineFromOptimizedMultifileClass() throws Exception {
            runTest("compiler/testData/codegen/boxInline/multifileClasses/inlineFromOptimizedMultifileClass.kt");
@@ -1924,6 +1924,11 @@ public class CompileKotlinAgainstInlineKotlinTestGenerated extends AbstractCompi
            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/multifileClasses"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM, true);
        }

+        @TestMetadata("defaultArguments.kt")
+        public void testDefaultArguments() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/multifileClasses/defaultArguments.kt");
+        }
+
        @TestMetadata("inlineFromOptimizedMultifileClass.kt")
        public void testInlineFromOptimizedMultifileClass() throws Exception {
            runTest("compiler/testData/codegen/boxInline/multifileClasses/inlineFromOptimizedMultifileClass.kt");
@@ -1924,6 +1924,11 @@ public class IrBlackBoxInlineCodegenTestGenerated extends AbstractIrBlackBoxInli
            KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxInline/multifileClasses"), Pattern.compile("^(.+)\\.kt$"), TargetBackend.JVM_IR, true);
        }

+        @TestMetadata("defaultArguments.kt")
+        public void testDefaultArguments() throws Exception {
+            runTest("compiler/testData/codegen/boxInline/multifileClasses/defaultArguments.kt");
+        }
+
        @TestMetadata("inlineFromOptimizedMultifileClass.kt")
        public void testInlineFromOptimizedMultifileClass() throws Exception {
            runTest("compiler/testData/codegen/boxInline/multifileClasses/inlineFromOptimizedMultifileClass.kt");