rewrite propagation of extension functions to subclasses to handle shadowing and applicability more correctly

3 files changed, 79 insertions, 17 deletions

diff --git a/core/src/main/kotlin/Formats/StructuredFormatService.kt b/core/src/main/kotlin/Formats/StructuredFormatService.kt
index 01f38605..f23c27fa 100644
--- a/core/src/main/kotlin/Formats/StructuredFormatService.kt
+++ b/core/src/main/kotlin/Formats/StructuredFormatService.kt
@@ -309,7 +309,7 @@ abstract class StructuredFormatService(locationService: LocationService,
                 )
             })
 
-            val allExtensions = collectAllExtensions(node)
+            val allExtensions = node.extensions
             appendSection("Extension Properties", allExtensions.filter { it.kind == NodeKind.Property })
             appendSection("Extension Functions", allExtensions.filter { it.kind == NodeKind.Function })
             appendSection("Companion Object Extension Properties", allExtensions.filter { it.kind == NodeKind.CompanionObjectProperty })
@@ -379,18 +379,3 @@ abstract class StructuredFormatService(locationService: LocationService,
         }
     }
 }
-
-private fun collectAllExtensions(node: DocumentationNode): Collection<DocumentationNode> {
-    val result = LinkedHashSet<DocumentationNode>()
-    val visited = hashSetOf<DocumentationNode>()
-
-    fun collect(node: DocumentationNode) {
-        if (!visited.add(node)) return
-        result.addAll(node.extensions)
-        node.references(RefKind.Superclass).forEach { collect(it.to) }
-    }
-
-    collect(node)
-
-    return result
-}
diff --git a/core/src/main/kotlin/Kotlin/DocumentationBuilder.kt b/core/src/main/kotlin/Kotlin/DocumentationBuilder.kt
index a9701778..2a69bb13 100644
--- a/core/src/main/kotlin/Kotlin/DocumentationBuilder.kt
+++ b/core/src/main/kotlin/Kotlin/DocumentationBuilder.kt
@@ -33,6 +33,8 @@ import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeProjection
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
 import org.jetbrains.kotlin.types.typeUtil.supertypes
 
 data class DocumentationOptions(val outputDir: String,
@@ -273,6 +275,69 @@ class DocumentationBuilder
             val packageNode = findOrCreatePackageNode(packageName.asString(), packageContent)
             packageDocumentationBuilder.buildPackageDocumentation(this@DocumentationBuilder, packageName, packageNode, declarations)
         }
+
+        propagateExtensionFunctionsToSubclasses(fragments)
+    }
+
+    private fun propagateExtensionFunctionsToSubclasses(fragments: Collection<PackageFragmentDescriptor>) {
+        val allDescriptors = fragments.flatMap { it.getMemberScope().getContributedDescriptors() }
+        val allClasses = allDescriptors.filterIsInstance<ClassDescriptor>()
+        val classHierarchy = buildClassHierarchy(allClasses)
+
+        val allExtensionFunctions = allDescriptors
+                .filterIsInstance<CallableMemberDescriptor>()
+                .filter { it.extensionReceiverParameter != null }
+        val extensionFunctionsByName = allExtensionFunctions.groupBy { it.name }
+
+        allExtensionFunctions.forEach { extensionFunction ->
+            val possiblyShadowingFunctions = extensionFunctionsByName[extensionFunction.name]
+                    ?.filter { fn -> fn.canShadow(extensionFunction) }
+                    ?: emptyList()
+
+            val classDescriptor = extensionFunction.getExtensionClassDescriptor() ?: return@forEach
+            val subclasses = classHierarchy[classDescriptor] ?: return@forEach
+            subclasses.forEach { subclass ->
+                if (subclass.defaultType.isSubtypeOf(extensionFunction.extensionReceiverParameter!!.type) &&
+                        possiblyShadowingFunctions.none { subclass.defaultType.isSubtypeOf(it.extensionReceiverParameter!!.type) }) {
+                    refGraph.link(subclass.signature(), extensionFunction.signature(), RefKind.Extension)
+                }
+            }
+        }
+    }
+
+    private fun buildClassHierarchy(classes: List<ClassDescriptor>): Map<ClassDescriptor, List<ClassDescriptor>> {
+        val result = hashMapOf<ClassDescriptor, MutableList<ClassDescriptor>>()
+        classes.forEach { cls ->
+            TypeUtils.getAllSupertypes(cls.defaultType).forEach { supertype ->
+                val classDescriptor = supertype.constructor.declarationDescriptor as? ClassDescriptor
+                if (classDescriptor != null) {
+                    val subtypesList = result.getOrPut(classDescriptor) { arrayListOf() }
+                    subtypesList.add(cls)
+                }
+            }
+        }
+        return result
+    }
+
+    private fun CallableMemberDescriptor.canShadow(other: CallableMemberDescriptor): Boolean {
+        if (this == other) return false
+        if (this is PropertyDescriptor && other is PropertyDescriptor) {
+            return true
+        }
+        if (this is FunctionDescriptor && other is FunctionDescriptor) {
+            val parameters1 = valueParameters
+            val parameters2 = other.valueParameters
+            if (parameters1.size != parameters2.size) {
+                return false
+            }
+            for ((p1, p2) in parameters1 zip parameters2) {
+                if (p1.type != p2.type) {
+                    return false
+                }
+            }
+            return true
+        }
+        return false
     }
 
     fun DeclarationDescriptor.build(): DocumentationNode = when (this) {
@@ -608,7 +673,11 @@ fun CallableMemberDescriptor.getExtensionClassDescriptor(): ClassifierDescriptor
     val extensionReceiver = extensionReceiverParameter
     if (extensionReceiver != null) {
         val type = extensionReceiver.type
-        return type.constructor.declarationDescriptor as? ClassDescriptor
+        val receiverClass = type.constructor.declarationDescriptor as? ClassDescriptor
+        if ((receiverClass as? ClassDescriptor)?.isCompanionObject ?: false) {
+            return receiverClass?.containingDeclaration as? ClassifierDescriptor
+        }
+        return receiverClass
     }
     return null
 }
diff --git a/core/src/test/kotlin/format/MarkdownFormatTest.kt b/core/src/test/kotlin/format/MarkdownFormatTest.kt
index 58b80cee..509b5f39 100644
--- a/core/src/test/kotlin/format/MarkdownFormatTest.kt
+++ b/core/src/test/kotlin/format/MarkdownFormatTest.kt
@@ -230,6 +230,14 @@ public class MarkdownFormatTest {
         verifyMarkdownNodeByName("inheritedCompanionObjectProperties", "C")
     }
 
+    @Test fun shadowedExtensionFunctions() {
+        verifyMarkdownNodeByName("shadowedExtensionFunctions", "Bar")
+    }
+
+    @Test fun inapplicableExtensionFunctions() {
+        verifyMarkdownNodeByName("inapplicableExtensionFunctions", "Bar")
+    }
+
     private fun verifyMarkdownPackage(fileName: String, withKotlinRuntime: Boolean = false) {
         verifyOutput("testdata/format/$fileName.kt", ".package.md", withKotlinRuntime = withKotlinRuntime) { model, output ->
             markdownService.appendNodes(tempLocation, output, model.members)