Extract inheritance map

3 files changed, 90 insertions, 1 deletions

diff --git a/plugins/base/src/main/kotlin/DokkaBase.kt b/plugins/base/src/main/kotlin/DokkaBase.kt
index f2adcbc1..4af12991 100644
--- a/plugins/base/src/main/kotlin/DokkaBase.kt
+++ b/plugins/base/src/main/kotlin/DokkaBase.kt
@@ -9,15 +9,16 @@ import org.jetbrains.dokka.base.resolvers.LocationProviderFactory
 import org.jetbrains.dokka.base.signatures.KotlinSignatureProvider
 import org.jetbrains.dokka.base.signatures.SignatureProvider
 import org.jetbrains.dokka.base.transformers.documentables.DefaultDocumentableMerger
+import org.jetbrains.dokka.base.transformers.documentables.InheritorsExtractorTransformer
 import org.jetbrains.dokka.base.transformers.pages.comments.CommentsToContentConverter
 import org.jetbrains.dokka.base.transformers.pages.comments.DocTagToContentConverter
 import org.jetbrains.dokka.base.transformers.pages.merger.FallbackPageMergerStrategy
 import org.jetbrains.dokka.base.transformers.pages.merger.PageMerger
 import org.jetbrains.dokka.base.transformers.pages.merger.PageMergerStrategy
 import org.jetbrains.dokka.base.transformers.pages.merger.SameMethodNamePageMergerStrategy
-import org.jetbrains.dokka.base.translators.psi.DefaultPsiToDocumentableTranslator
 import org.jetbrains.dokka.base.translators.descriptors.DefaultDescriptorToDocumentableTranslator
 import org.jetbrains.dokka.base.translators.documentables.DefaultDocumentableToPageTranslator
+import org.jetbrains.dokka.base.translators.psi.DefaultPsiToDocumentableTranslator
 import org.jetbrains.dokka.plugability.DokkaPlugin
 
 class DokkaBase : DokkaPlugin() {
@@ -45,6 +46,10 @@ class DokkaBase : DokkaPlugin() {
         }
     }
 
+    val inheritorsExtractor by extending {
+        CoreExtensions.documentableTransformer with InheritorsExtractorTransformer()
+    }
+
     val documentableToPageTranslator by extending(isFallback = true) {
         CoreExtensions.documentableToPageTranslator providing { ctx ->
             DefaultDocumentableToPageTranslator(
diff --git a/plugins/base/src/main/kotlin/transformers/documentables/InheritorsExtractorTransformer.kt b/plugins/base/src/main/kotlin/transformers/documentables/InheritorsExtractorTransformer.kt
new file mode 100644
index 00000000..e372ad9c
--- /dev/null
+++ b/plugins/base/src/main/kotlin/transformers/documentables/InheritorsExtractorTransformer.kt
@@ -0,0 +1,73 @@
+package org.jetbrains.dokka.base.transformers.documentables
+
+import org.jetbrains.dokka.links.DRI
+import org.jetbrains.dokka.model.*
+import org.jetbrains.dokka.model.properties.ExtraProperty
+import org.jetbrains.dokka.model.properties.MergeStrategy
+import org.jetbrains.dokka.pages.PlatformData
+import org.jetbrains.dokka.plugability.DokkaContext
+import org.jetbrains.dokka.transformers.documentation.DocumentableTransformer
+
+class InheritorsExtractorTransformer : DocumentableTransformer {
+    override fun invoke(original: DModule, context: DokkaContext): DModule =
+        original.generateInheritanceMap().let { inheritanceMap -> original.appendInheritors(inheritanceMap) as DModule }
+
+    private fun <T : Documentable> T.appendInheritors(inheritanceMap: Map<PlatformData, Map<DRI, List<DRI>>>): Documentable =
+        InheritorsInfo(PlatformDependent(inheritanceMap.getForDRI(dri))).let { info ->
+            when (this) {
+                is DModule -> copy(packages = packages.map { it.appendInheritors(inheritanceMap) as DPackage })
+                is DPackage -> copy(classlikes = classlikes.map { it.appendInheritors(inheritanceMap) as DClasslike })
+                is DClass -> copy(
+                    extra = extra + info,
+                    classlikes = classlikes.map { it.appendInheritors(inheritanceMap) as DClasslike })
+                is DEnum -> copy(
+                    extra = extra + info,
+                    classlikes = classlikes.map { it.appendInheritors(inheritanceMap) as DClasslike })
+                is DInterface -> copy(
+                    extra = extra + info,
+                    classlikes = classlikes.map { it.appendInheritors(inheritanceMap) as DClasslike })
+                is DObject -> copy(classlikes = classlikes.map { it.appendInheritors(inheritanceMap) as DClasslike })
+                is DAnnotation -> copy(classlikes = classlikes.map { it.appendInheritors(inheritanceMap) as DClasslike })
+                else -> this
+            }
+        }
+
+    private fun Map<PlatformData, Map<DRI, List<DRI>>>.getForDRI(dri: DRI) =
+        PlatformDependent(map { (v, k) ->
+            v to k[dri]
+        }.map { (k, v) -> k to v.orEmpty() }.toMap())
+
+    private fun DModule.generateInheritanceMap() =
+        getInheritanceEntriesRec().filterNot { it.second.isEmpty() }.groupBy({ it.first }) { it.second }
+            .map { (k, v) ->
+                k to v.flatMap { p -> p.groupBy({ it.first }) { it.second }.toList() }
+                    .groupBy({ it.first }) { it.second }.map { (k2, v2) -> k2 to v2.flatten() }.toMap()
+            }.toMap()
+
+    private fun <T : Documentable> T.getInheritanceEntriesRec(): List<Pair<PlatformData, List<Pair<DRI, DRI>>>> =
+        this.toInheritanceEntries() + children.flatMap { it.getInheritanceEntriesRec() }
+
+    private fun <T : Documentable> T.toInheritanceEntries() =
+        (this as? WithSupertypes)?.let {
+            it.supertypes.map.map { (k, v) -> k to v.map { it to dri } }
+        }.orEmpty()
+
+}
+
+class InheritorsInfo(val value: PlatformDependent<List<DRI>>) : ExtraProperty<Documentable> {
+    companion object : ExtraProperty.Key<Documentable, InheritorsInfo> {
+        override fun mergeStrategyFor(left: InheritorsInfo, right: InheritorsInfo): MergeStrategy<Documentable> =
+            MergeStrategy.Replace(
+                InheritorsInfo(
+                    PlatformDependent(
+                        (left.value.map.entries.toList() + right.value.map.entries.toList())
+                            .groupBy({ it.key }) { it.value }
+                            .map { (k, v) -> k to v.flatten() }.toMap()
+                    )
+                )
+            )
+    }
+
+    override val key: ExtraProperty.Key<Documentable, *> = InheritorsInfo
+}
+
diff --git a/plugins/base/src/main/kotlin/translators/documentables/DefaultPageCreator.kt b/plugins/base/src/main/kotlin/translators/documentables/DefaultPageCreator.kt
index a0b5a072..eba56625 100644
--- a/plugins/base/src/main/kotlin/translators/documentables/DefaultPageCreator.kt
+++ b/plugins/base/src/main/kotlin/translators/documentables/DefaultPageCreator.kt
@@ -1,12 +1,14 @@
 package org.jetbrains.dokka.base.translators.documentables
 
 import org.jetbrains.dokka.base.signatures.SignatureProvider
+import org.jetbrains.dokka.base.transformers.documentables.InheritorsInfo
 import org.jetbrains.dokka.base.transformers.pages.comments.CommentsToContentConverter
 import org.jetbrains.dokka.links.DRI
 import org.jetbrains.dokka.model.*
 import org.jetbrains.dokka.model.DFunction
 import org.jetbrains.dokka.model.doc.Property
 import org.jetbrains.dokka.model.doc.TagWrapper
+import org.jetbrains.dokka.model.properties.WithExtraProperties
 import org.jetbrains.dokka.pages.*
 import org.jetbrains.dokka.utilities.DokkaLogger
 
@@ -85,6 +87,15 @@ open class DefaultPageCreator(
                 text(it.briefDocumentation())
             }
         }
+        (s as? WithExtraProperties<Documentable>)?.let { it.extra[InheritorsInfo] }?.let { inheritors ->
+            val map = inheritors.value.map
+            text("Subclasses:")
+            platformDependentHint(dri = s.dri, platformData = map.keys) {
+                map.keys.forEach { pd ->
+                    linkTable(map[pd].orEmpty(), ContentKind.Classlikes, setOf(pd))
+                }
+            }
+        }
     }
 
     protected open fun contentForClasslike(c: DClasslike) = contentBuilder.contentFor(c) {