Implement ModuleAndPackageDocumentationReader

12 files changed, 277 insertions, 197 deletions

diff --git a/plugins/base/src/main/kotlin/DokkaBase.kt b/plugins/base/src/main/kotlin/DokkaBase.kt
index e74bc6d9..44a532f8 100644
--- a/plugins/base/src/main/kotlin/DokkaBase.kt
+++ b/plugins/base/src/main/kotlin/DokkaBase.kt
@@ -81,7 +81,9 @@ class DokkaBase : DokkaPlugin() {
 
     val modulesAndPackagesDocumentation by extending {
         CoreExtensions.preMergeDocumentableTransformer providing { ctx ->
-            ModuleAndPackageDocumentationTransformer(ctx, ctx.single(kotlinAnalysis))
+            ModuleAndPackageDocumentationTransformer(
+                ModuleAndPackageDocumentationReader(ctx, ctx.single(kotlinAnalysis))
+            )
         }
     }
 
@@ -127,7 +129,7 @@ class DokkaBase : DokkaPlugin() {
     }
 
     val sourceSetMerger by extending {
-        CoreExtensions.pageTransformer providing  ::SourceSetMergingPageTransformer
+        CoreExtensions.pageTransformer providing ::SourceSetMergingPageTransformer
     }
 
     val fallbackMerger by extending {
diff --git a/plugins/base/src/main/kotlin/parsers/ModuleAndPackageDocumentationParser.kt b/plugins/base/src/main/kotlin/parsers/ModuleAndPackageDocumentationParser.kt
deleted file mode 100644
index 4b46537b..00000000
--- a/plugins/base/src/main/kotlin/parsers/ModuleAndPackageDocumentationParser.kt
+++ /dev/null
@@ -1,76 +0,0 @@
-package org.jetbrains.dokka.base.parsers
-
-import org.jetbrains.dokka.DokkaException
-import java.io.File
-
-internal class IllegalModuleAndPackageDocumentation(
-    source: ModuleAndPackageDocumentationSource, message: String
-) : DokkaException("[$source] $message")
-
-data class ModuleAndPackageDocFragment(
-    val classifier: Classifier,
-    val name: String,
-    val documentation: String
-) {
-    enum class Classifier { Module, Package }
-}
-
-internal abstract class ModuleAndPackageDocumentationSource {
-    abstract val sourceDescription: String
-    abstract val documentation: String
-
-    override fun toString(): String {
-        return sourceDescription
-    }
-}
-
-internal class ModuleAndPackageDocumentationFile(private val file: File) : ModuleAndPackageDocumentationSource() {
-    override val sourceDescription: String = file.path
-    override val documentation: String by lazy(LazyThreadSafetyMode.PUBLICATION) { file.readText() }
-}
-
-internal fun parseModuleAndPackageDocFragments(source: File): List<ModuleAndPackageDocFragment> {
-    return parseModuleAndPackageDocFragments(ModuleAndPackageDocumentationFile(source))
-}
-
-internal fun parseModuleAndPackageDocFragments(source: ModuleAndPackageDocumentationSource): List<ModuleAndPackageDocFragment> {
-    val fragmentStrings = source.documentation.split(Regex("(|^)#\\s*(?=(Module|Package))"))
-    return fragmentStrings
-        .filter(String::isNotBlank)
-        .map { fragmentString -> parseModuleAndPackageDocFragment(source, fragmentString) }
-}
-
-private fun parseModuleAndPackageDocFragment(
-    source: ModuleAndPackageDocumentationSource,
-    fragment: String
-): ModuleAndPackageDocFragment {
-    val firstLineAndDocumentation = fragment.split("\r\n", "\n", "\r", limit = 2)
-    val firstLine = firstLineAndDocumentation[0]
-
-    val classifierAndName = firstLine.split(Regex("\\s+"), limit = 2)
-    if (classifierAndName.size != 2) {
-        throw IllegalModuleAndPackageDocumentation(source, "Missing ${classifierAndName.first()} name")
-    }
-
-    val classifier = when (classifierAndName[0].trim()) {
-        "Module" -> ModuleAndPackageDocFragment.Classifier.Module
-        "Package" -> ModuleAndPackageDocFragment.Classifier.Package
-        else -> throw IllegalStateException("Unexpected classifier ${classifierAndName[0]}")
-    }
-
-    val name = classifierAndName[1].trim()
-    if (name.contains(Regex("\\s"))) {
-        throw IllegalModuleAndPackageDocumentation(
-            source, "Module/Package name cannot contain whitespace in '$firstLine'"
-        )
-    }
-
-    return ModuleAndPackageDocFragment(
-        classifier = classifier,
-        name = name,
-        documentation = firstLineAndDocumentation.getOrNull(1)?.trim().orEmpty()
-    )
-}
-
-
-
diff --git a/plugins/base/src/main/kotlin/parsers/moduleAndPackage/IllegalModuleAndPackageDocumentation.kt b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/IllegalModuleAndPackageDocumentation.kt
new file mode 100644
index 00000000..f642c374
--- /dev/null
+++ b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/IllegalModuleAndPackageDocumentation.kt
@@ -0,0 +1,7 @@
+package org.jetbrains.dokka.base.parsers.moduleAndPackage
+
+import org.jetbrains.dokka.DokkaException
+
+internal class IllegalModuleAndPackageDocumentation(
+    source: ModuleAndPackageDocumentationSource, message: String
+) : DokkaException("[$source] $message")
diff --git a/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentation.kt b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentation.kt
new file mode 100644
index 00000000..5139c872
--- /dev/null
+++ b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentation.kt
@@ -0,0 +1,11 @@
+package org.jetbrains.dokka.base.parsers.moduleAndPackage
+
+import org.jetbrains.dokka.model.doc.DocumentationNode
+
+internal data class ModuleAndPackageDocumentation(
+    val name: String,
+    val classifier: Classifier,
+    val documentation: DocumentationNode
+) {
+    enum class Classifier { Module, Package }
+}
diff --git a/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationFragment.kt b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationFragment.kt
new file mode 100644
index 00000000..fa99a8e2
--- /dev/null
+++ b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationFragment.kt
@@ -0,0 +1,10 @@
+package org.jetbrains.dokka.base.parsers.moduleAndPackage
+
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.ModuleAndPackageDocumentation.*
+
+internal data class ModuleAndPackageDocumentationFragment(
+    val name: String,
+    val classifier: Classifier,
+    val documentation: String,
+    val source: ModuleAndPackageDocumentationSource
+)
diff --git a/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationParsingContext.kt b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationParsingContext.kt
new file mode 100644
index 00000000..f11a5d13
--- /dev/null
+++ b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationParsingContext.kt
@@ -0,0 +1,36 @@
+@file:Suppress("FunctionName")
+
+package org.jetbrains.dokka.base.parsers.moduleAndPackage
+
+import org.jetbrains.dokka.analysis.DokkaResolutionFacade
+import org.jetbrains.dokka.base.parsers.MarkdownParser
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.ModuleAndPackageDocumentation.Classifier.*
+import org.jetbrains.dokka.model.doc.DocumentationNode
+import org.jetbrains.dokka.utilities.DokkaLogger
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+
+internal fun interface ModuleAndPackageDocumentationParsingContext {
+    fun markdownParserFor(fragment: ModuleAndPackageDocumentationFragment): MarkdownParser
+}
+
+internal fun ModuleAndPackageDocumentationParsingContext.parse(
+    fragment: ModuleAndPackageDocumentationFragment
+): DocumentationNode {
+    return markdownParserFor(fragment).parse(fragment.documentation)
+}
+
+internal fun ModuleAndPackageDocumentationParsingContext(
+    logger: DokkaLogger,
+    facade: DokkaResolutionFacade? = null,
+) = ModuleAndPackageDocumentationParsingContext { fragment ->
+    val descriptor = when (fragment.classifier) {
+        Module -> facade?.moduleDescriptor?.getPackage(FqName.topLevel(Name.identifier("")))
+        Package -> facade?.moduleDescriptor?.getPackage(FqName(fragment.name))
+    }
+    MarkdownParser(
+        resolutionFacade = facade,
+        declarationDescriptor = descriptor,
+        logger = logger
+    )
+}
diff --git a/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationSource.kt b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationSource.kt
new file mode 100644
index 00000000..7c3f6d97
--- /dev/null
+++ b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/ModuleAndPackageDocumentationSource.kt
@@ -0,0 +1,17 @@
+package org.jetbrains.dokka.base.parsers.moduleAndPackage
+
+import java.io.File
+
+internal abstract class ModuleAndPackageDocumentationSource {
+    abstract val sourceDescription: String
+    abstract val documentation: String
+
+    override fun toString(): String {
+        return sourceDescription
+    }
+}
+
+internal data class ModuleAndPackageDocumentationFile(private val file: File) : ModuleAndPackageDocumentationSource() {
+    override val sourceDescription: String = file.path
+    override val documentation: String by lazy(LazyThreadSafetyMode.PUBLICATION) { file.readText() }
+}
diff --git a/plugins/base/src/main/kotlin/parsers/moduleAndPackage/parseModuleAndPackageDocumentation.kt b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/parseModuleAndPackageDocumentation.kt
new file mode 100644
index 00000000..a2876308
--- /dev/null
+++ b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/parseModuleAndPackageDocumentation.kt
@@ -0,0 +1,15 @@
+@file:Suppress("FunctionName")
+
+package org.jetbrains.dokka.base.parsers.moduleAndPackage
+
+internal fun parseModuleAndPackageDocumentation(
+    context: ModuleAndPackageDocumentationParsingContext,
+    fragment: ModuleAndPackageDocumentationFragment
+): ModuleAndPackageDocumentation {
+    return ModuleAndPackageDocumentation(
+        name = fragment.name,
+        classifier = fragment.classifier,
+        documentation = context.parse(fragment)
+    )
+}
+
diff --git a/plugins/base/src/main/kotlin/parsers/moduleAndPackage/parseModuleAndPackageDocumentationFragments.kt b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/parseModuleAndPackageDocumentationFragments.kt
new file mode 100644
index 00000000..2667681c
--- /dev/null
+++ b/plugins/base/src/main/kotlin/parsers/moduleAndPackage/parseModuleAndPackageDocumentationFragments.kt
@@ -0,0 +1,53 @@
+package org.jetbrains.dokka.base.parsers.moduleAndPackage
+
+import java.io.File
+
+
+internal fun parseModuleAndPackageDocumentationFragments(source: File): List<ModuleAndPackageDocumentationFragment> {
+    return parseModuleAndPackageDocumentationFragments(ModuleAndPackageDocumentationFile(source))
+}
+
+internal fun parseModuleAndPackageDocumentationFragments(
+    source: ModuleAndPackageDocumentationSource
+): List<ModuleAndPackageDocumentationFragment> {
+    val fragmentStrings = source.documentation.split(Regex("(|^)#\\s*(?=(Module|Package))"))
+    return fragmentStrings
+        .filter(String::isNotBlank)
+        .map { fragmentString -> parseModuleAndPackageDocFragment(source, fragmentString) }
+}
+
+private fun parseModuleAndPackageDocFragment(
+    source: ModuleAndPackageDocumentationSource,
+    fragment: String
+): ModuleAndPackageDocumentationFragment {
+    val firstLineAndDocumentation = fragment.split("\r\n", "\n", "\r", limit = 2)
+    val firstLine = firstLineAndDocumentation[0]
+
+    val classifierAndName = firstLine.split(Regex("\\s+"), limit = 2)
+    if (classifierAndName.size != 2) {
+        throw IllegalModuleAndPackageDocumentation(source, "Missing ${classifierAndName.first()} name")
+    }
+
+    val classifier = when (classifierAndName[0].trim()) {
+        "Module" -> ModuleAndPackageDocumentation.Classifier.Module
+        "Package" -> ModuleAndPackageDocumentation.Classifier.Package
+        else -> throw IllegalStateException("Unexpected classifier ${classifierAndName[0]}")
+    }
+
+    val name = classifierAndName[1].trim()
+    if (name.contains(Regex("\\s"))) {
+        throw IllegalModuleAndPackageDocumentation(
+            source, "Module/Package name cannot contain whitespace in '$firstLine'"
+        )
+    }
+
+    return ModuleAndPackageDocumentationFragment(
+        name = name,
+        classifier = classifier,
+        documentation = firstLineAndDocumentation.getOrNull(1)?.trim().orEmpty(),
+        source = source
+    )
+}
+
+
+
diff --git a/plugins/base/src/main/kotlin/transformers/documentables/ModuleAndPackageDocumentationReader.kt b/plugins/base/src/main/kotlin/transformers/documentables/ModuleAndPackageDocumentationReader.kt
new file mode 100644
index 00000000..e8297f3f
--- /dev/null
+++ b/plugins/base/src/main/kotlin/transformers/documentables/ModuleAndPackageDocumentationReader.kt
@@ -0,0 +1,72 @@
+@file:Suppress("FunctionName")
+
+package org.jetbrains.dokka.base.transformers.documentables
+
+import org.jetbrains.dokka.DokkaConfiguration
+import org.jetbrains.dokka.analysis.KotlinAnalysis
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.ModuleAndPackageDocumentation.Classifier
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.ModuleAndPackageDocumentationFragment
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.ModuleAndPackageDocumentationParsingContext
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.parseModuleAndPackageDocumentation
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.parseModuleAndPackageDocumentationFragments
+import org.jetbrains.dokka.model.DModule
+import org.jetbrains.dokka.model.DPackage
+import org.jetbrains.dokka.model.SourceSetDependent
+import org.jetbrains.dokka.model.doc.DocumentationNode
+import org.jetbrains.dokka.plugability.DokkaContext
+import org.jetbrains.dokka.utilities.associateWithNotNull
+
+internal interface ModuleAndPackageDocumentationReader {
+    operator fun get(module: DModule): SourceSetDependent<DocumentationNode>
+    operator fun get(pkg: DPackage): SourceSetDependent<DocumentationNode>
+}
+
+// TODO NOW: Test
+internal fun ModuleAndPackageDocumentationReader(
+    context: DokkaContext, kotlinAnalysis: KotlinAnalysis? = null
+): ModuleAndPackageDocumentationReader = ContextModuleAndPackageDocumentationReader(context, kotlinAnalysis)
+
+private class ContextModuleAndPackageDocumentationReader(
+    private val context: DokkaContext,
+    private val kotlinAnalysis: KotlinAnalysis?
+) : ModuleAndPackageDocumentationReader {
+
+    private val documentationFragments: SourceSetDependent<List<ModuleAndPackageDocumentationFragment>> =
+        context.configuration.sourceSets.associateWith { sourceSet ->
+            sourceSet.includes.flatMap { include -> parseModuleAndPackageDocumentationFragments(include) }
+        }
+
+    private fun findDocumentationNodes(
+        sourceSets: Set<DokkaConfiguration.DokkaSourceSet>,
+        predicate: (ModuleAndPackageDocumentationFragment) -> Boolean
+    ): SourceSetDependent<DocumentationNode> {
+        return sourceSets.associateWithNotNull { sourceSet ->
+            val fragments = documentationFragments[sourceSet].orEmpty().filter(predicate)
+            val resolutionFacade = kotlinAnalysis?.get(sourceSet)?.facade
+            val documentations = fragments.map { fragment ->
+                parseModuleAndPackageDocumentation(
+                    context = ModuleAndPackageDocumentationParsingContext(context.logger, resolutionFacade),
+                    fragment = fragment
+                )
+            }
+            when (documentations.size) {
+                0 -> null
+                1 -> documentations.single().documentation
+                else -> DocumentationNode(documentations.flatMap { it.documentation.children })
+            }
+        }
+    }
+
+    override fun get(module: DModule): SourceSetDependent<DocumentationNode> {
+        return findDocumentationNodes(module.sourceSets) { fragment ->
+            fragment.classifier == Classifier.Module && fragment.name == module.name
+        }
+    }
+
+    override fun get(pkg: DPackage): SourceSetDependent<DocumentationNode> {
+        return findDocumentationNodes(pkg.sourceSets) { fragment ->
+            // TODO NOW: handle JS Root thing
+            fragment.classifier == Classifier.Package && fragment.name == pkg.dri.packageName
+        }
+    }
+}
diff --git a/plugins/base/src/main/kotlin/transformers/documentables/ModuleAndPackageDocumentationTransformer.kt b/plugins/base/src/main/kotlin/transformers/documentables/ModuleAndPackageDocumentationTransformer.kt
index 5f1a540d..9bdec75a 100644
--- a/plugins/base/src/main/kotlin/transformers/documentables/ModuleAndPackageDocumentationTransformer.kt
+++ b/plugins/base/src/main/kotlin/transformers/documentables/ModuleAndPackageDocumentationTransformer.kt
@@ -1,101 +1,32 @@
 package org.jetbrains.dokka.base.transformers.documentables
 
-import org.jetbrains.dokka.analysis.KotlinAnalysis
+import org.jetbrains.dokka.DokkaConfiguration.DokkaSourceSet
 import org.jetbrains.dokka.model.DModule
 import org.jetbrains.dokka.model.doc.DocumentationNode
-import org.jetbrains.dokka.DokkaConfiguration.DokkaSourceSet
-import org.jetbrains.dokka.base.parsers.MarkdownParser
-import org.jetbrains.dokka.plugability.DokkaContext
 import org.jetbrains.dokka.transformers.documentation.PreMergeDocumentableTransformer
-import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
-import java.nio.file.Files
-import java.nio.file.Paths
-
 
+// TODO NOW: Test
 internal class ModuleAndPackageDocumentationTransformer(
-    private val context: DokkaContext,
-    private val kotlinAnalysis: KotlinAnalysis
+    private val moduleAndPackageDocumentationReader: ModuleAndPackageDocumentationReader
 ) : PreMergeDocumentableTransformer {
 
     override fun invoke(modules: List<DModule>): List<DModule> {
-
-        val modulesAndPackagesDocumentation =
-            context.configuration.sourceSets
-                .map { sourceSet ->
-                    Pair(sourceSet.moduleDisplayName, sourceSet) to
-                            sourceSet.includes.map { it.toPath() }
-                                .also {
-                                    it.forEach {
-                                        if (Files.notExists(it))
-                                            context.logger.warn("Not found file under this path ${it.toAbsolutePath()}")
-                                    }
-                                }
-                                .filter { Files.exists(it) }
-                                .flatMap {
-                                    it.toFile()
-                                        .readText()
-                                        .split(Regex("(\n|^)# (?=(Module|Package))")) // Matches heading with Module/Package to split by
-                                        .filter { it.isNotEmpty() }
-                                        .map { it.split(Regex(" "), 2) } // Matches space between Module/Package and fully qualified name
-                                }.groupBy({ it[0] }, {
-                                    it[1].split(Regex("\n"), 2) // Matches new line after fully qualified name
-                                        .let { it[0].trim() to it[1].trim() }
-                                }).mapValues {
-                                    it.value.toMap()
-                                }
-                }.toMap()
-
         return modules.map { module ->
-
-            val moduleDocumentation =
-                module.sourceSets.mapNotNull { pd ->
-                    val doc = modulesAndPackagesDocumentation[Pair(module.name, pd)]
-                    val facade = kotlinAnalysis[pd].facade
-                    try {
-                        doc?.get("Module")?.get(module.name)?.run {
-                            pd to MarkdownParser(
-                                facade,
-                                facade.moduleDescriptor.getPackage(FqName.topLevel(Name.identifier(""))),
-                                context.logger
-                            ).parse(this)
-                        }
-                    } catch (e: IllegalArgumentException) {
-                        context.logger.error(e.message.orEmpty())
-                        null
-                    }
-                }.toMap()
-
-            val packagesDocumentation = module.packages.map {
-                it.name to it.sourceSets.mapNotNull { pd ->
-                    val doc = modulesAndPackagesDocumentation[Pair(module.name, pd)]
-                    val facade = kotlinAnalysis[pd].facade
-                    val descriptor = facade.moduleDescriptor.getPackage(FqName(it.name.let { if(it == "[JS root]") "" else it }))
-                    doc?.get("Package")?.get(it.name)?.run {
-                        pd to MarkdownParser(
-                            facade,
-                            descriptor,
-                            context.logger
-                        ).parse(this)
-                    }
-                }.toMap()
-            }.toMap()
-
             module.copy(
-                documentation = mergeDocumentation(module.documentation, moduleDocumentation),
-                packages = module.packages.map {
-                    val packageDocumentation = packagesDocumentation[it.name]
-                    if (packageDocumentation != null && packageDocumentation.isNotEmpty())
-                        it.copy(documentation = mergeDocumentation(it.documentation, packageDocumentation))
-                    else
-                        it
+                documentation = module.documentation + moduleAndPackageDocumentationReader[module],
+                packages = module.packages.map { pkg ->
+                    pkg.copy(
+                        documentation = pkg.documentation + moduleAndPackageDocumentationReader[pkg]
+                    )
                 }
             )
         }
     }
 
-    private fun mergeDocumentation(origin: Map<DokkaSourceSet, DocumentationNode>, new: Map<DokkaSourceSet, DocumentationNode>): Map<DokkaSourceSet, DocumentationNode> =
-        (origin.asSequence() + new.asSequence())
+    private operator fun Map<DokkaSourceSet, DocumentationNode>.plus(
+        other: Map<DokkaSourceSet, DocumentationNode>
+    ): Map<DokkaSourceSet, DocumentationNode> =
+        (asSequence() + other.asSequence())
             .distinct()
             .groupBy({ it.key }, { it.value })
             .mapValues { (_, values) -> DocumentationNode(values.flatMap { it.children }) }
diff --git a/plugins/base/src/test/kotlin/parsers/ParseModuleAndPackageDocFragmentsTest.kt b/plugins/base/src/test/kotlin/parsers/ParseModuleAndPackageDocumentationFragmentsTest.kt
index 49eddaea..5c0a0646 100644
--- a/plugins/base/src/test/kotlin/parsers/ParseModuleAndPackageDocFragmentsTest.kt
+++ b/plugins/base/src/test/kotlin/parsers/ParseModuleAndPackageDocumentationFragmentsTest.kt
@@ -1,11 +1,8 @@
 package parsers
 
-import org.jetbrains.dokka.base.parsers.IllegalModuleAndPackageDocumentation
-import org.jetbrains.dokka.base.parsers.ModuleAndPackageDocFragment
-import org.jetbrains.dokka.base.parsers.ModuleAndPackageDocFragment.Classifier.Module
-import org.jetbrains.dokka.base.parsers.ModuleAndPackageDocFragment.Classifier.Package
-import org.jetbrains.dokka.base.parsers.ModuleAndPackageDocumentationSource
-import org.jetbrains.dokka.base.parsers.parseModuleAndPackageDocFragments
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.*
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.ModuleAndPackageDocumentation.Classifier.Module
+import org.jetbrains.dokka.base.parsers.moduleAndPackage.ModuleAndPackageDocumentation.Classifier.Package
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Assertions.assertTrue
 import org.junit.jupiter.api.Test
@@ -13,14 +10,13 @@ import org.junit.jupiter.api.assertThrows
 import org.junit.jupiter.api.io.TempDir
 import java.nio.file.Path
 
-class ParseModuleAndPackageDocFragmentsTest {
+class ParseModuleAndPackageDocumentationFragmentsTest {
 
     @Test
     fun `basic example`() {
 
-        val fragments = parseModuleAndPackageDocFragments(
-            source(
-                """
+        val source = source(
+            """
                 # Module kotlin-demo
                 Module description
         
@@ -32,25 +28,28 @@ class ParseModuleAndPackageDocFragmentsTest {
                 # Package org.jetbrains.kotlin.demo2
                 Package demo2 description
                 """.trimIndent()
-            )
         )
+        val fragments = parseModuleAndPackageDocumentationFragments(source)
 
         assertEquals(
             listOf(
-                ModuleAndPackageDocFragment(
+                ModuleAndPackageDocumentationFragment(
                     classifier = Module,
                     name = "kotlin-demo",
-                    documentation = "Module description"
+                    documentation = "Module description",
+                    source = source
                 ),
-                ModuleAndPackageDocFragment(
+                ModuleAndPackageDocumentationFragment(
                     classifier = Package,
                     name = "org.jetbrains.kotlin.demo",
-                    documentation = "Package demo description\n## Level 2 heading\nHeading 2"
+                    documentation = "Package demo description\n## Level 2 heading\nHeading 2",
+                    source = source
                 ),
-                ModuleAndPackageDocFragment(
+                ModuleAndPackageDocumentationFragment(
                     classifier = Package,
                     name = "org.jetbrains.kotlin.demo2",
-                    documentation = "Package demo2 description"
+                    documentation = "Package demo2 description",
+                    source = source
                 )
             ),
             fragments
@@ -60,7 +59,7 @@ class ParseModuleAndPackageDocFragmentsTest {
     @Test
     fun `no module name specified fails`() {
         val exception = assertThrows<IllegalModuleAndPackageDocumentation> {
-            parseModuleAndPackageDocFragments(
+            parseModuleAndPackageDocumentationFragments(
                 source(
                     """
                     # Module
@@ -79,7 +78,7 @@ class ParseModuleAndPackageDocFragmentsTest {
     @Test
     fun `no package name specified fails`() {
         val exception = assertThrows<IllegalModuleAndPackageDocumentation> {
-            parseModuleAndPackageDocFragments(
+            parseModuleAndPackageDocumentationFragments(
                 source(
                     """
                     # Package
@@ -98,7 +97,7 @@ class ParseModuleAndPackageDocFragmentsTest {
     @Test
     fun `white space in module name fails`() {
         val exception = assertThrows<IllegalModuleAndPackageDocumentation> {
-            parseModuleAndPackageDocFragments(
+            parseModuleAndPackageDocumentationFragments(
                 source(
                     """
                     # Module My Module
@@ -116,7 +115,7 @@ class ParseModuleAndPackageDocFragmentsTest {
     @Test
     fun `white space in package name fails`() {
         val exception = assertThrows<IllegalModuleAndPackageDocumentation> {
-            parseModuleAndPackageDocFragments(
+            parseModuleAndPackageDocumentationFragments(
                 source(
                     """
                     # Package my package
@@ -133,28 +132,29 @@ class ParseModuleAndPackageDocFragmentsTest {
 
     @Test
     fun `multiple whitespaces are supported in first line`() {
-        val fragments = parseModuleAndPackageDocFragments(
-            source(
-                """
-                    #    Module my-module
-                    My Module
-                    #   Package com.my.package
-                    My Package
+        val source = source(
+            """
+            #    Module my-module
+            My Module
+            #   Package com.my.package
+            My Package
                 """.trimIndent()
-            )
         )
+        val fragments = parseModuleAndPackageDocumentationFragments(source)
 
         assertEquals(
             listOf(
-                ModuleAndPackageDocFragment(
+                ModuleAndPackageDocumentationFragment(
                     classifier = Module,
                     name = "my-module",
-                    documentation = "My Module"
+                    documentation = "My Module",
+                    source = source
                 ),
-                ModuleAndPackageDocFragment(
+                ModuleAndPackageDocumentationFragment(
                     classifier = Package,
                     name = "com.my.package",
-                    documentation = "My Package"
+                    documentation = "My Package",
+                    source = source
                 )
             ),
             fragments
@@ -175,18 +175,20 @@ class ParseModuleAndPackageDocFragmentsTest {
 
         assertEquals(
             listOf(
-                ModuleAndPackageDocFragment(
+                ModuleAndPackageDocumentationFragment(
                     classifier = Module,
                     name = "MyModule",
-                    documentation = "D1"
+                    documentation = "D1",
+                    source = ModuleAndPackageDocumentationFile(file)
                 ),
-                ModuleAndPackageDocFragment(
+                ModuleAndPackageDocumentationFragment(
                     classifier = Package,
                     name = "com.sample",
-                    documentation = "D2"
+                    documentation = "D2",
+                    source = ModuleAndPackageDocumentationFile(file)
                 )
             ),
-            parseModuleAndPackageDocFragments(file)
+            parseModuleAndPackageDocumentationFragments(file)
         )
     }