[IR] Add module property to IrFile

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt

@@ -55,10 +55,11 @@ class Fir2IrConverter(
         processClassMembers(regularClass, irClass)
     }
 
-    fun registerFileAndClasses(file: FirFile): IrFile {
+    fun registerFileAndClasses(file: FirFile, moduleFragment: IrModuleFragment) {
         val irFile = IrFileImpl(
             PsiIrFileEntry(file.psi as KtFile),
-            moduleDescriptor.getPackage(file.packageFqName).fragments.first()
+            moduleDescriptor.getPackage(file.packageFqName).fragments.first(),
+            moduleFragment
         )
         declarationStorage.registerFile(file, irFile)
         file.declarations.forEach {
@@ -66,7 +67,7 @@ class Fir2IrConverter(
                 registerClassAndNestedClasses(it, irFile)
             }
         }
-        return irFile
+        moduleFragment.files += irFile
     }
 
     fun processClassHeaders(file: FirFile) {
@@ -279,12 +280,11 @@ class Fir2IrConverter(
             components.visibilityConverter = visibilityConverter
             components.builtIns = builtIns
             components.annotationGenerator = annotationGenerator
-            val irFiles = mutableListOf<IrFile>()
 
+            val irModuleFragment = IrModuleFragmentImpl(moduleDescriptor, irBuiltIns)
             for (firFile in firFiles) {
-                irFiles += converter.registerFileAndClasses(firFile)
+                converter.registerFileAndClasses(firFile, irModuleFragment)
             }
-            val irModuleFragment = IrModuleFragmentImpl(moduleDescriptor, irBuiltIns, irFiles)
             val irProviders =
                 generateTypicalIrProviderList(irModuleFragment.descriptor, irBuiltIns, symbolTable, extensions = generatorExtensions)
             val externalDependenciesGenerator = ExternalDependenciesGenerator(

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsIrBackendContext.kt

@@ -100,7 +100,7 @@ class JsIrBackendContext(
 
             override fun getLineNumber(offset: Int) = UNDEFINED_OFFSET
             override fun getColumnNumber(offset: Int) = UNDEFINED_OFFSET
-        }, internalPackageFragmentDescriptor).also {
+        }, internalPackageFragmentDescriptor, module).also {
             module.files += it
         }
     }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/MoveBodilessDeclarationsToSeparatePlace.kt

@@ -62,7 +62,7 @@ class MoveBodilessDeclarationsToSeparatePlaceLowering(private val context: JsIrB
 
         val externalPackageFragment by lazy {
             context.externalPackageFragment.getOrPut(irFile.symbol) {
-                IrFileImpl(fileEntry = irFile.fileEntry, fqName = irFile.fqName, symbol = IrFileSymbolImpl()).also {
+                IrFileImpl(fileEntry = irFile.fileEntry, fqName = irFile.fqName, symbol = IrFileSymbolImpl(), module = irFile.module).also {
                     it.annotations += irFile.annotations
                 }
             }

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

@@ -21,7 +21,6 @@ import org.jetbrains.kotlin.backend.jvm.ir.getKtFile
 import org.jetbrains.kotlin.config.JvmAnalysisFlags
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.descriptors.Modality
-import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.impl.EmptyPackageFragmentDescriptor
 import org.jetbrains.kotlin.ir.*
 import org.jetbrains.kotlin.ir.builders.*
@@ -57,7 +56,7 @@ internal val generateMultifileFacadesPhase = makeCustomPhase<JvmBackendContext,
         // we construct an inheritance chain such that all part members are present as fake overrides in the facade.
         val shouldGeneratePartHierarchy = context.state.languageVersionSettings.getFlag(JvmAnalysisFlags.inheritMultifileParts)
         input.files.addAll(
-            generateMultifileFacades(input.descriptor, context, shouldGeneratePartHierarchy, functionDelegates)
+            generateMultifileFacades(input, context, shouldGeneratePartHierarchy, functionDelegates)
         )
 
         UpdateFunctionCallSites(functionDelegates).lower(input)
@@ -90,7 +89,7 @@ class MultifileFacadeFileEntry(
 }
 
 private fun generateMultifileFacades(
-    module: ModuleDescriptor,
+    module: IrModuleFragment,
     context: JvmBackendContext,
     shouldGeneratePartHierarchy: Boolean,
     functionDelegates: MutableMap<IrSimpleFunction, IrSimpleFunction>
@@ -108,7 +107,7 @@ private fun generateMultifileFacades(
         }
 
         val fileEntry = MultifileFacadeFileEntry(jvmClassName, partClasses.map(IrClass::fileParent))
-        val file = IrFileImpl(fileEntry, EmptyPackageFragmentDescriptor(module, kotlinPackageFqName))
+        val file = IrFileImpl(fileEntry, EmptyPackageFragmentDescriptor(module.descriptor, kotlinPackageFqName), module)
 
         context.log {
             "Multifile facade $jvmClassName:\n  ${partClasses.joinToString("\n  ") { it.fqNameWhenAvailable!!.asString() }}\n"

compiler/ir/backend.wasm/src/org/jetbrains/kotlin/backend/wasm/WasmBackendContext.kt

@@ -76,7 +76,7 @@ class WasmBackendContext(
 
             override fun getLineNumber(offset: Int) = UNDEFINED_OFFSET
             override fun getColumnNumber(offset: Int) = UNDEFINED_OFFSET
-        }, internalPackageFragmentDescriptor).also {
+        }, internalPackageFragmentDescriptor, irModuleFragment).also {
             irModuleFragment.files += it
         }
     }

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ModuleGenerator.kt

@@ -47,7 +47,7 @@ class ModuleGenerator(
         IrModuleFragmentImpl(context.moduleDescriptor, context.irBuiltIns).also { irModule ->
             val irDeclarationGenerator = DeclarationGenerator(context)
             ktFiles.toSet().mapTo(irModule.files) { ktFile ->
-                generateSingleFile(irDeclarationGenerator, ktFile)
+                generateSingleFile(irDeclarationGenerator, ktFile, irModule)
             }
         }
 
@@ -69,8 +69,8 @@ class ModuleGenerator(
             .generateUnboundSymbolsAsDependencies()
     }
 
-    private fun generateSingleFile(irDeclarationGenerator: DeclarationGenerator, ktFile: KtFile): IrFileImpl {
+    private fun generateSingleFile(irDeclarationGenerator: DeclarationGenerator, ktFile: KtFile, module: IrModuleFragment): IrFileImpl {
-        val irFile = createEmptyIrFile(ktFile)
+        val irFile = createEmptyIrFile(ktFile, module)
 
         for (ktAnnotationEntry in ktFile.annotationEntries) {
             val annotationDescriptor = getOrFail(BindingContext.ANNOTATION, ktAnnotationEntry)
@@ -101,10 +101,10 @@ class ModuleGenerator(
         return irFile
     }
 
-    private fun createEmptyIrFile(ktFile: KtFile): IrFileImpl {
+    private fun createEmptyIrFile(ktFile: KtFile, module: IrModuleFragment): IrFileImpl {
         val fileEntry = PsiIrFileEntry(ktFile)
         val packageFragmentDescriptor = context.moduleDescriptor.findPackageFragmentForFile(ktFile)!!
-        return IrFileImpl(fileEntry, packageFragmentDescriptor).apply {
+        return IrFileImpl(fileEntry, packageFragmentDescriptor, module).apply {
             metadata = DescriptorMetadataSource.File(CodegenUtil.getMemberDescriptorsToGenerate(ktFile, context.bindingContext))
         }
     }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrFile.kt

@@ -42,6 +42,8 @@ abstract class IrExternalPackageFragment : IrPackageFragment() {
 abstract class IrFile : IrPackageFragment(), IrMutableAnnotationContainer, IrMetadataSourceOwner {
     abstract override val symbol: IrFileSymbol
 
+    abstract val module: IrModuleFragment
+
     abstract val fileEntry: IrFileEntry
 
     override fun <D> transform(transformer: IrElementTransformer<D>, data: D): IrFile =

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/impl/IrFileImpl.kt

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.ir.IrFileEntry
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
 import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.declarations.MetadataSource
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
 import org.jetbrains.kotlin.ir.symbols.IrFileSymbol
@@ -39,10 +40,27 @@ class IrFileImpl(
         packageFragmentDescriptor: PackageFragmentDescriptor
     ) : this(fileEntry, IrFileSymbolImpl(packageFragmentDescriptor), packageFragmentDescriptor.fqName)
 
+    constructor(
+        fileEntry: IrFileEntry,
+        packageFragmentDescriptor: PackageFragmentDescriptor,
+        module: IrModuleFragment,
+    ) : this(fileEntry, IrFileSymbolImpl(packageFragmentDescriptor), packageFragmentDescriptor.fqName, module)
+
+    constructor(
+        fileEntry: IrFileEntry,
+        symbol: IrFileSymbol,
+        fqName: FqName,
+        module: IrModuleFragment
+    ) : this(fileEntry, symbol, fqName) {
+        this.module = module
+    }
+
     init {
         symbol.bind(this)
     }
 
+    override lateinit var module: IrModuleFragment
+
     override val startOffset: Int
         get() = 0
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/DeepCopyIrTreeWithSymbols.kt

@@ -53,6 +53,8 @@ open class DeepCopyIrTreeWithSymbols(
     private val symbolRenamer: SymbolRenamer
 ) : IrElementTransformerVoid() {
 
+    private var transformedModule: IrModuleFragment? = null
+
     constructor(symbolRemapper: SymbolRemapper, typeRemapper: TypeRemapper) : this(symbolRemapper, typeRemapper, SymbolRenamer.DEFAULT)
 
     init {
@@ -82,12 +84,16 @@ open class DeepCopyIrTreeWithSymbols(
     override fun visitElement(element: IrElement): IrElement =
         throw IllegalArgumentException("Unsupported element type: $element")
 
-    override fun visitModuleFragment(declaration: IrModuleFragment): IrModuleFragment =
+    override fun visitModuleFragment(declaration: IrModuleFragment): IrModuleFragment {
-        IrModuleFragmentImpl(
+        val result = IrModuleFragmentImpl(
             declaration.descriptor,
             declaration.irBuiltins,
-            declaration.files.transform()
         )
+        transformedModule = result
+        result.files += declaration.files.transform()
+        transformedModule = null
+        return result
+    }
 
     override fun visitExternalPackageFragment(declaration: IrExternalPackageFragment): IrExternalPackageFragment =
         IrExternalPackageFragmentImpl(
@@ -101,7 +107,8 @@ open class DeepCopyIrTreeWithSymbols(
         IrFileImpl(
             declaration.fileEntry,
             symbolRemapper.getDeclaredFile(declaration.symbol),
-            symbolRenamer.getFileName(declaration.symbol)
+            symbolRenamer.getFileName(declaration.symbol),
+            transformedModule ?: declaration.module
         ).apply {
             transformAnnotations(declaration)
             declaration.transformDeclarationsTo(this)

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/BasicIrModuleDeserializer.kt

@@ -103,7 +103,7 @@ abstract class BasicIrModuleDeserializer(
     private fun deserializeIrFile(fileProto: ProtoFile, fileIndex: Int, moduleDeserializer: IrModuleDeserializer, allowErrorNodes: Boolean): IrFile {
 
         val fileReader = IrLibraryFileFromKlib(moduleDeserializer.klib, fileIndex)
-        val file = fileReader.createFile(moduleDescriptor, fileProto)
+        val file = fileReader.createFile(moduleFragment, fileProto)
 
         val fileDeserializationState = FileDeserializationState(
             linker,

compiler/ir/serialization.common/src/org/jetbrains/kotlin/backend/common/serialization/IrFileDeserializer.kt

@@ -7,8 +7,7 @@ package org.jetbrains.kotlin.backend.common.serialization
 
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.impl.EmptyPackageFragmentDescriptor
-import org.jetbrains.kotlin.ir.declarations.IrDeclaration
+import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrFileSymbolImpl
 import org.jetbrains.kotlin.ir.util.IdSignature
@@ -154,11 +153,11 @@ internal fun IrLibraryFile.deserializeString(index: Int): String = WobblyTF8.dec
 internal fun IrLibraryFile.deserializeFqName(fqn: List<Int>): String =
     fqn.joinToString(".", transform = ::deserializeString)
 
-internal fun IrLibraryFile.createFile(moduleDescriptor: ModuleDescriptor, fileProto: ProtoFile): IrFile {
+internal fun IrLibraryFile.createFile(module: IrModuleFragment, fileProto: ProtoFile): IrFile {
     val fileName = fileProto.fileEntry.name
     val fileEntry = NaiveSourceBasedFileEntryImpl(fileName, fileProto.fileEntry.lineStartOffsetsList.toIntArray())
     val fqName = FqName(deserializeFqName(fileProto.fqNameList))
-    val packageFragmentDescriptor = EmptyPackageFragmentDescriptor(moduleDescriptor, fqName)
+    val packageFragmentDescriptor = EmptyPackageFragmentDescriptor(module.descriptor, fqName)
     val symbol = IrFileSymbolImpl(packageFragmentDescriptor)
-    return IrFileImpl(fileEntry, symbol, fqName)
+    return IrFileImpl(fileEntry, symbol, fqName, module)
 }