Skip to content

Commit

Permalink
Display inherited extensions (#2625)
Browse files Browse the repository at this point in the history
  • Loading branch information
vmishenev committed Aug 29, 2022
1 parent 34a8ae1 commit e68eea6
Show file tree
Hide file tree
Showing 5 changed files with 358 additions and 97 deletions.
@@ -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()
Expand All @@ -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) } }
}
9 changes: 5 additions & 4 deletions plugins/base/api/base.api
Expand Up @@ -1217,10 +1217,6 @@ public final class org/jetbrains/dokka/base/transformers/documentables/Extension
public fun invoke (Lorg/jetbrains/dokka/model/DModule;Lorg/jetbrains/dokka/plugability/DokkaContext;)Lorg/jetbrains/dokka/model/DModule;
}

public final class org/jetbrains/dokka/base/transformers/documentables/ExtensionExtractorTransformerKt {
public static final fun consumeAsFlow (Lkotlinx/coroutines/channels/ReceiveChannel;)Lkotlinx/coroutines/flow/Flow;
}

public final class org/jetbrains/dokka/base/transformers/documentables/InheritedEntriesDocumentableFilterTransformer : org/jetbrains/dokka/base/transformers/documentables/SuppressedByConditionDocumentableFilterTransformer {
public fun <init> (Lorg/jetbrains/dokka/plugability/DokkaContext;)V
public fun shouldBeSuppressed (Lorg/jetbrains/dokka/model/Documentable;)Z
Expand Down Expand Up @@ -1278,6 +1274,11 @@ public final class org/jetbrains/dokka/base/transformers/documentables/UtilsKt {
public static final fun isException (Lorg/jetbrains/dokka/model/properties/WithExtraProperties;)Z
}

public final class org/jetbrains/dokka/base/transformers/documentables/utils/FullClassHierarchyBuilder {
public fun <init> ()V
public final fun invoke (Lorg/jetbrains/dokka/model/DModule;Lkotlin/coroutines/Continuation;)Ljava/lang/Object;
}

public final class org/jetbrains/dokka/base/transformers/pages/annotations/SinceKotlinTransformer : org/jetbrains/dokka/transformers/documentation/DocumentableTransformer {
public fun <init> (Lorg/jetbrains/dokka/plugability/DokkaContext;)V
public final fun getContext ()Lorg/jetbrains/dokka/plugability/DokkaContext;
Expand Down
@@ -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.*
Expand All @@ -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) =
Expand All @@ -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)
@@ -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) }
}
}
}
}
}

0 comments on commit e68eea6

Please sign in to comment.