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

7 files changed, 150 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)
diff --git a/core/testdata/format/inapplicableExtensionFunctions.kt b/core/testdata/format/inapplicableExtensionFunctions.kt
new file mode 100644
index 00000000..d2c65b46
--- /dev/null
+++ b/core/testdata/format/inapplicableExtensionFunctions.kt
@@ -0,0 +1,11 @@
+open class Foo<T> {
+}
+
+class Bar : Foo<Char>() {
+}
+
+fun Foo<Int>.shazam() {
+}
+
+fun Bar.xyzzy() {
+}
diff --git a/core/testdata/format/inapplicableExtensionFunctions.md b/core/testdata/format/inapplicableExtensionFunctions.md
new file mode 100644
index 00000000..8460cf2f
--- /dev/null
+++ b/core/testdata/format/inapplicableExtensionFunctions.md
@@ -0,0 +1,20 @@
+[test](test/index) / [Bar](test/-bar/index)
+
+
+# Bar
+
+`class Bar&nbsp;:&nbsp;[Foo](test/-foo/index)&lt;Char&gt;`
+
+
+
+### Constructors
+
+
+| [&lt;init&gt;](test/-bar/-init-) | `Bar()` |
+
+
+### Extension Functions
+
+
+| [xyzzy](test/xyzzy) | `fun [Bar](test/-bar/index).xyzzy(): Unit` |
+
diff --git a/core/testdata/format/shadowedExtensionFunctions.kt b/core/testdata/format/shadowedExtensionFunctions.kt
new file mode 100644
index 00000000..64df1ecb
--- /dev/null
+++ b/core/testdata/format/shadowedExtensionFunctions.kt
@@ -0,0 +1,18 @@
+open class Foo {
+}
+
+class Bar : Foo() {
+}
+
+fun Foo.xyzzy() {
+}
+
+fun Foo.shazam() {
+
+}
+
+fun Bar.xyzzy() {
+}
+
+fun Bar.shazam(i: Int) {
+}
diff --git a/core/testdata/format/shadowedExtensionFunctions.md b/core/testdata/format/shadowedExtensionFunctions.md
new file mode 100644
index 00000000..5218d299
--- /dev/null
+++ b/core/testdata/format/shadowedExtensionFunctions.md
@@ -0,0 +1,22 @@
+[test](test/index) / [Bar](test/-bar/index)
+
+
+# Bar
+
+`class Bar&nbsp;:&nbsp;[Foo](test/-foo/index)`
+
+
+
+### Constructors
+
+
+| [&lt;init&gt;](test/-bar/-init-) | `Bar()` |
+
+
+### Extension Functions
+
+
+| [shazam](test/shazam) | `fun [Bar](test/-bar/index).shazam(i:&nbsp;Int): Unit`
+`fun [Foo](test/-foo/index).shazam(): Unit` |
+| [xyzzy](test/xyzzy) | `fun [Bar](test/-bar/index).xyzzy(): Unit` |
+