[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.

22 files changed, 246 insertions, 206 deletions

diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 9a25490b..c7202839 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -11,7 +11,7 @@ kotlinx-bcv = "0.13.2"
 
 ## Analysis
 kotlin-compiler = "1.9.10"
-kotlin-compiler-k2 = "1.9.30-dev-3330"
+kotlin-compiler-k2 = "2.0.0-dev-5387"
 
 # MUST match the version of the intellij platform used in the kotlin compiler,
 # otherwise this will lead to different versions of psi API and implementations
diff --git a/plugins/base/src/test/kotlin/content/annotations/SinceKotlinTest.kt b/plugins/base/src/test/kotlin/content/annotations/SinceKotlinTest.kt
index 4cb25704..6ee95bbd 100644
--- a/plugins/base/src/test/kotlin/content/annotations/SinceKotlinTest.kt
+++ b/plugins/base/src/test/kotlin/content/annotations/SinceKotlinTest.kt
@@ -14,10 +14,7 @@ import org.jetbrains.dokka.model.doc.CustomTagWrapper
 import org.jetbrains.dokka.model.doc.Text
 import org.jetbrains.dokka.pages.ContentPage
 import signatures.AbstractRenderingTest
-import utils.ParamAttributes
-import utils.TestOutputWriterPlugin
-import utils.assertNotNull
-import utils.bareSignature
+import utils.*
 import kotlin.test.*
 
 
@@ -125,30 +122,58 @@ class SinceKotlinTest : AbstractRenderingTest() {
         val configuration =   dokkaConfiguration {
             sourceSets {
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/jvm/")
                     analysisPlatform = "jvm"
                 }
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/native/")
                     analysisPlatform = "native"
+                    name = "native"
                 }
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/common/")
                     analysisPlatform = "common"
+                    name = "common"
                 }
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/js/")
                     analysisPlatform = "js"
+                    name = "js"
                 }
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/wasm/")
                     analysisPlatform = "wasm"
+                    name = "wasm"
                 }
             }
         }
         testInline(
             """
-            |/src/main/kotlin/test/source.kt
+            |/src/jvm/kotlin/test/source.kt
+            |package test
+            |
+            |fun ring(abc: String): String {
+            |    return "My precious " + abc
+            |}
+            |/src/native/kotlin/test/source.kt
+            |package test
+            |
+            |fun ring(abc: String): String {
+            |    return "My precious " + abc
+            |}
+            |/src/common/kotlin/test/source.kt
+            |package test
+            |
+            |fun ring(abc: String): String {
+            |    return "My precious " + abc
+            |}
+            |/src/js/kotlin/test/source.kt
+            |package test
+            |
+            |fun ring(abc: String): String {
+            |    return "My precious " + abc
+            |}
+            |/src/wasm/kotlin/test/source.kt
             |package test
             |
             |fun ring(abc: String): String {
@@ -185,33 +210,69 @@ class SinceKotlinTest : AbstractRenderingTest() {
         val configuration =   dokkaConfiguration {
             sourceSets {
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/jvm/")
                     classpath = listOfNotNull(jvmStdlibPath)
                     analysisPlatform = "jvm"
                 }
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/native/")
                     analysisPlatform = "native"
+                    name = "native"
                 }
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/common/")
                     classpath = listOfNotNull(commonStdlibPath)
                     analysisPlatform = "common"
+                    name = "common"
                 }
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/js/")
                     classpath = listOfNotNull(jsStdlibPath)
                     analysisPlatform = "js"
+                    name = "js"
                 }
                 sourceSet {
-                    sourceRoots = listOf("src/")
+                    sourceRoots = listOf("src/wasm/")
                     analysisPlatform = "wasm"
+                    name = "wasm"
                 }
             }
         }
         testInline(
             """
-            |/src/main/kotlin/test/source.kt
+            |/src/jvm/kotlin/test/source.kt
+            |package test
+            |
+            |/** dssdd */
+            |@SinceKotlin("1.3")
+            |fun ring(abc: String): String {
+            |    return "My precious " + abc
+            |}            
+            |/src/native/kotlin/test/source.kt
+            |package test
+            |
+            |/** dssdd */
+            |@SinceKotlin("1.3")
+            |fun ring(abc: String): String {
+            |    return "My precious " + abc
+            |}
+            |/src/common/kotlin/test/source.kt
+            |package test
+            |
+            |/** dssdd */
+            |@SinceKotlin("1.3")
+            |fun ring(abc: String): String {
+            |    return "My precious " + abc
+            |}            
+            |/src/js/kotlin/test/source.kt
+            |package test
+            |
+            |/** dssdd */
+            |@SinceKotlin("1.3")
+            |fun ring(abc: String): String {
+            |    return "My precious " + abc
+            |}            
+            |/src/wasm/kotlin/test/source.kt
             |package test
             |
             |/** dssdd */
@@ -238,7 +299,7 @@ class SinceKotlinTest : AbstractRenderingTest() {
                             find { it.sourceSets.first().analysisPlatform == i.key }?.documentation?.values?.first()
                                 ?.dfs { it is CustomTagWrapper && it.name == "Since Kotlin" }
                                 .assertNotNull("SinceKotlin[${i.key}]")
-                        assertEquals((tag.children.first() as Text).body, i.value.toString())
+                        assertEquals(i.value.toString(), (tag.children.first() as Text).body , "Platform ${i.key}")
                     }
                 }
             }
diff --git a/plugins/base/src/test/kotlin/content/exceptions/ContentForExceptions.kt b/plugins/base/src/test/kotlin/content/exceptions/ContentForExceptions.kt
index 9a9fd8b4..22becb93 100644
--- a/plugins/base/src/test/kotlin/content/exceptions/ContentForExceptions.kt
+++ b/plugins/base/src/test/kotlin/content/exceptions/ContentForExceptions.kt
@@ -339,7 +339,6 @@ class ContentForExceptions : BaseAbstractTest() {
         }
     }
 
-    @OnlyDescriptorsMPP("Return type for native `function` should be null rather than kotlin/Unit")
     @Test
     fun `throws in merged functions`() {
         testInline(
diff --git a/plugins/base/src/test/kotlin/content/seealso/ContentForSeeAlsoTest.kt b/plugins/base/src/test/kotlin/content/seealso/ContentForSeeAlsoTest.kt
index 2790d8b6..fb72178b 100644
--- a/plugins/base/src/test/kotlin/content/seealso/ContentForSeeAlsoTest.kt
+++ b/plugins/base/src/test/kotlin/content/seealso/ContentForSeeAlsoTest.kt
@@ -759,7 +759,6 @@ class ContentForSeeAlsoTest : BaseAbstractTest() {
         }
     }
 
-    @OnlyDescriptorsMPP
     @Test
     fun `multiplatform class with seealso in few platforms`() {
         testInline(
diff --git a/plugins/base/src/test/kotlin/linkableContent/LinkableContentTest.kt b/plugins/base/src/test/kotlin/linkableContent/LinkableContentTest.kt
index 504a1ebe..1b73ffee 100644
--- a/plugins/base/src/test/kotlin/linkableContent/LinkableContentTest.kt
+++ b/plugins/base/src/test/kotlin/linkableContent/LinkableContentTest.kt
@@ -20,12 +20,11 @@ import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertNotEquals
 import kotlin.test.assertNotNull
-import utils.OnlyDescriptors
 import utils.OnlyDescriptorsMPP
 
 class LinkableContentTest : BaseAbstractTest() {
 
-    @OnlyDescriptorsMPP
+    @OnlyDescriptorsMPP("#3238")
     @Test
     fun `Include module and package documentation`() {
 
@@ -151,7 +150,7 @@ class LinkableContentTest : BaseAbstractTest() {
         }
     }
 
-    @OnlyDescriptorsMPP
+    @OnlyDescriptorsMPP("#3238")
     @Test
     fun `Samples multiplatform documentation`() {
 
@@ -286,7 +285,7 @@ class LinkableContentTest : BaseAbstractTest() {
             sourceSets {
                 sourceSet {
                     analysisPlatform = "js"
-                    sourceRoots = listOf("jsMain", "commonMain", "jvmAndJsSecondCommonMain").map {
+                    sourceRoots = listOf("jsMain").map {
                         Paths.get("$testDataDir/$it/kotlin").toString()
                     }
                     name = "js"
@@ -294,7 +293,7 @@ class LinkableContentTest : BaseAbstractTest() {
                 }
                 sourceSet {
                     analysisPlatform = "jvm"
-                    sourceRoots = listOf("jvmMain", "commonMain", "jvmAndJsSecondCommonMain").map {
+                    sourceRoots = listOf("jvmMain").map {
                         Paths.get("$testDataDir/$it/kotlin").toString()
                     }
                     name = "jvm"
diff --git a/plugins/base/src/test/kotlin/model/InheritorsTest.kt b/plugins/base/src/test/kotlin/model/InheritorsTest.kt
index 5c8e7c74..459dd9ac 100644
--- a/plugins/base/src/test/kotlin/model/InheritorsTest.kt
+++ b/plugins/base/src/test/kotlin/model/InheritorsTest.kt
@@ -60,15 +60,22 @@ class InheritorsTest : AbstractModelTest("/src/main/kotlin/inheritors/Test.kt",
     fun multiplatform() {
         val configuration = dokkaConfiguration {
             sourceSets {
+                val commonSourceSet = sourceSet {
+                    name = "common"
+                    sourceRoots = listOf("common/src/")
+                    analysisPlatform = "common"
+                }
                 sourceSet {
                     name = "jvm"
-                    sourceRoots = listOf("common/src/", "jvm/src/")
+                    sourceRoots = listOf("jvm/src/")
                     analysisPlatform = "jvm"
+                    dependentSourceSets =  setOf(commonSourceSet.value.sourceSetID)
                 }
                 sourceSet {
                     name = "js"
-                    sourceRoots = listOf("common/src/", "js/src/")
+                    sourceRoots = listOf("js/src/")
                     analysisPlatform = "js"
+                    dependentSourceSets =  setOf(commonSourceSet.value.sourceSetID)
                 }
             }
         }
diff --git a/plugins/base/src/test/kotlin/signatures/DivergentSignatureTest.kt b/plugins/base/src/test/kotlin/signatures/DivergentSignatureTest.kt
index 45a5523f..509dd6e7 100644
--- a/plugins/base/src/test/kotlin/signatures/DivergentSignatureTest.kt
+++ b/plugins/base/src/test/kotlin/signatures/DivergentSignatureTest.kt
@@ -7,10 +7,8 @@ package signatures
 import utils.TestOutputWriterPlugin
 import kotlin.test.Test
 import kotlin.test.assertEquals
-import utils.OnlyDescriptors
-import utils.OnlyDescriptorsMPP
 
-@OnlyDescriptorsMPP
+
 class DivergentSignatureTest : AbstractRenderingTest() {
 
     @Test
diff --git a/plugins/base/src/test/kotlin/signatures/SignatureTest.kt b/plugins/base/src/test/kotlin/signatures/SignatureTest.kt
index 37efd6df..80a043fe 100644
--- a/plugins/base/src/test/kotlin/signatures/SignatureTest.kt
+++ b/plugins/base/src/test/kotlin/signatures/SignatureTest.kt
@@ -555,7 +555,6 @@ class SignatureTest : BaseAbstractTest() {
             }
         }
     }
-    @OnlyDescriptorsMPP
     @Test
     fun `actual typealias should have generic parameters and fully qualified name of the expansion type`() {
         val writerPlugin = TestOutputWriterPlugin()
@@ -590,7 +589,6 @@ class SignatureTest : BaseAbstractTest() {
         }
     }
 
-    @OnlyDescriptorsMPP
     @Test
     fun `type with an actual typealias`() {
         val writerPlugin = TestOutputWriterPlugin()
diff --git a/plugins/base/src/test/kotlin/transformers/ModuleAndPackageDocumentationTransformerFunctionalTest.kt b/plugins/base/src/test/kotlin/transformers/ModuleAndPackageDocumentationTransformerFunctionalTest.kt
index 85ba5ba4..54f0120a 100644
--- a/plugins/base/src/test/kotlin/transformers/ModuleAndPackageDocumentationTransformerFunctionalTest.kt
+++ b/plugins/base/src/test/kotlin/transformers/ModuleAndPackageDocumentationTransformerFunctionalTest.kt
@@ -15,7 +15,7 @@ import kotlin.test.assertEquals
 
 class ModuleAndPackageDocumentationTransformerFunctionalTest : BaseAbstractTest() {
 
-    @OnlyDescriptorsMPP
+    @OnlyDescriptorsMPP("#3238")
     @Test
     fun `multiplatform project`(@TempDir tempDir: Path) {
         val include = tempDir.resolve("include.md").toFile()
diff --git a/plugins/base/src/test/kotlin/transformers/SourceLinkTransformerTest.kt b/plugins/base/src/test/kotlin/transformers/SourceLinkTransformerTest.kt
index c23a0885..87424120 100644
--- a/plugins/base/src/test/kotlin/transformers/SourceLinkTransformerTest.kt
+++ b/plugins/base/src/test/kotlin/transformers/SourceLinkTransformerTest.kt
@@ -9,7 +9,6 @@ import org.jetbrains.dokka.SourceLinkDefinitionImpl
 import org.jetbrains.dokka.base.testApi.testRunner.BaseAbstractTest
 import org.jsoup.nodes.Element
 import signatures.renderedContent
-import utils.OnlyDescriptorsMPP
 import utils.TestOutputWriterPlugin
 import java.net.URL
 import kotlin.test.Test
@@ -71,7 +70,6 @@ class SourceLinkTransformerTest : BaseAbstractTest() {
         }
     }
 
-    @OnlyDescriptorsMPP
     @Test
     fun `source link should be for actual typealias`() {
         val mppConfiguration = dokkaConfiguration {
diff --git a/subprojects/analysis-kotlin-symbols/build.gradle.kts b/subprojects/analysis-kotlin-symbols/build.gradle.kts
index 6733b778..b6626f92 100644
--- a/subprojects/analysis-kotlin-symbols/build.gradle.kts
+++ b/subprojects/analysis-kotlin-symbols/build.gradle.kts
@@ -72,6 +72,9 @@ dependencies {
             isTransitive = false // see KTIJ-19820
         }
     }
+    // copy-pasted from Analysis API https://github.com/JetBrains/kotlin/blob/a10042f9099e20a656dec3ecf1665eea340a3633/analysis/low-level-api-fir/build.gradle.kts#L37
+    runtimeOnly("com.github.ben-manes.caffeine:caffeine:2.9.3")
+
     runtimeOnly(libs.kotlinx.collections.immutable)
     implementation(libs.kotlin.compiler.k2) {
         isTransitive = false
@@ -79,6 +82,7 @@ dependencies {
 
     // TODO [beresnev] get rid of it
     compileOnly(libs.kotlinx.coroutines.core)
+
 }
 
 tasks {
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))