Display inherited extensions

Kotlin · Aug 16, 2022 · 1cb34ac · 1cb34ac
1 parent 018af7d
commit 1cb34ac
Show file tree

Hide file tree

Showing 4 changed files with 293 additions and 93 deletions.
diff --git a/core/src/main/kotlin/utilities/parallelCollectionOperations.kt b/core/src/main/kotlin/utilities/parallelCollectionOperations.kt
@@ -1,8 +1,6 @@
 package org.jetbrains.dokka.utilities
 
-import kotlinx.coroutines.async
-import kotlinx.coroutines.awaitAll
-import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.*
 
 suspend inline fun <A, B> Iterable<A>.parallelMap(crossinline f: suspend (A) -> B): List<B> = coroutineScope {
     map { async { f(it) } }.awaitAll()
@@ -13,5 +11,5 @@ suspend inline fun <A, B> Iterable<A>.parallelMapNotNull(crossinline f: suspend
 }
 
 suspend inline fun <A> Iterable<A>.parallelForEach(crossinline f: suspend (A) -> Unit): Unit = coroutineScope {
-    map { async { f(it) } }.awaitAll()
+    forEach { launch { f(it) } }
 }
diff --git 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.ClassGraphBuilder
 import org.jetbrains.dokka.links.DRI
 import org.jetbrains.dokka.links.DriOfAny
 import org.jetbrains.dokka.model.*
@@ -17,97 +11,138 @@ 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.kotlin.utils.addToStdlib.popLast
+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 { ClassGraphBuilder()(original) }
+
         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())
+    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) }
+            }
         }
-        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 fun Map<DRI, List<Callable>>.find(dri: DRI) = get(dri)?.toSet()?.let(::CallableExtensions)
+    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 suspend fun collectExtensions(
-    documentable: Documentable,
-    channel: SendChannel<Pair<DRI, Callable>>
-): Unit = coroutineScope {
-    if (documentable is WithScope) {
-        documentable.classlikes.forEach {
-            launch { collectExtensions(it, channel) }
-        }
+            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
+    }
 
-        if (documentable is DObject || documentable is DPackage) {
-            (documentable.properties.asSequence() + documentable.functions.asSequence())
-                .flatMap(Callable::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>()
+        extensionMap[dri]?.let { resultSet.addAll(it) }
+
+        val queue = mutableListOf<DRI>()
+        sourceSets.forEach { classGraph[it]?.get(dri)?.let { supertypesDRIs -> queue.addAll(supertypesDRIs) } }
+        while (queue.isNotEmpty()) {
+            val element = queue.popLast()
+            sourceSets.forEach { classGraph[it]?.get(element)?.let { supertypesDRIs -> queue.addAll(supertypesDRIs) } }
+            extensionMap[element]?.let { resultSet.addAll(it) }
         }
+        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 +151,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/ClassGraphBuilder.kt b/plugins/base/src/main/kotlin/transformers/documentables/utils/ClassGraphBuilder.kt
@@ -0,0 +1,66 @@
+package org.jetbrains.dokka.base.transformers.documentables.utils
+
+import kotlinx.coroutines.*
+import org.jetbrains.dokka.DokkaConfiguration
+import org.jetbrains.dokka.analysis.DescriptorDocumentableSource
+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
+
+
+/**
+ * The class allows lo build a full class hierarchy via descriptors
+ */
+class ClassGraphBuilder {
+    suspend operator fun invoke(original: DModule,): SourceSetDependent<Map<DRI, List<DRI>>> = 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>,
+        sourceSet: DokkaConfiguration.DokkaSourceSet,
+        supersMap: SourceSetDependent<MutableMap<DRI, List<DRI>>>
+    ): Unit = coroutineScope {
+        val supertypes = DRIWithKType.second.immediateSupertypes().filterNot { it.isAnyOrNullableAny() }
+        val supertypesDRIWithKType = supertypes.mapNotNull { supertype ->
+            supertype.constructor.declarationDescriptor?.let {
+                DRI.from(it) to supertype
+            }
+        }
+
+        supersMap[sourceSet]?.let {
+            if (it[DRIWithKType.first] == null) {
+                // another thread can rewrite the same value, but it isn't a problem
+                it[DRIWithKType.first] = supertypesDRIWithKType.map { it.first }
+                supertypesDRIWithKType.parallelForEach { collectSupertypesFromKotlinType(it, sourceSet, supersMap) }
+            }
+        }
+    }
+
+    private suspend fun visitDocumentable(
+        documentable: Documentable,
+        supersMap: SourceSetDependent<MutableMap<DRI, List<DRI>>>
+    ): Unit = coroutineScope {
+        if (documentable is WithScope) {
+            documentable.classlikes.
+            parallelForEach{ visitDocumentable(it, supersMap) }
+        }
+        if(documentable is DClasslike) {
+            documentable.sources.forEach { (sourceSet, source) ->
+                if (source is DescriptorDocumentableSource) {
+                    val descriptor = source.descriptor as ClassDescriptor
+                    val type = descriptor.defaultType
+                    collectSupertypesFromKotlinType( documentable.dri to type, sourceSet, supersMap)
+                }
+            }
+        }
+    }
+}