[K/N][gradle] Supported friend modules for two stage compilation

Fixes https://youtrack.jetbrains.com/issue/KT-49248

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/KlibResolvedModuleDescriptorsFactory.kt

@@ -5,6 +5,7 @@ import org.jetbrains.kotlin.backend.common.serialization.metadata.KlibMetadataMo
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
+import org.jetbrains.kotlin.konan.file.File
 import org.jetbrains.kotlin.library.KotlinLibrary
 import org.jetbrains.kotlin.storage.StorageManager
 
@@ -29,23 +30,25 @@ interface KlibResolvedModuleDescriptorsFactory {
         storageManager: StorageManager,
         builtIns: KotlinBuiltIns?,
         languageVersionSettings: LanguageVersionSettings,
-        customAction: ((KotlinLibrary, ModuleDescriptorImpl) -> Unit)? = null,
+        friendModuleFiles: Set<File>,
+        includedLibraryFiles: Set<File>,
         additionalDependencyModules: Iterable<ModuleDescriptorImpl>
     ): KotlinResolvedModuleDescriptors
 }
 
 class KotlinResolvedModuleDescriptors(
+    /**
+     * The list of modules each representing an individual Kotlin/Native library. All modules
+     * in this list have properly installed dependencies, i.e. module has all necessary dependencies
+     * on other modules plus a dependency on the [forwardDeclarationsModule].
+     */
+    val resolvedDescriptors: List<ModuleDescriptorImpl>,
 
-        /**
-         * The list of modules each representing an individual Kotlin/Native library. All modules
-         * in this list have properly installed dependencies, i.e. module has all necessary dependencies
-         * on other modules plus a dependency on the [forwardDeclarationsModule].
-         */
-        val resolvedDescriptors: List<ModuleDescriptorImpl>,
+    /**
+     * This is a module which "contains" forward declarations.
+     * Note: this module should be unique per compilation and should always be the last dependency of any module.
+     */
+    val forwardDeclarationsModule: ModuleDescriptorImpl,
 
-        /**
-         * This is a module which "contains" forward declarations.
-         * Note: this module should be unique per compilation and should always be the last dependency of any module.
-         */
-        val forwardDeclarationsModule: ModuleDescriptorImpl
+    val friendModules: Set<ModuleDescriptorImpl>
 )

compiler/util-klib-metadata/src/org/jetbrains/kotlin/library/metadata/impl/KlibResolvedModuleDescriptorsFactoryImpl.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.descriptors.konan.KlibModuleOrigin
 import org.jetbrains.kotlin.descriptors.konan.SyntheticModulesOrigin
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
+import org.jetbrains.kotlin.konan.file.File
 import org.jetbrains.kotlin.util.profile
 import org.jetbrains.kotlin.library.KotlinLibrary
 import org.jetbrains.kotlin.library.metadata.PackageAccessHandler
@@ -37,14 +38,19 @@ class KlibResolvedModuleDescriptorsFactoryImpl(
         storageManager: StorageManager,
         builtIns: KotlinBuiltIns?,
         languageVersionSettings: LanguageVersionSettings,
-        customAction: ((KotlinLibrary, ModuleDescriptorImpl) -> Unit)?,
+        friendModuleFiles: Set<File>,
+        includedLibraryFiles: Set<File>,
         additionalDependencyModules: Iterable<ModuleDescriptorImpl>
     ): KotlinResolvedModuleDescriptors {
 
         val moduleDescriptors = mutableListOf<ModuleDescriptorImpl>()
+
         @Suppress("NAME_SHADOWING")
         var builtIns = builtIns
 
+        val friendModuleDescriptors = mutableSetOf<ModuleDescriptorImpl>()
+        val includedLibraryDescriptors = mutableSetOf<ModuleDescriptorImpl>()
+
         // Build module descriptors.
         resolvedLibraries.forEach { library, packageAccessHandler ->
             profile("Loading ${library.libraryName}") {
@@ -56,7 +62,10 @@ class KlibResolvedModuleDescriptorsFactoryImpl(
                 builtIns = moduleDescriptor.builtIns
                 moduleDescriptors.add(moduleDescriptor)
 
-                customAction?.invoke(library, moduleDescriptor)
+                if (friendModuleFiles.contains(library.libraryFile))
+                    friendModuleDescriptors.add(moduleDescriptor)
+                if (includedLibraryFiles.contains(library.libraryFile))
+                    includedLibraryDescriptors.add(moduleDescriptor)
             }
         }
 
@@ -64,14 +73,17 @@ class KlibResolvedModuleDescriptorsFactoryImpl(
 
         // Set inter-dependencies between module descriptors, add forwarding declarations module.
         for (module in moduleDescriptors) {
-            // Yes, just to all of them.
+            val friends = additionalDependencyModules.toMutableSet()
+            if (module in includedLibraryDescriptors)
+                friends.addAll(friendModuleDescriptors)
             module.setDependencies(
+                // Yes, just to all of them.
                 moduleDescriptors + additionalDependencyModules + forwardDeclarationsModule,
-                additionalDependencyModules.toSet()
+                friends
             )
         }
 
-        return KotlinResolvedModuleDescriptors(moduleDescriptors, forwardDeclarationsModule)
+        return KotlinResolvedModuleDescriptors(moduleDescriptors, forwardDeclarationsModule, friendModuleDescriptors)
     }
 
     fun createForwardDeclarationsModule(