1 files changed, 65 insertions, 40 deletions
diff --git a/plugins/base/src/main/kotlin/transformers/pages/sourcelinks/SourceLinksTransformer.kt b/plugins/base/src/main/kotlin/transformers/pages/sourcelinks/SourceLinksTransformer.kt
index bf411a1f..58e52c43 100644
--- a/plugins/base/src/main/kotlin/transformers/pages/sourcelinks/SourceLinksTransformer.kt
+++ b/plugins/base/src/main/kotlin/transformers/pages/sourcelinks/SourceLinksTransformer.kt
@@ -5,6 +5,8 @@ import com.intellij.psi.PsiDocumentManager
 import org.jetbrains.dokka.DokkaConfiguration
 import org.jetbrains.dokka.base.translators.documentables.PageContentBuilder
 import org.jetbrains.dokka.model.DescriptorDocumentableSource
+import org.jetbrains.dokka.model.DocumentableSource
+import org.jetbrains.dokka.model.PsiDocumentableSource
 import org.jetbrains.dokka.model.WithExpectActual
 import org.jetbrains.dokka.pages.*
 import org.jetbrains.dokka.plugability.DokkaContext
@@ -15,52 +17,75 @@ import org.jetbrains.kotlin.utils.addToStdlib.cast
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class SourceLinksTransformer(val context: DokkaContext, val builder: PageContentBuilder) : PageTransformer {
+    private val sourceLinks = getSourceLinks()
 
-    override fun invoke(input: RootPageNode): RootPageNode {
+    override fun invoke(input: RootPageNode) =
+        input.transformContentPagesTree { node ->
+            when(val documentable = node.documentable){
+                is WithExpectActual -> resolveSources(documentable)
+                    .takeIf{ it.isNotEmpty() }
+                    ?.let { node.addSourcesContent(it) }
+                    ?: node
+                else -> node
+            }
+        }
 
-        val sourceLinks = context.configuration.passesConfigurations
-            .flatMap { it.sourceLinks.map { sl -> SourceLink(sl, it.platformData) } }
+    private fun getSourceLinks() = context.configuration.passesConfigurations
+        .flatMap { it.sourceLinks.map { sl -> SourceLink(sl, it.platformData) } }
 
-        return input.transformContentPagesTree { node ->
-            node.documentable.safeAs<WithExpectActual>()?.sources?.map?.entries?.let { entries ->
-                val resolvedSources = entries.mapNotNull { entry ->
-                    sourceLinks.find { entry.value.path.contains(it.path) && it.platformData == entry.key }?.let {
-                        Pair(
-                            entry.key,
-                            entry.value.cast<DescriptorDocumentableSource>().toLink(it)
-                        )
-                    }
-                }
-                if (resolvedSources.isNotEmpty()) {
-                    val table = builder.contentFor(node.dri.first(), node.documentable!!.platformData.toSet()) {
-                        header(2) { text("Sources") }
-                        +ContentTable(
-                            emptyList(),
-                            resolvedSources.map {
-                                buildGroup(node.dri.first(), setOf(it.first), ContentKind.Source) {
-                                    platformDependentHint(node.dri.first(), setOf(it.first)) {
-                                        +link("(source)", it.second, ContentKind.Source, mainPlatformData, mainStyles, mainExtra)
-                                    }
-                                }
-                            },
-                            DCI(node.dri, ContentKind.Subtypes),
-                            node.documentable!!.platformData.toSet(),
-                            style = emptySet()
-                        )
+    private fun resolveSources(documentable: WithExpectActual) = documentable.sources.map.entries
+        .mapNotNull { entry ->
+            sourceLinks.find { entry.value.path.contains(it.path) && it.platformData == entry.key }?.let {
+                Pair(
+                    entry.key,
+                    entry.value.toLink(it)
+                )
+            }
+        }
+
+    private fun ContentPage.addSourcesContent(sources: List<Pair<PlatformData, String>>) = builder
+        .buildSourcesContent(this, sources)
+        .let {
+            this.modified(
+                content = this.content.addTable(it)
+            )
+        }
+
+    private fun PageContentBuilder.buildSourcesContent(
+        node: ContentPage,
+        sources: List<Pair<PlatformData,String>>
+    ) = contentFor(
+            node.dri.first(),
+            node.documentable!!.platformData.toSet()
+        ) {
+            header(2) { text("Sources") }
+            +ContentTable(
+                emptyList(),
+                sources.map {
+                    buildGroup(node.dri.first(), setOf(it.first)) {
+                        +link("(source)", it.second)
                     }
-                    node.modified(content = node.content.addTable(table))
-                } else {
-                    node
-                }
-            } ?: node
+                },
+                DCI(node.dri, ContentKind.Source),
+                node.documentable!!.platformData.toSet(),
+                style = emptySet()
+            )
         }
-    }
 
-    private fun DescriptorDocumentableSource.toLink(sourceLink: SourceLink): String =
-        sourceLink.url +
-        this.path.split(sourceLink.path)[1] +
-        sourceLink.lineSuffix +
-        "${this.descriptor.cast<DeclarationDescriptorWithSource>().source.getPsi()?.lineNumber() ?: 1}"
+    private fun DocumentableSource.toLink(sourceLink: SourceLink): String {
+        val lineNumber = when(this){
+            is DescriptorDocumentableSource -> this.descriptor
+                .cast<DeclarationDescriptorWithSource>()
+                .source.getPsi()
+                ?.lineNumber()
+            is PsiDocumentableSource -> this.psi.lineNumber()
+            else -> null
+        }
+        return sourceLink.url +
+                this.path.split(sourceLink.path)[1] +
+                sourceLink.lineSuffix +
+                "${lineNumber ?: 1}"
+    }
 
     private fun ContentNode.addTable(table: ContentGroup): ContentNode =
         when (this) {