Display inherited extensions (#2625)

2 files changed, 199 insertions, 89 deletions

diff --git a/plugins/base/src/main/kotlin/transformers/documentables/ExtensionExtractorTransformer.kt b/plugins/base/src/main/kotlin/transformers/documentables/ExtensionExtractorTransformer.kt
index 73023a86..19af0564 100644
--- a/plugins/base/src/main/kotlin/transformers/documentables/ExtensionExtractorTransformer.kt
+++ b/plugins/base/src/main/kotlin/transformers/documentables/ExtensionExtractorTransformer.kt
@@ -1,14 +1,8 @@
 package org.jetbrains.dokka.base.transformers.documentables
 
 import kotlinx.coroutines.*
-import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.channels.ClosedReceiveChannelException
-import kotlinx.coroutines.channels.ReceiveChannel
-import kotlinx.coroutines.channels.SendChannel
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.flow
-import kotlinx.coroutines.flow.flowOn
-import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.channels.*
+import org.jetbrains.dokka.base.transformers.documentables.utils.FullClassHierarchyBuilder
 import org.jetbrains.dokka.links.DRI
 import org.jetbrains.dokka.links.DriOfAny
 import org.jetbrains.dokka.model.*
@@ -17,97 +11,140 @@ import org.jetbrains.dokka.model.properties.MergeStrategy
 import org.jetbrains.dokka.model.properties.plus
 import org.jetbrains.dokka.plugability.DokkaContext
 import org.jetbrains.dokka.transformers.documentation.DocumentableTransformer
+import org.jetbrains.dokka.utilities.parallelForEach
+import org.jetbrains.dokka.utilities.parallelMap
 
 
 class ExtensionExtractorTransformer : DocumentableTransformer {
     override fun invoke(original: DModule, context: DokkaContext): DModule = runBlocking(Dispatchers.Default) {
+        val classGraph = async {
+            if (!context.configuration.suppressInheritedMembers)
+                FullClassHierarchyBuilder()(original)
+            else
+                emptyMap()
+        }
+
         val channel = Channel<Pair<DRI, Callable>>(10)
         launch {
-            coroutineScope {
-                original.packages.forEach { launch { collectExtensions(it, channel) } }
-            }
+            original.packages.parallelForEach { collectExtensions(it, channel) }
             channel.close()
         }
-        val extensionMap = channel.consumeAsFlow().toList().toMultiMap()
+        val extensionMap = channel.toList().toMultiMap()
 
-        val newPackages = original.packages.map { async { it.addExtensionInformation(extensionMap) } }
-        original.copy(packages = newPackages.awaitAll())
+        val newPackages = original.packages.parallelMap { it.addExtensionInformation(classGraph.await(), extensionMap) }
+        original.copy(packages = newPackages)
     }
-}
 
-private suspend fun <T : Documentable> T.addExtensionInformation(
-    extensionMap: Map<DRI, List<Callable>>
-): T = coroutineScope {
-    val newClasslikes = (this@addExtensionInformation as? WithScope)
-        ?.classlikes
-        ?.map { async { it.addExtensionInformation(extensionMap) } }
-        .orEmpty()
-
-    @Suppress("UNCHECKED_CAST")
-    when (this@addExtensionInformation) {
-        is DPackage -> {
-            val newTypealiases = typealiases.map { async { it.addExtensionInformation(extensionMap) } }
-            copy(classlikes = newClasslikes.awaitAll(), typealiases = newTypealiases.awaitAll())
-        }
-        is DClass -> copy(classlikes = newClasslikes.awaitAll(), extra = extra + extensionMap.find(dri))
-        is DEnum -> copy(classlikes = newClasslikes.awaitAll(), extra = extra + extensionMap.find(dri))
-        is DInterface -> copy(classlikes = newClasslikes.awaitAll(), extra = extra + extensionMap.find(dri))
-        is DObject -> copy(classlikes = newClasslikes.awaitAll(), extra = extra + extensionMap.find(dri))
-        is DAnnotation -> copy(classlikes = newClasslikes.awaitAll(), extra = extra + extensionMap.find(dri))
-        is DTypeAlias -> copy(extra = extra + extensionMap.find(dri))
-        else -> throw IllegalStateException(
-            "${this@addExtensionInformation::class.simpleName} is not expected to have extensions"
-        )
-    } as T
-}
+    private suspend fun <T : Documentable> T.addExtensionInformation(
+        classGraph: SourceSetDependent<Map<DRI, List<DRI>>>,
+        extensionMap: Map<DRI, List<Callable>>
+    ): T = coroutineScope {
+        val newClasslikes = (this@addExtensionInformation as? WithScope)
+            ?.classlikes
+            ?.map { async { it.addExtensionInformation(classGraph, extensionMap) } }
+            .orEmpty()
+
+        @Suppress("UNCHECKED_CAST")
+        when (this@addExtensionInformation) {
+            is DPackage -> {
+                val newTypealiases = typealiases.map { async { it.addExtensionInformation(classGraph, extensionMap) } }
+                copy(classlikes = newClasslikes.awaitAll(), typealiases = newTypealiases.awaitAll())
+            }
 
-private fun Map<DRI, List<Callable>>.find(dri: DRI) = get(dri)?.toSet()?.let(::CallableExtensions)
+            is DClass -> copy(
+                classlikes = newClasslikes.awaitAll(),
+                extra = extra + findExtensions(classGraph, extensionMap)
+            )
+
+            is DEnum -> copy(
+                classlikes = newClasslikes.awaitAll(),
+                extra = extra + findExtensions(classGraph, extensionMap)
+            )
+
+            is DInterface -> copy(
+                classlikes = newClasslikes.awaitAll(),
+                extra = extra + findExtensions(classGraph, extensionMap)
+            )
+
+            is DObject -> copy(
+                classlikes = newClasslikes.awaitAll(),
+                extra = extra + findExtensions(classGraph, extensionMap)
+            )
+
+            is DAnnotation -> copy(
+                classlikes = newClasslikes.awaitAll(),
+                extra = extra + findExtensions(classGraph, extensionMap)
+            )
+
+            is DTypeAlias -> copy(extra = extra + findExtensions(classGraph, extensionMap))
+            else -> throw IllegalStateException(
+                "${this@addExtensionInformation::class.simpleName} is not expected to have extensions"
+            )
+        } as T
+    }
+
+    private suspend fun collectExtensions(
+        documentable: Documentable,
+        channel: SendChannel<Pair<DRI, Callable>>
+    ): Unit = coroutineScope {
+        if (documentable is WithScope) {
+            documentable.classlikes.forEach {
+                launch { collectExtensions(it, channel) }
+            }
 
-private suspend fun collectExtensions(
-    documentable: Documentable,
-    channel: SendChannel<Pair<DRI, Callable>>
-): Unit = coroutineScope {
-    if (documentable is WithScope) {
-        documentable.classlikes.forEach {
-            launch { collectExtensions(it, channel) }
+            if (documentable is DObject || documentable is DPackage) {
+                (documentable.properties.asSequence() + documentable.functions.asSequence())
+                    .flatMap { it.asPairsWithReceiverDRIs() }
+                    .forEach { channel.send(it) }
+            }
         }
+    }
+
+    private fun <T : Documentable> T.findExtensions(
+        classGraph: SourceSetDependent<Map<DRI, List<DRI>>>,
+        extensionMap: Map<DRI, List<Callable>>
+    ): CallableExtensions? {
+        val resultSet = mutableSetOf<Callable>()
 
-        if (documentable is DObject || documentable is DPackage) {
-            (documentable.properties.asSequence() + documentable.functions.asSequence())
-                .flatMap(Callable::asPairsWithReceiverDRIs)
-                .forEach { channel.send(it) }
+        fun collectFrom(element: DRI) {
+            extensionMap[element]?.let { resultSet.addAll(it) }
+            sourceSets.forEach { sourceSet -> classGraph[sourceSet]?.get(element)?.forEach { collectFrom(it) } }
         }
+        collectFrom(dri)
+
+        return if (resultSet.isEmpty()) null else CallableExtensions(resultSet)
     }
-}
 
+    private fun Callable.asPairsWithReceiverDRIs(): Sequence<Pair<DRI, Callable>> =
+        receiver?.type?.let { findReceiverDRIs(it) }.orEmpty().map { it to this }
+
+    // In normal cases we return at max one DRI, but sometimes receiver type can be bound by more than one type constructor
+    // for example `fun <T> T.example() where T: A, T: B` is extension of both types A and B
+    // another one `typealias A = B`
+    // Note: in some cases returning empty sequence doesn't mean that we cannot determine the DRI but only that we don't
+    // care about it since there is nowhere to put documentation of given extension.
+    private fun Callable.findReceiverDRIs(bound: Bound): Sequence<DRI> = when (bound) {
+        is Nullable -> findReceiverDRIs(bound.inner)
+        is DefinitelyNonNullable -> findReceiverDRIs(bound.inner)
+        is TypeParameter ->
+            if (this is DFunction && bound.dri == this.dri)
+                generics.find { it.name == bound.name }?.bounds?.asSequence()?.flatMap { findReceiverDRIs(it) }.orEmpty()
+            else
+                emptySequence()
+
+        is TypeConstructor -> sequenceOf(bound.dri)
+        is PrimitiveJavaType -> emptySequence()
+        is Void -> emptySequence()
+        is JavaObject -> sequenceOf(DriOfAny)
+        is Dynamic -> sequenceOf(DriOfAny)
+        is UnresolvedBound -> emptySequence()
+        is TypeAliased -> findReceiverDRIs(bound.typeAlias) + findReceiverDRIs(bound.inner)
+    }
 
-private fun Callable.asPairsWithReceiverDRIs(): Sequence<Pair<DRI, Callable>> =
-    receiver?.type?.let(::findReceiverDRIs).orEmpty().map { it to this }
-
-// In normal cases we return at max one DRI, but sometimes receiver type can be bound by more than one type constructor
-// for example `fun <T> T.example() where T: A, T: B` is extension of both types A and B
-// Note: in some cases returning empty sequence doesn't mean that we cannot determine the DRI but only that we don't
-// care about it since there is nowhere to put documentation of given extension.
-private fun Callable.findReceiverDRIs(bound: Bound): Sequence<DRI> = when (bound) {
-    is Nullable -> findReceiverDRIs(bound.inner)
-    is DefinitelyNonNullable -> findReceiverDRIs(bound.inner)
-    is TypeParameter ->
-        if (this is DFunction && bound.dri == this.dri)
-            generics.find { it.name == bound.name }?.bounds?.asSequence()?.flatMap(::findReceiverDRIs).orEmpty()
-        else
-            emptySequence()
-    is TypeConstructor -> sequenceOf(bound.dri)
-    is PrimitiveJavaType -> emptySequence()
-    is Void -> emptySequence()
-    is JavaObject -> sequenceOf(DriOfAny)
-    is Dynamic -> sequenceOf(DriOfAny)
-    is UnresolvedBound -> emptySequence()
-    is TypeAliased -> findReceiverDRIs(bound.typeAlias)
+    private fun <T, U> Iterable<Pair<T, U>>.toMultiMap(): Map<T, List<U>> =
+        groupBy(Pair<T, *>::first, Pair<*, U>::second)
 }
 
-private fun <T, U> Iterable<Pair<T, U>>.toMultiMap(): Map<T, List<U>> =
-    groupBy(Pair<T, *>::first, Pair<*, U>::second)
-
 data class CallableExtensions(val extensions: Set<Callable>) : ExtraProperty<Documentable> {
     companion object Key : ExtraProperty.Key<Documentable, CallableExtensions> {
         override fun mergeStrategyFor(left: CallableExtensions, right: CallableExtensions) =
@@ -116,14 +153,3 @@ data class CallableExtensions(val extensions: Set<Callable>) : ExtraProperty<Doc
 
     override val key = Key
 }
-
-//TODO IMPORTANT remove this terrible hack after updating to 1.4-M3
-fun <T : Any> ReceiveChannel<T>.consumeAsFlow(): Flow<T> = flow {
-    try {
-        while (true) {
-            emit(receive())
-        }
-    } catch (_: ClosedReceiveChannelException) {
-        // cool and good
-    }
-}.flowOn(Dispatchers.Default)
diff --git a/plugins/base/src/main/kotlin/transformers/documentables/utils/FullClassHierarchyBuilder.kt b/plugins/base/src/main/kotlin/transformers/documentables/utils/FullClassHierarchyBuilder.kt
new file mode 100644
index 00000000..d657fa32
--- /dev/null
+++ b/plugins/base/src/main/kotlin/transformers/documentables/utils/FullClassHierarchyBuilder.kt
@@ -0,0 +1,84 @@
+package org.jetbrains.dokka.base.transformers.documentables.utils
+
+import com.intellij.psi.PsiClass
+import kotlinx.coroutines.*
+import org.jetbrains.dokka.analysis.DescriptorDocumentableSource
+import org.jetbrains.dokka.analysis.PsiDocumentableSource
+import org.jetbrains.dokka.analysis.from
+import org.jetbrains.dokka.links.DRI
+import org.jetbrains.dokka.model.*
+import org.jetbrains.dokka.utilities.parallelForEach
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.typeUtil.immediateSupertypes
+import org.jetbrains.kotlin.types.typeUtil.isAnyOrNullableAny
+import java.util.concurrent.ConcurrentHashMap
+
+typealias Supertypes = List<DRI>
+typealias ClassHierarchy = SourceSetDependent<Map<DRI, Supertypes>>
+
+class FullClassHierarchyBuilder {
+    suspend operator fun invoke(original: DModule): ClassHierarchy = coroutineScope {
+        val map = original.sourceSets.associateWith { ConcurrentHashMap<DRI, List<DRI>>() }
+        original.packages.parallelForEach { visitDocumentable(it, map) }
+        map
+    }
+
+    private suspend fun collectSupertypesFromKotlinType(
+        driWithKType: Pair<DRI, KotlinType>,
+        supersMap: MutableMap<DRI, Supertypes>
+    ): Unit = coroutineScope {
+        val (dri, kotlinType) = driWithKType
+        val supertypes = kotlinType.immediateSupertypes().filterNot { it.isAnyOrNullableAny() }
+        val supertypesDriWithKType = supertypes.mapNotNull { supertype ->
+            supertype.constructor.declarationDescriptor?.let {
+                DRI.from(it) to supertype
+            }
+        }
+
+        if (supersMap[dri] == null) {
+            // another thread can rewrite the same value, but it isn't a problem
+            supersMap[dri] = supertypesDriWithKType.map { it.first }
+            supertypesDriWithKType.parallelForEach { collectSupertypesFromKotlinType(it, supersMap) }
+        }
+    }
+
+    private suspend fun collectSupertypesFromPsiClass(
+        driWithPsiClass: Pair<DRI, PsiClass>,
+        supersMap: MutableMap<DRI, Supertypes>
+    ): Unit = coroutineScope {
+        val (dri, psiClass) = driWithPsiClass
+        val supertypes = psiClass.superTypes.mapNotNull { it.resolve() }
+            .filterNot { it.qualifiedName == "java.lang.Object" }
+        val supertypesDriWithPsiClass = supertypes.map { DRI.from(it) to it }
+
+        if (supersMap[dri] == null) {
+            // another thread can rewrite the same value, but it isn't a problem
+            supersMap[dri] = supertypesDriWithPsiClass.map { it.first }
+            supertypesDriWithPsiClass.parallelForEach { collectSupertypesFromPsiClass(it, supersMap) }
+        }
+    }
+
+    private suspend fun visitDocumentable(
+        documentable: Documentable,
+        hierarchy: SourceSetDependent<MutableMap<DRI, List<DRI>>>
+    ): Unit = coroutineScope {
+        if (documentable is WithScope) {
+            documentable.classlikes.parallelForEach { visitDocumentable(it, hierarchy) }
+        }
+        if (documentable is DClasslike) {
+            // to build a full class graph, using supertypes from Documentable
+            // is not enough since it keeps only one level of hierarchy
+            documentable.sources.forEach { (sourceSet, source) ->
+                if (source is DescriptorDocumentableSource) {
+                    val descriptor = source.descriptor as ClassDescriptor
+                    val type = descriptor.defaultType
+                    hierarchy[sourceSet]?.let { collectSupertypesFromKotlinType(documentable.dri to type, it) }
+                } else if (source is PsiDocumentableSource) {
+                    val psi = source.psi as PsiClass
+                    hierarchy[sourceSet]?.let { collectSupertypesFromPsiClass(documentable.dri to psi, it) }
+                }
+            }
+        }
+    }
+}
\ No newline at end of file