fix rendering of extension function types

2 files changed, 16 insertions, 17 deletions

diff --git a/src/Kotlin/DocumentationBuilder.kt b/src/Kotlin/DocumentationBuilder.kt
index 9321e5ca..b526e0ed 100644
--- a/src/Kotlin/DocumentationBuilder.kt
+++ b/src/Kotlin/DocumentationBuilder.kt
@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant
 import org.jetbrains.kotlin.resolve.constants.ConstantValue
 import org.jetbrains.kotlin.resolve.constants.TypedCompileTimeConstant
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.resolve.descriptorUtil.isDocumentedAnnotation
 import org.jetbrains.kotlin.resolve.lazy.ResolveSession
 import org.jetbrains.kotlin.resolve.source.PsiSourceElement
@@ -327,6 +328,7 @@ class DocumentationBuilder(val resolutionFacade: ResolutionFacade,
         }
 
         append(node, DocumentationReference.Kind.Detail)
+        node.appendAnnotations(jetType)
         for (typeArgument in jetType.arguments)
             node.appendProjection(typeArgument)
     }
@@ -335,7 +337,7 @@ class DocumentationBuilder(val resolutionFacade: ResolutionFacade,
         DescriptorUtils.getFqName(this).asString() in boringBuiltinClasses
 
     fun DocumentationNode.appendAnnotations(annotated: Annotated) {
-        annotated.annotations.filter { it.source.getPsi() != null && it.mustBeDocumented() }.forEach {
+        annotated.annotations.filter { it.isDocumented() }.forEach {
             val annotationNode = it.build()
             if (annotationNode != null) {
                 append(annotationNode,
@@ -344,6 +346,12 @@ class DocumentationBuilder(val resolutionFacade: ResolutionFacade,
         }
     }
 
+    private fun AnnotationDescriptor.isDocumented(): Boolean {
+        if (source.getPsi() != null && mustBeDocumented()) return true
+        val annotationClassName = type.constructor.declarationDescriptor?.fqNameSafe?.asString()
+        return annotationClassName == "kotlin.Extension"
+    }
+
     fun AnnotationDescriptor.mustBeDocumented(): Boolean {
         val annotationClass = type.constructor.declarationDescriptor as? Annotated ?: return false
         return annotationClass.isDocumentedAnnotation()
diff --git a/src/Kotlin/KotlinLanguageService.kt b/src/Kotlin/KotlinLanguageService.kt
index bd98579a..6541e46c 100644
--- a/src/Kotlin/KotlinLanguageService.kt
+++ b/src/Kotlin/KotlinLanguageService.kt
@@ -76,26 +76,17 @@ class KotlinLanguageService : LanguageService {
     }
 
     private fun ContentBlock.renderType(node: DocumentationNode) {
-        val typeArguments = node.details(DocumentationNode.Kind.Type)
+        var typeArguments = node.details(DocumentationNode.Kind.Type)
         if (node.name == "Function${typeArguments.count() - 1}") {
             // lambda
-            symbol("(")
-            renderList(typeArguments.take(typeArguments.size() - 1), noWrap = true) {
-                renderType(it)
+            val isExtension = node.annotations.any { it.name == "Extension" }
+            if (isExtension) {
+                renderType(typeArguments.first())
+                symbol(".")
+                typeArguments = typeArguments.drop(1)
             }
-            symbol(")")
-            nbsp()
-            symbol("->")
-            nbsp()
-            renderType(typeArguments.last())
-            return
-        }
-        if (node.name == "ExtensionFunction${typeArguments.count() - 2}") {
-            // extension lambda
-            renderType(typeArguments.first())
-            symbol(".")
             symbol("(")
-            renderList(typeArguments.drop(1).take(typeArguments.size() - 2), noWrap = true) {
+            renderList(typeArguments.take(typeArguments.size() - 1), noWrap = true) {
                 renderType(it)
             }
             symbol(")")