IR: fixes for IR by-file copying

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/CommonBackendContext.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.ir.builders.irString
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.IrFileSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 
 interface LoggingContext {
@@ -40,6 +41,7 @@ interface CommonBackendContext : BackendContext, LoggingContext {
 
     // Adjust internal structures after a deep copy of some declarations.
     fun handleDeepCopy(
+        fileSymbolMap: MutableMap<IrFileSymbol, IrFileSymbol>,
         classSymbolMap: MutableMap<IrClassSymbol, IrClassSymbol>,
         functionSymbolMap: MutableMap<IrSimpleFunctionSymbol, IrSimpleFunctionSymbol>
     ) {}

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/phaser/PhaseBuilders.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.backend.common.FileLoweringPass
 import org.jetbrains.kotlin.backend.common.lower
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
@@ -18,6 +19,7 @@ import org.jetbrains.kotlin.ir.expressions.IrCall
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.impl.IrCallImpl
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.symbols.IrFileSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.util.DeepCopySymbolRemapper
 import org.jetbrains.kotlin.ir.util.copyTypeAndValueArgumentsFrom
@@ -166,12 +168,13 @@ private class PerformByIrFilePhase<Context : CommonBackendContext>(
         // We can only report one exception through ISE
         val thrownFromThread = AtomicReference<Pair<Throwable, IrFile>?>(null)
 
+        val remappedFiles = mutableMapOf<IrFileSymbol, IrFileSymbol>()
         val remappedFunctions = mutableMapOf<IrSimpleFunctionSymbol, IrSimpleFunctionSymbol>()
         val remappedClasses = mutableMapOf<IrClassSymbol, IrClassSymbol>()
 
         // Each thread needs its own copy of phaserState.alreadyDone
         val filesAndStates = input.files.map {
-            it.copySavingMappings(remappedFunctions, remappedClasses) to phaserState.copyOf()
+            it.copySavingMappings(remappedFiles, remappedFunctions, remappedClasses) to phaserState.copyOf()
         }
 
         val executor = Executors.newFixedThreadPool(nThreads)
@@ -201,8 +204,8 @@ private class PerformByIrFilePhase<Context : CommonBackendContext>(
         input.files.addAll(filesAndStates.map { (irFile, _) -> irFile }.toMutableList())
 
         adjustDefaultArgumentStubs(context, remappedFunctions)
+        context.handleDeepCopy(remappedFiles, remappedClasses, remappedFunctions)
         input.transformChildrenVoid(CrossFileCallAdjuster(remappedFunctions))
-        context.handleDeepCopy(remappedClasses, remappedFunctions)
 
         // TODO: no guarantee that module identity is preserved by `lower`
         return input
@@ -283,6 +286,7 @@ fun <Context : CommonBackendContext, OldData, NewData> transform(op: (OldData) -
 // We need to remap inline function calls after lowering files
 
 fun IrFile.copySavingMappings(
+    remappedFiles: MutableMap<IrFileSymbol, IrFileSymbol>,
     remappedFunctions: MutableMap<IrSimpleFunctionSymbol, IrSimpleFunctionSymbol>,
     remappedClasses: MutableMap<IrClassSymbol, IrClassSymbol>,
 ): IrFile {
@@ -297,6 +301,8 @@ fun IrFile.copySavingMappings(
         remappedClasses[klass] = symbolRemapper.getReferencedClass(klass)
     }
 
+    remappedFiles[symbol] = newIrFile.symbol
+
     return newIrFile
 }
 
@@ -331,6 +337,12 @@ fun adjustDefaultArgumentStubs(
 private class CrossFileCallAdjuster(
     val remappedFunctions: Map<IrSimpleFunctionSymbol, IrSimpleFunctionSymbol>
 ) : IrElementTransformerVoid() {
+
+    override fun visitSimpleFunction(declaration: IrSimpleFunction): IrStatement {
+        declaration.overriddenSymbols = declaration.overriddenSymbols.map { remappedFunctions[it] ?: it }
+        return super.visitSimpleFunction(declaration)
+    }
+
     override fun visitCall(expression: IrCall): IrExpression {
         expression.transformChildrenVoid(this)
         return remappedFunctions[expression.symbol]?.let { newSymbol ->

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

@@ -155,9 +155,18 @@ class JvmBackendContext(
         }
 
     override fun handleDeepCopy(
+        fileSymbolMap: MutableMap<IrFileSymbol, IrFileSymbol>,
         classSymbolMap: MutableMap<IrClassSymbol, IrClassSymbol>,
         functionSymbolMap: MutableMap<IrSimpleFunctionSymbol, IrSimpleFunctionSymbol>
     ) {
+        for ((oldFileSymbol, newFileSymbol) in fileSymbolMap) {
+            val psiFileEntry = psiSourceManager.getFileEntry(oldFileSymbol.owner) as? PsiSourceManager.PsiFileEntry ?: continue
+            psiSourceManager.putFileEntry(
+                newFileSymbol.owner,
+                psiFileEntry
+            )
+        }
+
         val oldClassesWithNameOverride = classNameOverride.keys.toList()
         for (klass in oldClassesWithNameOverride) {
             classSymbolMap[klass.symbol]?.let { newSymbol ->
@@ -170,7 +179,29 @@ class JvmBackendContext(
             multifileFacade.setValue(newPartClasses.toMutableList())
         }
 
-        super.handleDeepCopy(classSymbolMap, functionSymbolMap)
+        for ((staticReplacement, original) in inlineClassReplacements.originalFunctionForStaticReplacement) {
+            val newOriginal = functionSymbolMap[original.symbol]?.owner ?: continue
+            val newStaticReplacement = inlineClassReplacements.getReplacementFunction(newOriginal) ?: continue
+            staticReplacement as IrSimpleFunction
+            functionSymbolMap[staticReplacement.symbol] = newStaticReplacement.symbol
+        }
+
+        for ((original, suspendView) in suspendFunctionOriginalToView) {
+            val newOriginal = functionSymbolMap[original.symbol]?.owner ?: continue
+            val newSuspendView = suspendFunctionOriginalToView[newOriginal] ?: continue
+            suspendView as? IrSimpleFunction ?: continue
+            newSuspendView as? IrSimpleFunction ?: continue
+            functionSymbolMap[suspendView.symbol] = newSuspendView.symbol
+        }
+
+        for ((nonStaticDefaultSymbol, staticDefault) in staticDefaultStubs) {
+            val staticDefaultSymbol = staticDefault.symbol
+            val newNonStaticDefaultSymbol = functionSymbolMap[nonStaticDefaultSymbol] ?: continue
+            val newStaticDefaultSymbol = staticDefaultStubs[newNonStaticDefaultSymbol]?.symbol ?: continue
+            functionSymbolMap[staticDefaultSymbol] = newStaticDefaultSymbol
+        }
+
+        super.handleDeepCopy(fileSymbolMap, classSymbolMap, functionSymbolMap)
     }
 
     inner class JvmIr(