Support KLIB for Common platform (#2441)

3 files changed, 204 insertions, 19 deletions

diff --git a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/AnalysisEnvironment.kt b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/AnalysisEnvironment.kt
index 393a2c62..452b821c 100644
--- a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/AnalysisEnvironment.kt
+++ b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/AnalysisEnvironment.kt
@@ -20,6 +20,7 @@ import org.jetbrains.dokka.Platform
 import org.jetbrains.dokka.analysis.resolve.*
 import org.jetbrains.kotlin.analyzer.*
 import org.jetbrains.kotlin.analyzer.common.CommonAnalysisParameters
+import org.jetbrains.kotlin.analyzer.common.CommonDependenciesContainer
 import org.jetbrains.kotlin.analyzer.common.CommonPlatformAnalyzerServices
 import org.jetbrains.kotlin.analyzer.common.CommonResolverForModuleFactory
 import org.jetbrains.kotlin.builtins.DefaultBuiltIns
@@ -50,6 +51,7 @@ import org.jetbrains.kotlin.idea.klib.getCompatibilityInfo
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
 import org.jetbrains.kotlin.js.resolve.JsPlatformAnalyzerServices
 import org.jetbrains.kotlin.library.KLIB_FILE_EXTENSION
+import org.jetbrains.kotlin.library.KotlinLibrary
 import org.jetbrains.kotlin.library.ToolingSingleFileKlibResolveStrategy
 import org.jetbrains.kotlin.library.resolveSingleFileKlib
 import org.jetbrains.kotlin.load.java.structure.impl.JavaClassImpl
@@ -184,7 +186,15 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
             Platform.jvm -> JvmPlatforms.defaultJvmPlatform
         }
 
-        val nativeLibraries: Map<AbsolutePathString, LibraryModuleInfo> = loadNativeLibraries()
+        val kotlinLibraries: Map<AbsolutePathString, KotlinLibrary> = resolveKotlinLibraries()
+
+        val commonDependencyContainer = if (analysisPlatform == Platform.common) DokkaKlibMetadataCommonDependencyContainer(
+            kotlinLibraries.values.toList(),
+            environment.configuration,
+            projectContext.storageManager
+        ) else null
+
+        val extraModuleDependencies = kotlinLibraries.values.registerLibraries() + commonDependencyContainer?.moduleInfos.orEmpty()
 
         val library = object : LibraryModuleInfo {
             override val analyzerServices: PlatformDependentAnalyzerServices =
@@ -194,7 +204,7 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
             override fun dependencies(): List<ModuleInfo> = listOf(this)
             override fun getLibraryRoots(): Collection<String> = classpath
                 .map { libraryFile -> libraryFile.absolutePath }
-                .filter { path -> path !in nativeLibraries }
+                .filter { path -> path !in kotlinLibraries }
         }
 
         val module = object : ModuleInfo {
@@ -202,7 +212,8 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
                 analysisPlatform.analyzerServices()
             override val name: Name = Name.special("<module>")
             override val platform: TargetPlatform = targetPlatform
-            override fun dependencies(): List<ModuleInfo> = listOf(this, library) + nativeLibraries.values
+            override fun dependencies(): List<ModuleInfo> =
+                listOf(this, library) + extraModuleDependencies
         }
 
         val sourcesScope = createSourceModuleSearchScope(environment.project, sourceFiles)
@@ -211,10 +222,11 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
                 library -> ModuleContent(it, emptyList(), GlobalSearchScope.notScope(sourcesScope))
                 module -> ModuleContent(it, emptyList(), GlobalSearchScope.allScope(environment.project))
                 is DokkaKlibLibraryInfo -> {
-                    if (it.libraryRoot in nativeLibraries)
+                    if (it.libraryRoot in kotlinLibraries)
                         ModuleContent(it, emptyList(), GlobalSearchScope.notScope(sourcesScope))
                     else null
                 }
+                is CommonKlibModuleInfo -> ModuleContent(it, emptyList(), GlobalSearchScope.notScope(sourcesScope))
                 else -> null
             } ?: throw IllegalArgumentException("Unexpected module info")
         }
@@ -240,7 +252,8 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
                 projectContext,
                 module,
                 modulesContent,
-                environment
+                environment,
+                commonDependencyContainer
             )
             Platform.js -> createJsResolverForProject(projectContext, module, modulesContent)
             Platform.native -> createNativeResolverForProject(projectContext, module, modulesContent)
@@ -282,16 +295,26 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
         Platform.jvm -> JvmPlatformAnalyzerServices
     }
 
-    @OptIn(ExperimentalStdlibApi::class)
-    private fun loadNativeLibraries(): Map<AbsolutePathString, LibraryModuleInfo> {
-        if (analysisPlatform != Platform.native && analysisPlatform != Platform.js) return emptyMap()
-
+    fun Collection<KotlinLibrary>.registerLibraries(): List<DokkaKlibLibraryInfo> {
+        if (analysisPlatform != Platform.native && analysisPlatform != Platform.js) return emptyList()
         val dependencyResolver = DokkaKlibLibraryDependencyResolver()
         val analyzerServices = analysisPlatform.analyzerServices()
 
-        return buildMap {
+        return map { kotlinLibrary ->
+            if (analysisPlatform == org.jetbrains.dokka.Platform.native) DokkaNativeKlibLibraryInfo(
+                kotlinLibrary,
+                analyzerServices,
+                dependencyResolver
+            )
+            else DokkaJsKlibLibraryInfo(kotlinLibrary, analyzerServices, dependencyResolver)
+        }
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun resolveKotlinLibraries(): Map<AbsolutePathString, KotlinLibrary> {
+        return if (analysisPlatform == Platform.jvm) emptyMap() else buildMap {
             classpath
-                .filter { it.isDirectory || (it.extension == "jar" || it.extension == KLIB_FILE_EXTENSION) }
+                .filter { it.isDirectory || it.extension == KLIB_FILE_EXTENSION }
                 .forEach { libraryFile ->
                     try {
                         val kotlinLibrary = resolveSingleFileKlib(
@@ -303,12 +326,7 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
                             // exists, is KLIB, has compatible format
                             put(
                                 libraryFile.absolutePath,
-                                if (analysisPlatform == Platform.native) DokkaNativeKlibLibraryInfo(
-                                    kotlinLibrary,
-                                    analyzerServices,
-                                    dependencyResolver
-                                )
-                                else DokkaJsKlibLibraryInfo(kotlinLibrary, analyzerServices, dependencyResolver)
+                                kotlinLibrary
                             )
                         }
                     } catch (e: Throwable) {
@@ -323,7 +341,8 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
         projectContext: ProjectContext,
         module: ModuleInfo,
         modulesContent: (ModuleInfo) -> ModuleContent<ModuleInfo>,
-        environment: KotlinCoreEnvironment
+        environment: KotlinCoreEnvironment,
+        dependencyContainer: CommonDependenciesContainer?
     ): ResolverForProject<ModuleInfo> {
         return object : AbstractResolverForProject<ModuleInfo>(
             "Dokka",
@@ -344,7 +363,8 @@ class AnalysisEnvironment(val messageCollector: MessageCollector, val analysisPl
                     },
                     CompilerEnvironment,
                     unspecifiedJvmPlatform,
-                    true
+                    true,
+                    dependencyContainer
                 ).createResolverForModule(
                     descriptor as ModuleDescriptorImpl,
                     projectContext.withModule(descriptor),
diff --git a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/resolve/CommonKlibModuleInfo.kt b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/resolve/CommonKlibModuleInfo.kt
new file mode 100644
index 00000000..22c86dd7
--- /dev/null
+++ b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/resolve/CommonKlibModuleInfo.kt
@@ -0,0 +1,25 @@
+package org.jetbrains.dokka.analysis.resolve
+
+import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.analyzer.common.CommonPlatformAnalyzerServices
+import org.jetbrains.kotlin.library.KotlinLibrary
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.platform.CommonPlatforms
+import org.jetbrains.kotlin.platform.TargetPlatform
+import org.jetbrains.kotlin.resolve.PlatformDependentAnalyzerServices
+
+internal class CommonKlibModuleInfo(
+    override val name: Name,
+    val kotlinLibrary: KotlinLibrary,
+    private val dependOnModules: List<ModuleInfo>
+) : ModuleInfo {
+    override fun dependencies(): List<ModuleInfo> = dependOnModules
+
+    override fun dependencyOnBuiltIns(): ModuleInfo.DependencyOnBuiltIns = ModuleInfo.DependencyOnBuiltIns.LAST
+
+    override val platform: TargetPlatform
+        get() = CommonPlatforms.defaultCommonPlatform
+
+    override val analyzerServices: PlatformDependentAnalyzerServices
+        get() = CommonPlatformAnalyzerServices
+}
\ No newline at end of file
diff --git a/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/resolve/DokkaKlibMetadataCommonDependencyContainer.kt b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/resolve/DokkaKlibMetadataCommonDependencyContainer.kt
new file mode 100644
index 00000000..1a987a1f
--- /dev/null
+++ b/kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/resolve/DokkaKlibMetadataCommonDependencyContainer.kt
@@ -0,0 +1,140 @@
+package org.jetbrains.dokka.analysis.resolve
+
+import org.jetbrains.kotlin.analyzer.ModuleInfo
+import org.jetbrains.kotlin.analyzer.common.CommonDependenciesContainer
+import org.jetbrains.kotlin.builtins.DefaultBuiltIns
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.languageVersionSettings
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.descriptors.PackageFragmentProvider
+import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
+import org.jetbrains.kotlin.descriptors.konan.DeserializedKlibModuleOrigin
+import org.jetbrains.kotlin.incremental.components.LookupTracker
+import org.jetbrains.kotlin.konan.util.KlibMetadataFactories
+import org.jetbrains.kotlin.library.KotlinLibrary
+import org.jetbrains.kotlin.library.metadata.NativeTypeTransformer
+import org.jetbrains.kotlin.library.metadata.NullFlexibleTypeDeserializer
+import org.jetbrains.kotlin.library.metadata.parseModuleHeader
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.CompilerDeserializationConfiguration
+import org.jetbrains.kotlin.serialization.konan.impl.KlibMetadataModuleDescriptorFactoryImpl
+import org.jetbrains.kotlin.storage.LockBasedStorageManager
+import org.jetbrains.kotlin.storage.StorageManager
+import org.jetbrains.kotlin.utils.keysToMap
+
+/**
+ * Adapted from org.jetbrains.kotlin.cli.metadata.KlibMetadataDependencyContainer
+ */
+class DokkaKlibMetadataCommonDependencyContainer(
+    kotlinLibraries: List<KotlinLibrary>,
+    private val configuration: CompilerConfiguration,
+    private val storageManager: StorageManager
+) : CommonDependenciesContainer {
+
+    private val builtIns
+        get() = DefaultBuiltIns.Instance
+
+    private val mutableDependenciesForAllModuleDescriptors = mutableListOf<ModuleDescriptorImpl>().apply {
+        add(builtIns.builtInsModule)
+    }
+
+    private val mutableDependenciesForAllModules = mutableListOf<ModuleInfo>()
+
+    private val moduleDescriptorsForKotlinLibraries: Map<KotlinLibrary, ModuleDescriptorImpl> =
+        kotlinLibraries.keysToMap { library ->
+            val moduleHeader = parseModuleHeader(library.moduleHeaderData)
+            val moduleName = Name.special(moduleHeader.moduleName)
+            val moduleOrigin = DeserializedKlibModuleOrigin(library)
+            MetadataFactories.DefaultDescriptorFactory.createDescriptor(
+                moduleName, storageManager, builtIns, moduleOrigin
+            )
+        }.also { result ->
+            val resultValues = result.values
+            resultValues.forEach { module ->
+                module.setDependencies(mutableDependenciesForAllModuleDescriptors)
+            }
+            mutableDependenciesForAllModuleDescriptors.addAll(resultValues)
+        }
+
+    private val moduleInfosImpl: List<CommonKlibModuleInfo> = mutableListOf<CommonKlibModuleInfo>().apply {
+        addAll(
+            moduleDescriptorsForKotlinLibraries.map { (kotlinLibrary, moduleDescriptor) ->
+                CommonKlibModuleInfo(moduleDescriptor.name, kotlinLibrary, mutableDependenciesForAllModules)
+            }
+        )
+        mutableDependenciesForAllModules.addAll(this@apply)
+    }
+
+    override val moduleInfos: List<ModuleInfo> get() = moduleInfosImpl
+
+    /* not used in Dokka */
+    override val friendModuleInfos: List<ModuleInfo> = emptyList()
+
+    /* not used in Dokka */
+    override val refinesModuleInfos: List<ModuleInfo> = emptyList()
+
+    override fun moduleDescriptorForModuleInfo(moduleInfo: ModuleInfo): ModuleDescriptor {
+        if (moduleInfo !in moduleInfos)
+            error("Unknown module info $moduleInfo")
+
+        // Ensure that the package fragment provider has been created and the module descriptor has been
+        // initialized with the package fragment provider:
+        packageFragmentProviderForModuleInfo(moduleInfo)
+
+        return moduleDescriptorsForKotlinLibraries.getValue((moduleInfo as CommonKlibModuleInfo).kotlinLibrary)
+    }
+
+    override fun registerDependencyForAllModules(
+        moduleInfo: ModuleInfo,
+        descriptorForModule: ModuleDescriptorImpl
+    ) {
+        mutableDependenciesForAllModules.add(moduleInfo)
+        mutableDependenciesForAllModuleDescriptors.add(descriptorForModule)
+    }
+
+    override fun packageFragmentProviderForModuleInfo(
+        moduleInfo: ModuleInfo
+    ): PackageFragmentProvider? {
+        if (moduleInfo !in moduleInfos)
+            return null
+        return packageFragmentProviderForKotlinLibrary((moduleInfo as CommonKlibModuleInfo).kotlinLibrary)
+    }
+
+    private val klibMetadataModuleDescriptorFactory by lazy {
+        KlibMetadataModuleDescriptorFactoryImpl(
+            MetadataFactories.DefaultDescriptorFactory,
+            MetadataFactories.DefaultPackageFragmentsFactory,
+            MetadataFactories.flexibleTypeDeserializer,
+            MetadataFactories.platformDependentTypeTransformer
+        )
+    }
+
+    private fun packageFragmentProviderForKotlinLibrary(
+        library: KotlinLibrary
+    ): PackageFragmentProvider {
+        val languageVersionSettings = configuration.languageVersionSettings
+
+        val libraryModuleDescriptor = moduleDescriptorsForKotlinLibraries.getValue(library)
+        val packageFragmentNames = parseModuleHeader(library.moduleHeaderData).packageFragmentNameList
+
+        return klibMetadataModuleDescriptorFactory.createPackageFragmentProvider(
+            library,
+            packageAccessHandler = null,
+            packageFragmentNames = packageFragmentNames,
+            storageManager = LockBasedStorageManager("KlibMetadataPackageFragmentProvider"),
+            moduleDescriptor = libraryModuleDescriptor,
+            configuration = CompilerDeserializationConfiguration(languageVersionSettings),
+            compositePackageFragmentAddend = null,
+            lookupTracker = LookupTracker.DO_NOTHING
+        ).also {
+            libraryModuleDescriptor.initialize(it)
+        }
+    }
+}
+
+private val MetadataFactories =
+    KlibMetadataFactories(
+        { DefaultBuiltIns.Instance },
+        NullFlexibleTypeDeserializer,
+        NativeTypeTransformer()
+    )
\ No newline at end of file