[K2] Reorganize project model for MPP (#3236)

* Reorganize project model for MPP

We map Dokka's source set directly to a source module of Analysis API inside one Analysis Standalone session.

Analysis API session is created in src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt (see fun createAnalysisSession)

Before the PR, one Dokka's source set corresponded to one Standalone API session with one source module that has source roots from dependent source sets.

* Add 'caffeine' dependency from Analysis API

* Fix sample provider

* Fix tests

* Enable tests and update the version of Analysis API.

The PR allows the enabling of some tests annotated with OnlyDescriptorsMPP.
Also, tests with OnlyDescriptorsMPP that have unresolved common symbols are fixed by the new version of Analysis API.

11 files changed, 148 insertions, 171 deletions

diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/java/KotlinDocCommentParser.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/java/KotlinDocCommentParser.kt
index 33cc4305..0ee95e45 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/java/KotlinDocCommentParser.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/java/KotlinDocCommentParser.kt
@@ -39,7 +39,7 @@ internal class KotlinDocCommentParser(
         }
         val kotlinAnalysis = context.plugin<SymbolsAnalysisPlugin>().querySingle { kotlinAnalysis }
         val elementName = element.resolveDocContext.ktElement.name
-        return analyze(kotlinAnalysis[sourceSet].mainModule) {
+        return analyze(kotlinAnalysis.getModule(sourceSet)) {
             parseFromKDocTag(
                 kDocTag = element.comment,
                 externalDri = { link -> resolveKDocLink(link).ifUnresolved { context.logger.logUnresolvedLink(link.getLinkText(), elementName) } },
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationParsingContext.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationParsingContext.kt
index ef59aa33..f5cfbdb9 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationParsingContext.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationParsingContext.kt
@@ -36,8 +36,8 @@ internal fun ModuleAndPackageDocumentationParsingContext(
     if (kotlinAnalysis == null || sourceSet == null) {
         MarkdownParser(externalDri = { null }, sourceLocation)
     } else {
-        val analysisContext = kotlinAnalysis[sourceSet]
-        val contextPsi = analyze(analysisContext.mainModule) {
+        val sourceModule = kotlinAnalysis.getModule(sourceSet)
+        val contextPsi = analyze(sourceModule) {
             val contextSymbol = when (fragment.classifier) {
                 Module -> ROOT_PACKAGE_SYMBOL
                 Package -> getPackageSymbolIfPackageExists(FqName(fragment.name))
@@ -46,7 +46,7 @@ internal fun ModuleAndPackageDocumentationParsingContext(
         }
         MarkdownParser(
             externalDri = { link ->
-                analyze(analysisContext.mainModule) {
+                analyze(sourceModule) {
                     resolveKDocTextLink(
                         link,
                         contextPsi
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationReader.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationReader.kt
index c581c7a8..ef79e885 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationReader.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/kdoc/moduledocs/ModuleAndPackageDocumentationReader.kt
@@ -38,7 +38,7 @@ private class ContextModuleAndPackageDocumentationReader(
     ): SourceSetDependent<DocumentationNode> {
         return sourceSets.associateWithNotNull { sourceSet ->
             val fragments = documentationFragments[sourceSet].orEmpty().filter(predicate)
-            kotlinAnalysis[sourceSet] // test: to throw exception for unknown sourceSet
+            kotlinAnalysis.getModule(sourceSet)// test: to throw exception for unknown sourceSet
             val documentations = fragments.map { fragment ->
                 parseModuleAndPackageDocumentation(
                     context = ModuleAndPackageDocumentationParsingContext(context.logger, kotlinAnalysis, sourceSet),
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/AnalysisContext.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/AnalysisContext.kt
index cf57e815..191c5f92 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/AnalysisContext.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/AnalysisContext.kt
@@ -5,140 +5,49 @@
 package org.jetbrains.dokka.analysis.kotlin.symbols.plugin
 
 import com.intellij.openapi.Disposable
-import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
 import org.jetbrains.dokka.DokkaConfiguration
-import org.jetbrains.dokka.DokkaSourceSetID
 import org.jetbrains.dokka.model.SourceSetDependent
 import org.jetbrains.dokka.plugability.DokkaContext
 import org.jetbrains.kotlin.analysis.api.standalone.StandaloneAnalysisAPISession
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import java.io.Closeable
-import java.io.File
 
-@Suppress("FunctionName", "UNUSED_PARAMETER")
 internal fun SamplesKotlinAnalysis(
     sourceSets: List<DokkaConfiguration.DokkaSourceSet>,
     context: DokkaContext,
-    projectKotlinAnalysis: KotlinAnalysis
-): KotlinAnalysis {
-    val environments = sourceSets
-        .filter { it.samples.isNotEmpty() }
-        .associateWith { sourceSet ->
-            createAnalysisContext(
-                classpath = sourceSet.classpath,
-                sourceRoots = sourceSet.samples,
-                sourceSet = sourceSet
-            )
-        }
-
-    return EnvironmentKotlinAnalysis(environments, projectKotlinAnalysis)
-}
+): KotlinAnalysis = createAnalysisSession(
+    sourceSets = sourceSets,
+    logger = context.logger,
+    isSampleProject = true
+)
 
 internal fun ProjectKotlinAnalysis(
     sourceSets: List<DokkaConfiguration.DokkaSourceSet>,
     context: DokkaContext,
-): KotlinAnalysis {
-    val environments = sourceSets.associateWith { sourceSet ->
-        createAnalysisContext(
-            context = context,
-            sourceSets = sourceSets,
-            sourceSet = sourceSet
-        )
-    }
-    return EnvironmentKotlinAnalysis(environments)
-}
-
-
-@Suppress("UNUSED_PARAMETER")
-internal fun createAnalysisContext(
-    context: DokkaContext,
-    sourceSets: List<DokkaConfiguration.DokkaSourceSet>,
-    sourceSet: DokkaConfiguration.DokkaSourceSet
-): AnalysisContext {
-    val parentSourceSets = sourceSets.filter { it.sourceSetID in sourceSet.dependentSourceSets }
-    val classpath = sourceSet.classpath + parentSourceSets.flatMap { it.classpath }
-    val sources = sourceSet.sourceRoots + parentSourceSets.flatMap { it.sourceRoots }
-
-    return createAnalysisContext(classpath, sources, sourceSet)
-}
-
-internal fun createAnalysisContext(
-    classpath: List<File>,
-    sourceRoots: Set<File>,
-    sourceSet: DokkaConfiguration.DokkaSourceSet
-): AnalysisContext {
-    val applicationDisposable: Disposable = Disposer.newDisposable("StandaloneAnalysisAPISession.application")
-    val projectDisposable: Disposable = Disposer.newDisposable("StandaloneAnalysisAPISession.project")
-
-    val analysis= createAnalysisSession(
-        classpath = classpath,
-        sourceRoots = sourceRoots,
-        analysisPlatform = sourceSet.analysisPlatform,
-        languageVersion = sourceSet.languageVersion,
-        apiVersion = sourceSet.apiVersion,
-        applicationDisposable = applicationDisposable,
-        projectDisposable = projectDisposable
-    )
-    return AnalysisContextImpl(
-        mainModule = analysis.second,
-        analysisSession = analysis.first,
-        applicationDisposable = applicationDisposable,
-        projectDisposable = projectDisposable
-    )
-}
-
-
-/**
- * First child delegation. It does not close [parent].
- */
-internal abstract class KotlinAnalysis(
-    private val parent: KotlinAnalysis? = null
+): KotlinAnalysis = createAnalysisSession(
+    sourceSets = sourceSets,
+    logger = context.logger
+)
+
+internal class KotlinAnalysis(
+    private val sourceModules: SourceSetDependent<KtSourceModule>,
+    private val analysisSession: StandaloneAnalysisAPISession,
+    private val applicationDisposable: Disposable,
+    private val projectDisposable: Disposable
 ) : Closeable {
 
-    operator fun get(key: DokkaConfiguration.DokkaSourceSet): AnalysisContext {
-        return get(key.sourceSetID)
-    }
-
-    internal operator fun get(key: DokkaSourceSetID): AnalysisContext {
-        return find(key)
-            ?: parent?.get(key)
-            ?: throw IllegalStateException("Missing EnvironmentAndFacade for sourceSet $key")
-    }
-
-    internal abstract fun find(sourceSetID: DokkaSourceSetID): AnalysisContext?
-}
-
-internal open class EnvironmentKotlinAnalysis(
-    private val environments: SourceSetDependent<AnalysisContext>,
-    parent: KotlinAnalysis? = null,
-) : KotlinAnalysis(parent = parent) {
+    fun getModule(sourceSet: DokkaConfiguration.DokkaSourceSet) =
+        sourceModules[sourceSet] ?: error("Missing a source module for sourceSet ${sourceSet.displayName} with id ${sourceSet.sourceSetID}")
 
-    override fun find(sourceSetID: DokkaSourceSetID): AnalysisContext? =
-        environments.entries.firstOrNull { (sourceSet, _) -> sourceSet.sourceSetID == sourceSetID }?.value
+    fun getModuleOrNull(sourceSet: DokkaConfiguration.DokkaSourceSet) =
+        sourceModules[sourceSet]
 
-    override fun close() {
-        environments.values.forEach(AnalysisContext::close)
-    }
-}
-
-internal interface AnalysisContext: Closeable {
-    val project: Project
-    val mainModule: KtSourceModule
-    val analysisSession: StandaloneAnalysisAPISession
-}
-
-private class AnalysisContextImpl(
-    override val mainModule: KtSourceModule,
-    override val analysisSession: StandaloneAnalysisAPISession,
-    private val applicationDisposable: Disposable,
-    private val projectDisposable: Disposable
-) : AnalysisContext {
-    override val project: Project
-        get() = analysisSession.project
+    val modulesWithFiles
+        get() = analysisSession.modulesWithFiles
 
     override fun close() {
         Disposer.dispose(applicationDisposable)
         Disposer.dispose(projectDisposable)
     }
-}
+}
\ No newline at end of file
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt
index a6155fb0..e074a142 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/plugin/KotlinAnalysis.kt
@@ -5,11 +5,14 @@
 package org.jetbrains.dokka.analysis.kotlin.symbols.plugin
 
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.util.Disposer
+import org.jetbrains.dokka.DokkaConfiguration
+import org.jetbrains.dokka.DokkaSourceSetID
 import org.jetbrains.dokka.Platform
+import org.jetbrains.dokka.utilities.DokkaLogger
 import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenProvider
 import org.jetbrains.kotlin.analysis.api.standalone.KtAlwaysAccessibleLifetimeTokenProvider
-import org.jetbrains.kotlin.analysis.api.standalone.StandaloneAnalysisAPISession
 import org.jetbrains.kotlin.analysis.api.standalone.buildStandaloneAnalysisAPISession
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.builder.KtModuleBuilder
@@ -30,10 +33,10 @@ internal fun Platform.toTargetPlatform() = when (this) {
     Platform.jvm -> JvmPlatforms.defaultJvmPlatform
 }
 
-private fun getJdkHomeFromSystemProperty(): File? {
+private fun getJdkHomeFromSystemProperty(logger: DokkaLogger): File? {
     val javaHome = File(System.getProperty("java.home"))
     if (!javaHome.exists()) {
-        // messageCollector.report(CompilerMessageSeverity.WARNING, "Set existed java.home to use JDK")
+        logger.error("Set existed java.home to use JDK")
         return null
     }
     return javaHome
@@ -56,58 +59,117 @@ internal fun getLanguageVersionSettings(
     )
 }
 
-// it should be changed after https://github.com/Kotlin/dokka/issues/3114
 @OptIn(KtAnalysisApiInternals::class)
 internal fun createAnalysisSession(
-    classpath: List<File>,
-    sourceRoots: Set<File>,
-    analysisPlatform: Platform,
-    languageVersion: String?,
-    apiVersion: String?,
-    applicationDisposable: Disposable,
-    projectDisposable: Disposable
-): Pair<StandaloneAnalysisAPISession, KtSourceModule> {
-
-    var sourceModule: KtSourceModule? = null
+    sourceSets: List<DokkaConfiguration.DokkaSourceSet>,
+    logger: DokkaLogger,
+    applicationDisposable: Disposable = Disposer.newDisposable("StandaloneAnalysisAPISession.application"),
+    projectDisposable: Disposable = Disposer.newDisposable("StandaloneAnalysisAPISession.project"),
+    isSampleProject: Boolean = false
+): KotlinAnalysis {
+    val sourcesModule = mutableMapOf<DokkaConfiguration.DokkaSourceSet, KtSourceModule>()
+
     val analysisSession = buildStandaloneAnalysisAPISession(
         applicationDisposable = applicationDisposable,
         projectDisposable = projectDisposable,
         withPsiDeclarationFromBinaryModuleProvider = false
     ) {
         registerProjectService(KtLifetimeTokenProvider::class.java, KtAlwaysAccessibleLifetimeTokenProvider())
-        val targetPlatform = analysisPlatform.toTargetPlatform()
+
+        val sortedSourceSets = topologicalSortByDependantSourceSets(sourceSets, logger)
+
+        val sourcesModuleBySourceSetId = mutableMapOf<DokkaSourceSetID, KtSourceModule>()
 
         buildKtModuleProvider {
-            val libraryRoots = classpath
-            fun KtModuleBuilder.addModuleDependencies(moduleName: String) {
+            val jdkModule = getJdkHomeFromSystemProperty(logger)?.let { jdkHome ->
+                buildKtSdkModule {
+                    this.platform = Platform.jvm.toTargetPlatform()
+                    addBinaryRootsFromJdkHome(jdkHome.toPath(), isJre = true)
+                    sdkName = "JDK"
+                }
+            }
+
+            fun KtModuleBuilder.addModuleDependencies(sourceSet: DokkaConfiguration.DokkaSourceSet) {
+                val targetPlatform = sourceSet.analysisPlatform.toTargetPlatform()
                 addRegularDependency(
                     buildKtLibraryModule {
                         this.platform = targetPlatform
-                        addBinaryRoots(libraryRoots.map { it.toPath() })
-                        libraryName = "Library for $moduleName"
+                        addBinaryRoots(sourceSet.classpath.map { it.toPath() })
+                        libraryName = "Library for ${sourceSet.displayName}"
                     }
                 )
-                getJdkHomeFromSystemProperty()?.let { jdkHome ->
+                if (sourceSet.analysisPlatform == Platform.jvm) {
+                    jdkModule?.let { addRegularDependency(it) }
+                }
+                sourceSet.dependentSourceSets.forEach {
                     addRegularDependency(
-                        buildKtSdkModule {
-                            this.platform = targetPlatform
-                            addBinaryRootsFromJdkHome(jdkHome.toPath(), isJre = true)
-                            sdkName = "JDK for $moduleName"
-                        }
+                        sourcesModuleBySourceSetId[it]
+                            ?: error("There is no source module for $it")
+                    )
+                }
+            }
+
+            for (sourceSet in sortedSourceSets) {
+                val targetPlatform = sourceSet.analysisPlatform.toTargetPlatform()
+                val sourceModule = buildKtSourceModule {
+                    languageVersionSettings =
+                        getLanguageVersionSettings(sourceSet.languageVersion, sourceSet.apiVersion)
+                    platform = targetPlatform
+                    moduleName = "<module ${sourceSet.displayName}>"
+                    if (isSampleProject)
+                        addSourceRoots(sourceSet.samples.map { it.toPath() })
+                    else
+                        addSourceRoots(sourceSet.sourceRoots.map { it.toPath() })
+                    addModuleDependencies(
+                        sourceSet,
                     )
                 }
+                sourcesModule[sourceSet] = sourceModule
+                sourcesModuleBySourceSetId[sourceSet.sourceSetID] = sourceModule
+                addModule(sourceModule)
             }
-            sourceModule = buildKtSourceModule {
-                languageVersionSettings = getLanguageVersionSettings(languageVersion, apiVersion)
-                platform = targetPlatform
-                moduleName = "<module>"
-                // TODO: We should handle (virtual) file changes announced via LSP with the VFS
-                addSourceRoots(sourceRoots.map { it.toPath() })
-                addModuleDependencies(moduleName)
+            platform = sourceSets.map { it.analysisPlatform }.distinct().singleOrNull()?.toTargetPlatform()
+                ?: Platform.common.toTargetPlatform()
+        }
+    }
+    return KotlinAnalysis(sourcesModule, analysisSession, applicationDisposable, projectDisposable)
+}
+
+private enum class State {
+    UNVISITED,
+    VISITING,
+    VISITED;
+}
+
+internal fun topologicalSortByDependantSourceSets(
+    sourceSets: List<DokkaConfiguration.DokkaSourceSet>,
+    logger: DokkaLogger
+): List<DokkaConfiguration.DokkaSourceSet> {
+    val result = mutableListOf<DokkaConfiguration.DokkaSourceSet>()
+
+    val verticesAssociatedWithState = sourceSets.associateWithTo(mutableMapOf()) { State.UNVISITED }
+    fun dfs(souceSet: DokkaConfiguration.DokkaSourceSet) {
+        when (verticesAssociatedWithState[souceSet]) {
+            State.VISITED -> return
+            State.VISITING -> {
+                logger.error("Detected cycle in source set graph")
+                return
+            }
+
+            else -> {
+                val dependentSourceSets =
+                    souceSet.dependentSourceSets.mapNotNull { dependentSourceSetId ->
+                        sourceSets.find { it.sourceSetID == dependentSourceSetId }
+                        // just skip
+                            ?: null.also { logger.error("Unknown source set Id $dependentSourceSetId in dependencies of ${souceSet.sourceSetID}") }
+                    }
+                verticesAssociatedWithState[souceSet] = State.VISITING
+                dependentSourceSets.forEach(::dfs)
+                verticesAssociatedWithState[souceSet] = State.VISITED
+                result += souceSet
             }
-            platform = targetPlatform
-            addModule(sourceModule!!)
         }
     }
-    return Pair(analysisSession, sourceModule ?: throw IllegalStateException())
+    sourceSets.forEach(::dfs)
+    return result
 }
\ No newline at end of file
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinAnalysisProjectProvider.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinAnalysisProjectProvider.kt
index 398004a4..398d48ee 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinAnalysisProjectProvider.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinAnalysisProjectProvider.kt
@@ -15,6 +15,6 @@ import org.jetbrains.dokka.plugability.querySingle
 internal class KotlinAnalysisProjectProvider : ProjectProvider {
     override fun getProject(sourceSet: DokkaConfiguration.DokkaSourceSet, context: DokkaContext): Project {
         val kotlinAnalysis = context.plugin<SymbolsAnalysisPlugin>().querySingle { kotlinAnalysis }
-        return kotlinAnalysis[sourceSet].project
+        return kotlinAnalysis.getModule(sourceSet).project
     }
 }
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinSampleProvider.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinSampleProvider.kt
index c17ad75f..e453c72d 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinSampleProvider.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/KotlinSampleProvider.kt
@@ -15,6 +15,7 @@ import org.jetbrains.dokka.analysis.kotlin.internal.SampleProviderFactory
 import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.SamplesKotlinAnalysis
 import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.SymbolsAnalysisPlugin
 import org.jetbrains.kotlin.analysis.api.analyze
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtBlockExpression
@@ -37,10 +38,10 @@ public class KotlinSampleProviderFactory(
 public open class KotlinSampleProvider(
     public val context: DokkaContext
 ): SampleProvider {
-    private val kotlinAnalysis = SamplesKotlinAnalysis(
-        sourceSets = context.configuration.sourceSets,
-        context = context,
-        projectKotlinAnalysis = context.plugin<SymbolsAnalysisPlugin>().querySingle { kotlinAnalysis }
+    private val kotlinAnalysisOfRegularSources = context.plugin<SymbolsAnalysisPlugin>().querySingle { kotlinAnalysis }
+
+    private val kotlinAnalysisOfSamples = SamplesKotlinAnalysis(
+        sourceSets = context.configuration.sourceSets, context = context
     )
 
     protected open fun processBody(psiElement: PsiElement): String {
@@ -72,8 +73,13 @@ public open class KotlinSampleProvider(
      * @return [SampleProvider.SampleSnippet] or null if it has not found by [fqLink]
      */
     override fun getSample(sourceSet: DokkaConfiguration.DokkaSourceSet, fqLink: String): SampleProvider.SampleSnippet? {
-        val analysisContext = kotlinAnalysis[sourceSet]
-        val psiElement = analyze(analysisContext.mainModule) {
+        return kotlinAnalysisOfSamples.getModuleOrNull(sourceSet)?.let { getSampleFromModule(it, fqLink) }
+            ?: getSampleFromModule(
+                kotlinAnalysisOfRegularSources.getModule(sourceSet), fqLink
+            )
+    }
+    private fun getSampleFromModule(module: KtSourceModule, fqLink: String): SampleProvider.SampleSnippet? {
+        val psiElement = analyze(module) {
             val lastDotIndex = fqLink.lastIndexOf('.')
 
             val functionName = if (lastDotIndex == -1) fqLink else fqLink.substring(lastDotIndex + 1, fqLink.length)
@@ -87,7 +93,8 @@ public open class KotlinSampleProvider(
 
         return SampleProvider.SampleSnippet(imports, body)
     }
+
     override fun close() {
-        kotlinAnalysis.close()
+        kotlinAnalysisOfSamples.close()
     }
 }
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolExternalDocumentablesProvider.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolExternalDocumentablesProvider.kt
index 52496b78..1473a7da 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolExternalDocumentablesProvider.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolExternalDocumentablesProvider.kt
@@ -24,10 +24,9 @@ internal class SymbolExternalDocumentablesProvider(val context: DokkaContext) :
     override fun findClasslike(dri: DRI, sourceSet: DokkaSourceSet): DClasslike? {
         val classId = getClassIdFromDRI(dri)
 
-        val analysisContext = kotlinAnalysis[sourceSet]
-        return analyze(analysisContext.mainModule) {
+        return analyze(kotlinAnalysis.getModule(sourceSet)) {
             val symbol = getClassOrObjectSymbolByClassId(classId) as? KtNamedClassOrObjectSymbol?: return@analyze null
-            val translator = DokkaSymbolVisitor(sourceSet, sourceSet.displayName, analysisContext, logger = context.logger)
+            val translator = DokkaSymbolVisitor(sourceSet, sourceSet.displayName, kotlinAnalysis, logger = context.logger)
 
             val parentDRI = symbol.getContainingSymbol()?.let { getDRIFromSymbol(it) } ?: /* top level */ DRI(dri.packageName)
             with(translator) {
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolFullClassHierarchyBuilder.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolFullClassHierarchyBuilder.kt
index 0e90bec8..497d7946 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolFullClassHierarchyBuilder.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/services/SymbolFullClassHierarchyBuilder.kt
@@ -79,7 +79,7 @@ internal class SymbolFullClassHierarchyBuilder(val context: DokkaContext) : Full
             documentable.sources.forEach { (sourceSet, source) ->
                 if (source is KtPsiDocumentableSource) {
                     (source.psi as? KtClassOrObject)?.let { psi ->
-                        analyze(kotlinAnalysis[sourceSet].mainModule) {
+                        analyze(kotlinAnalysis.getModule(sourceSet)) {
                             val type = psi.getNamedClassOrObjectSymbol()?.buildSelfClassType() ?: return@analyze
                             hierarchy[sourceSet]?.let { collectSupertypesFromKtType(documentable.dri to type, it) }
                         }
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DRIFactory.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DRIFactory.kt
index 5c879156..a2cb423a 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DRIFactory.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DRIFactory.kt
@@ -103,7 +103,7 @@ internal fun KtAnalysisSession.getDRIFromSymbol(symbol: KtSymbol): DRI =
         is KtFunctionLikeSymbol -> getDRIFromFunctionLike(symbol)
         is KtClassLikeSymbol -> getDRIFromClassLike(symbol)
         is KtPackageSymbol -> getDRIFromPackage(symbol)
-        else -> throw IllegalStateException("Unknown symbol while creating DRI ")
+        else -> throw IllegalStateException("Unknown symbol while creating DRI $symbol")
     }
 
 private fun KtAnalysisSession.getDRIFromNonCallablePossibleLocalSymbol(symbol: KtSymbol): DRI {
diff --git a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DefaultSymbolToDocumentableTranslator.kt b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DefaultSymbolToDocumentableTranslator.kt
index 74d31269..6c7071f5 100644
--- a/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DefaultSymbolToDocumentableTranslator.kt
+++ b/subprojects/analysis-kotlin-symbols/src/main/kotlin/org/jetbrains/dokka/analysis/kotlin/symbols/translators/DefaultSymbolToDocumentableTranslator.kt
@@ -11,7 +11,6 @@ import org.jetbrains.dokka.DokkaConfiguration
 import org.jetbrains.dokka.analysis.java.JavaAnalysisPlugin
 import org.jetbrains.dokka.analysis.java.parsers.JavadocParser
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.getGeneratedKDocDocumentationFrom
-import org.jetbrains.dokka.analysis.kotlin.symbols.plugin.AnalysisContext
 import org.jetbrains.dokka.analysis.kotlin.symbols.services.KtPsiDocumentableSource
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.getJavaDocDocumentationFrom
 import org.jetbrains.dokka.analysis.kotlin.symbols.kdoc.getKDocDocumentationFrom
@@ -59,7 +58,7 @@ internal class DefaultSymbolToDocumentableTranslator(context: DokkaContext) : As
         sourceSet: DokkaConfiguration.DokkaSourceSet,
         context: DokkaContext
     ): DModule {
-        val analysisContext = kotlinAnalysis[sourceSet]
+        val analysisContext = kotlinAnalysis
         @Suppress("unused")
         return DokkaSymbolVisitor(
             sourceSet = sourceSet,
@@ -84,7 +83,7 @@ internal fun <T : Bound> T.wrapWithVariance(variance: org.jetbrains.kotlin.types
 internal class DokkaSymbolVisitor(
     private val sourceSet: DokkaConfiguration.DokkaSourceSet,
     private val moduleName: String,
-    private val analysisContext: AnalysisContext,
+    private val analysisContext: KotlinAnalysis,
     private val logger: DokkaLogger,
     private val javadocParser: JavadocParser? = null
 ) {
@@ -118,9 +117,10 @@ internal class DokkaSymbolVisitor(
     }
 
     fun visitModule(): DModule {
-        val ktFiles = analysisContext.analysisSession.modulesWithFiles.entries.single().value.filterIsInstance<KtFile>()
+        val sourceModule = analysisContext.getModule(sourceSet)
+        val ktFiles = analysisContext.modulesWithFiles[sourceModule]?.filterIsInstance<KtFile>() ?: throw IllegalStateException("No source files for a source module ${sourceModule.moduleName} of source set ${sourceSet.sourceSetID}")
         val processedPackages: MutableSet<FqName> = mutableSetOf()
-        return analyze(analysisContext.mainModule) {
+        return analyze(sourceModule) {
             val packageSymbols: List<DPackage> = ktFiles
                 .mapNotNull {
                     if (processedPackages.contains(it.packageFqName))