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.

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendContext.kt

@@ -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
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/IrSourceCompilerForInline.kt

@@ -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
     }
 

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/GenerateMultifileFacades.kt

@@ -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) {
+private fun IrFunction.createMultifileDelegateIfNeeded(context: JvmBackendContext, facadeClass: IrClass): IrFunction? {
-    if (this !is IrFunction ||
+    if (Visibilities.isPrivate(visibility) ||
-        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))
+                }
+            }
+        }
+    }
+}

compiler/testData/codegen/boxInline/multifileClasses/defaultArguments.kt

@@ -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 }

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxInlineCodegenTestGenerated.java

@@ -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");

compiler/tests/org/jetbrains/kotlin/codegen/CompileKotlinAgainstInlineKotlinTestGenerated.java

@@ -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");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxInlineCodegenTestGenerated.java

@@ -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");