generate ExternalClass nodes to hold extension functions and properties for classes from other packages

3 files changed, 46 insertions, 3 deletions

diff --git a/src/Formats/StructuredFormatService.kt b/src/Formats/StructuredFormatService.kt
index b75f39d1..cb510f80 100644
--- a/src/Formats/StructuredFormatService.kt
+++ b/src/Formats/StructuredFormatService.kt
@@ -178,10 +178,15 @@ public abstract class StructuredFormatService(val locationService: LocationServi
         for ((breadcrumbs, items) in breakdownByLocation) {
             appendLine(to, breadcrumbs)
             appendLine(to)
-            appendLocation(location, to, items)
+            appendLocation(location, to, items.filter { it.kind != DocumentationNode.Kind.ExternalClass })
         }
 
         for (node in nodes) {
+            if (node.kind == DocumentationNode.Kind.ExternalClass) {
+                appendSection(location, "Extensions for ${node.name}", node.members, node, to)
+                continue
+            }
+
             appendSection(location, "Packages", node.members(DocumentationNode.Kind.Package), node, to)
             appendSection(location, "Types", node.members.filter {
                 it.kind in setOf(
@@ -191,6 +196,7 @@ public abstract class StructuredFormatService(val locationService: LocationServi
                         DocumentationNode.Kind.Object,
                         DocumentationNode.Kind.AnnotationClass)
             }, node, to)
+            appendSection(location, "Extensions for External Classes", node.members(DocumentationNode.Kind.ExternalClass), node, to)
             appendSection(location, "Constructors", node.members(DocumentationNode.Kind.Constructor), node, to)
             appendSection(location, "Properties", node.members(DocumentationNode.Kind.Property), node, to)
             appendSection(location, "Functions", node.members(DocumentationNode.Kind.Function), node, to)
@@ -210,7 +216,8 @@ public abstract class StructuredFormatService(val locationService: LocationServi
                         DocumentationNode.Kind.Function,
                         DocumentationNode.Kind.PropertyAccessor,
                         DocumentationNode.Kind.ClassObjectProperty,
-                        DocumentationNode.Kind.ClassObjectFunction
+                        DocumentationNode.Kind.ClassObjectFunction,
+                        DocumentationNode.Kind.ExternalClass
                         )
             }, node, to)
             appendSection(location, "Extensions", node.extensions, node, to)
diff --git a/src/Kotlin/DocumentationBuilder.kt b/src/Kotlin/DocumentationBuilder.kt
index c00580c5..99c81760 100644
--- a/src/Kotlin/DocumentationBuilder.kt
+++ b/src/Kotlin/DocumentationBuilder.kt
@@ -15,6 +15,12 @@ import org.jetbrains.jet.lang.descriptors.impl.EnumEntrySyntheticClassDescriptor
 
 public data class DocumentationOptions(val includeNonPublic: Boolean = false)
 
+private fun isSamePackage(descriptor1: DeclarationDescriptor, descriptor2: DeclarationDescriptor): Boolean {
+    val package1 = DescriptorUtils.getParentOfType(descriptor1, javaClass<PackageFragmentDescriptor>())
+    val package2 = DescriptorUtils.getParentOfType(descriptor2, javaClass<PackageFragmentDescriptor>())
+    return package1 != null && package2 != null && package1.fqName == package2.fqName
+}
+
 class DocumentationBuilder(val session: ResolveSession, val options: DocumentationOptions) {
     val visibleToDocumentation = setOf(Visibilities.INTERNAL, Visibilities.PROTECTED, Visibilities.PUBLIC)
     val descriptorToNode = hashMapOf<DeclarationDescriptor, DocumentationNode>()
@@ -141,6 +147,22 @@ class DocumentationBuilder(val session: ResolveSession, val options: Documentati
         descriptors.forEach { descriptor -> appendChild(descriptor, kind) }
     }
 
+    fun DocumentationNode.getParentForPackageMember(descriptor: DeclarationDescriptor,
+                                                    externalClassNodes: MutableMap<FqName, DocumentationNode>): DocumentationNode {
+        if (descriptor is CallableMemberDescriptor) {
+            val extensionClassDescriptor = descriptor.getExtensionClassDescriptor()
+            if (extensionClassDescriptor != null && !isSamePackage(descriptor, extensionClassDescriptor)) {
+                val fqName = DescriptorUtils.getFqNameFromTopLevelClass(extensionClassDescriptor)
+                return externalClassNodes.getOrPut(fqName, {
+                    val newNode = DocumentationNode(fqName.asString(), Content.Empty, Kind.ExternalClass)
+                    append(newNode, DocumentationReference.Kind.Member)
+                    newNode
+                })
+            }
+        }
+        return this
+    }
+
     fun DocumentationNode.appendFragments(fragments: Collection<PackageFragmentDescriptor>) {
         val descriptors = hashMapOf<String, List<DeclarationDescriptor>>()
         for ((name, parts) in fragments.groupBy { it.fqName }) {
@@ -149,7 +171,11 @@ class DocumentationBuilder(val session: ResolveSession, val options: Documentati
         for ((packageName, declarations) in descriptors) {
             println("  package $packageName: ${declarations.count()} nodes")
             val packageNode = DocumentationNode(packageName, Content.Empty, Kind.Package)
-            packageNode.appendChildren(declarations, DocumentationReference.Kind.Member)
+            val externalClassNodes = hashMapOf<FqName, DocumentationNode>()
+            declarations.forEach { descriptor ->
+                val parent = packageNode.getParentForPackageMember(descriptor, externalClassNodes)
+                parent.appendChild(descriptor, DocumentationReference.Kind.Member)
+            }
             append(packageNode, DocumentationReference.Kind.Member)
         }
     }
@@ -205,6 +231,15 @@ class DocumentationBuilder(val session: ResolveSession, val options: Documentati
     private fun DeclarationDescriptor.inClassObject() =
             getContainingDeclaration().let { it is ClassDescriptor && it.getKind() == ClassKind.CLASS_OBJECT }
 
+    fun CallableMemberDescriptor.getExtensionClassDescriptor(): ClassifierDescriptor? {
+        val extensionReceiver = getExtensionReceiverParameter()
+        if (extensionReceiver != null) {
+            val type = extensionReceiver.getType()
+            return type.getConstructor().getDeclarationDescriptor() as? ClassDescriptor
+        }
+        return null
+    }
+
     fun FunctionDescriptor.build(): DocumentationNode {
         val node = DocumentationNode(this, if (inClassObject()) Kind.ClassObjectFunction else Kind.Function)
 
diff --git a/src/Model/DocumentationNode.kt b/src/Model/DocumentationNode.kt
index caae77a8..5f9aabab 100644
--- a/src/Model/DocumentationNode.kt
+++ b/src/Model/DocumentationNode.kt
@@ -89,6 +89,7 @@ public open class DocumentationNode(val name: String,
 
         Module
 
+        ExternalClass
         Annotation
 
         Value