Wrote separate descriptor renderer

It was easier to write my own renderer than change current one

compiler/visualizer/render-psi/src/org/jetbrains/kotlin/compiler/visualizer/PsiRenderer.kt

@@ -5,47 +5,45 @@
 
 package org.jetbrains.kotlin.compiler.visualizer
 
-import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.analyzer.AnalysisResult
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.compiler.visualizer.Annotator.annotate
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
 import org.jetbrains.kotlin.renderer.ClassifierNamePolicy
 import org.jetbrains.kotlin.renderer.DescriptorRenderer
 import org.jetbrains.kotlin.renderer.ParameterNameRenderingPolicy
+import org.jetbrains.kotlin.renderer.render
 import org.jetbrains.kotlin.resolve.BindingContext.*
+import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getAbbreviatedTypeOrType
 import org.jetbrains.kotlin.resolve.calls.callUtil.getCall
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
+import org.jetbrains.kotlin.resolve.descriptorUtil.declaresOrInheritsDefaultValue
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.types.*
+import java.util.ArrayList
 
 class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : BaseRenderer {
     val bindingContext = analysisResult.bindingContext
     private val annotations = mutableListOf<Annotator.AnnotationInfo>()
 
-    val descriptorRenderer: DescriptorRenderer = DescriptorRenderer.withOptions {
+    val descriptorRenderer = PsiDescriptorRenderer()
-        withDefinedIn = false
-        modifiers = emptySet()
-        receiverAfterName = false
-        startFromName = false
-        withSourceFileForTopLevel = false
-        renderConstructorKeyword = false
-        classifierNamePolicy = ClassifierNamePolicy.SOURCE_CODE_QUALIFIED
-        includeAdditionalModifiers = false
-        renderContextNearLocalVariable = true
-        fullContextForLocalVariable = false
-        useBaseClassAsReceiver = true
-        parameterNameRenderingPolicy = ParameterNameRenderingPolicy.NONE
-    }
 
     private val unnecessaryData = mapOf(
         "kotlin." to "",
-        "collections." to ""
+        "kotlin/" to ""
     )
 
-    private fun addAnnotation(annotationText: String, element: PsiElement) {
+    private fun addAnnotation(annotationText: String, element: PsiElement, deleteDuplicate: Boolean = true) {
-        annotations.removeIf { it.range.startOffset == element.textRange.startOffset }
+        annotations.removeIf { it.range.startOffset == element.textRange.startOffset && deleteDuplicate }
 
         var textWithOutUnnecessaryData = annotationText
         for ((key, value) in unnecessaryData) {
@@ -62,6 +60,9 @@ class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : Ba
     }
 
     inner class Renderer : KtVisitorVoid() {
+        private val implicitReceivers = mutableListOf<ReceiverValue>()
+        private var lastCallWithLambda = ""
+
         private fun renderType(type: KotlinType?): String {
             return type?.let { descriptorRenderer.renderType(it) } ?: "[ERROR: unknown type]"
         }
@@ -107,8 +108,25 @@ class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : Ba
             renderVariableType(multiDeclarationEntry)
 
         override fun visitTypeReference(typeReference: KtTypeReference) {
-            val type = typeReference.getAbbreviatedTypeOrType(bindingContext)
+            if (typeReference.text.isNotEmpty()) {
-            addAnnotation(renderType(type), typeReference)
+                val hasResolvedCall = with(object : KtVisitorVoid() {
+                    var hasCall: Boolean = false
+                    override fun visitKtElement(element: KtElement) {
+                        if (!hasCall) {
+                            hasCall = element.getResolvedCall(bindingContext) != null
+                            element.acceptChildren(this)
+                        }
+                    }
+                }) {
+                    typeReference.accept(this)
+                    this.hasCall
+                }
+
+                if (!hasResolvedCall) {
+                    val type = typeReference.getAbbreviatedTypeOrType(bindingContext)
+                    addAnnotation(renderType(type), typeReference)
+                }
+            }
             super.visitTypeReference(typeReference)
         }
 
@@ -116,47 +134,41 @@ class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : Ba
             addAnnotation(renderType(expression), expression)
         }
 
-        private fun renderCall(call: Call): String {
+        private fun renderCall(expression: KtExpression): ResolvedCall<out CallableDescriptor>? {
-            val resolvedCall = bindingContext[RESOLVED_CALL, call] ?: return "[ERROR: not resolved]"
+            val call = expression.getCall(bindingContext)
+            val resolvedCall = expression.getResolvedCall(bindingContext)
+            if (call == null) {
+                return null
+            } else if (resolvedCall == null) {
+                addAnnotation("[ERROR: not resolved]", expression)
+                return null
+            }
 
             val descriptor = resolvedCall.resultingDescriptor
-            var annotation =
+            val annotation = descriptorRenderer.render(descriptor)
-                if (descriptor.isSpecial()) descriptorRenderer.renderSpecialFunction(descriptor)
+            addAnnotation(annotation, expression, deleteDuplicate = false)
-                else descriptorRenderer.render(descriptor)
-            if (descriptor is ClassConstructorDescriptor) annotation = annotation.replace("(", "<init>(")
 
-            resolvedCall.typeArguments.forEach { (type, value) ->
+            fun addReceiverAnnotation(receiver: ReceiverValue?, receiverKind: ExplicitReceiverKind) {
-                val newValue = if (value is ErrorType) "???" else "$value"
+                if (receiver != null && resolvedCall.explicitReceiverKind != receiverKind) {
-                annotation = annotation.replaceFirst(type.name.toString(), newValue)
+                    val index = implicitReceivers.indexOf(receiver)
+                    if (index != -1) {
+                        addAnnotation("this@$index", expression, deleteDuplicate = false)
+                    }
+                }
             }
-            return annotation
-        }
 
-        private fun CallableDescriptor.isSpecial(): Boolean {
+            addReceiverAnnotation(resolvedCall.extensionReceiver, ExplicitReceiverKind.EXTENSION_RECEIVER)
-            return this.name.asString().contains("SPECIAL-FUNCTION")
+            addReceiverAnnotation(resolvedCall.dispatchReceiver, ExplicitReceiverKind.DISPATCH_RECEIVER)
-        }
 
-        private fun DescriptorRenderer.renderSpecialFunction(descriptor: CallableDescriptor): String {
+            return resolvedCall
-            val descriptorName = descriptor.name.asString()
-            val name = when {
-                descriptorName.contains("ELVIS") -> "?:"
-                descriptorName.contains("EXCLEXCL") -> "!!"
-                else -> "UNKNOWN"
-            }
-            return "fun $name " + this.renderValueParameters(descriptor.valueParameters, false) +
-                    ": " + renderType(descriptor)
         }
 
         override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
-            val call = expression.getCall(bindingContext)
+            val qualifierDescriptor = bindingContext[QUALIFIER, expression]?.descriptor
-            if (call != null) {
+            if (qualifierDescriptor != null) {
-                addAnnotation(renderCall(call), expression)
+                addAnnotation(descriptorRenderer.render(qualifierDescriptor), expression)
-                super.visitSimpleNameExpression(expression)
             } else {
-                val qualifierDescriptor = bindingContext[QUALIFIER, expression]?.descriptor
+                renderCall(expression)
-                if (qualifierDescriptor != null) {
-                    addAnnotation(descriptorRenderer.render(qualifierDescriptor), expression)
-                }
             }
         }
 
@@ -174,5 +186,301 @@ class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : Ba
             addAnnotation(renderType(jetWhenEntry.expression), jetWhenEntry.expression!!)
             super.visitWhenEntry(jetWhenEntry)
         }
+
+        override fun visitLambdaExpression(lambdaExpression: KtLambdaExpression) {
+            val descriptor = bindingContext[DECLARATION_TO_DESCRIPTOR, lambdaExpression.functionLiteral] as AnonymousFunctionDescriptor
+            val extensionReceiver = descriptor.extensionReceiverParameter ?: return super.visitLambdaExpression(lambdaExpression)
+            addAnnotation("$lastCallWithLambda@${implicitReceivers.size}", lambdaExpression)
+
+            implicitReceivers += extensionReceiver.value
+            super.visitLambdaExpression(lambdaExpression)
+            implicitReceivers -= extensionReceiver.value
+        }
+
+        override fun visitCallExpression(expression: KtCallExpression) {
+            val resolvedCall = renderCall(expression) ?: return super.visitCallExpression(expression)
+
+            if (expression.getChildOfType<KtLambdaArgument>() != null) {
+                lastCallWithLambda = resolvedCall.resultingDescriptor.name.asString()
+            }
+            for (child in expression.children) {
+                if (child.node.elementType != KtNodeTypes.REFERENCE_EXPRESSION) {
+                    child.accept(this)
+                }
+            }
+        }
+    }
+
+    inner class PsiDescriptorRenderer : DeclarationDescriptorVisitor<Unit, StringBuilder> {
+        private val typeRenderer: DescriptorRenderer = DescriptorRenderer.withOptions {
+            withDefinedIn = false
+            modifiers = emptySet()
+            classifierNamePolicy = /*ClassifierNamePolicy.SOURCE_CODE_QUALIFIED*/object : ClassifierNamePolicy {
+                override fun renderClassifier(classifier: ClassifierDescriptor, renderer: DescriptorRenderer): String {
+                    return (if (classifier is TypeParameterDescriptor) renderer.renderName(classifier.name, false)
+                    else renderer.renderFqName(DescriptorUtils.getFqName(classifier))).replace(".", "/")
+                }
+            }
+            includeAdditionalModifiers = false
+            parameterNameRenderingPolicy = ParameterNameRenderingPolicy.NONE
+
+            withoutTypeParameters = true
+        }
+
+        private fun CallableDescriptor.isSpecial(): Boolean {
+            return this.name.asString().contains("SPECIAL-FUNCTION")
+        }
+
+        fun render(declarationDescriptor: DeclarationDescriptor): String {
+            if (declarationDescriptor is CallableDescriptor && declarationDescriptor.isSpecial()) {
+                return this.renderSpecialFunction(declarationDescriptor)
+            }
+            return buildString {
+                declarationDescriptor.accept(this@PsiDescriptorRenderer, this)
+            }
+        }
+
+        fun renderType(type: KotlinType): String {
+            return typeRenderer.renderType(type)
+        }
+
+        private fun renderName(descriptor: DeclarationDescriptor, hasReceiver: Boolean = false): String {
+            return if (hasReceiver) {
+                descriptor.name.asString()
+            } else {
+                renderFqName(descriptor)
+            }
+        }
+
+        private fun renderFqName(descriptor: DeclarationDescriptor): String {
+            /*generateSequence(descriptor) { it.containingDeclaration }
+                .takeWhile {  }*/
+            return descriptor.fqNameUnsafe.render().let {
+                if (descriptor !is PropertyDescriptor && descriptor !is VariableDescriptor)
+                    it.replace(".", "/")
+                else it
+            }
+        }
+
+        private fun renderReceiver(descriptor: CallableDescriptor, data: StringBuilder): ReceiverParameterDescriptor? {
+            return descriptor.extensionReceiverParameter?.also {
+                visitReceiverParameterDescriptor(it, data)
+                data.append(".")
+            } ?: descriptor.dispatchReceiverParameter?.also {
+                data.append("(")
+                visitReceiverParameterDescriptor(it, data)
+                data.append(").")
+            }
+        }
+
+        private fun renderSuperTypes(klass: ClassDescriptor, builder: StringBuilder) {
+            if (KotlinBuiltIns.isNothing(klass.defaultType)) return
+
+            val supertypes = klass.typeConstructor.supertypes
+            if (supertypes.isEmpty() || supertypes.size == 1 && KotlinBuiltIns.isAnyOrNullableAny(supertypes.iterator().next())) return
+
+            builder.append(": ")
+            supertypes.joinTo(builder, ", ") { renderType(it) }
+        }
+
+        private fun renderValueParameter(
+            parameter: ValueParameterDescriptor, includeNames: Boolean, data: StringBuilder, topLevel: Boolean
+        ) {
+            renderVariable(parameter, includeNames, data, topLevel)
+            if (parameter.declaresOrInheritsDefaultValue()) {
+                data.append(" = ...")
+            }
+        }
+
+        private fun renderVariable(variable: VariableDescriptor, includeName: Boolean, data: StringBuilder, topLevel: Boolean) {
+            val realType = variable.type
+
+            val varargElementType = (variable as? ValueParameterDescriptor)?.varargElementType
+            val typeToRender = varargElementType ?: realType
+            if (varargElementType != null) data.append("vararg ")
+
+            if (topLevel && variable !is ValueParameterDescriptor) data.append(if (variable.isVar) "var" else "val").append(" ")
+            if (includeName) {
+                data.append(renderName(variable)).append(": ")
+            }
+
+            data.append(renderType(typeToRender))
+        }
+
+        private fun renderSpecialFunction(descriptor: CallableDescriptor): String {
+            val descriptorName = descriptor.name.asString()
+            val name = when {
+                descriptorName.contains("ELVIS") -> "?:"
+                descriptorName.contains("EXCLEXCL") -> "!!"
+                else -> "UNKNOWN"
+            }
+            val valueParameters = buildString { visitValueParameters(descriptor.valueParameters, this) }
+            val returnType = descriptor.returnType?.let { renderType(it) } ?: "[ERROR: unknown type]"
+
+            return "fun $name $valueParameters: $returnType"
+        }
+
+        override fun visitPackageFragmentDescriptor(descriptor: PackageFragmentDescriptor, data: StringBuilder) {
+            data.append("package-fragment ${renderFqName(descriptor)}")
+        }
+
+        override fun visitPackageViewDescriptor(descriptor: PackageViewDescriptor, data: StringBuilder) {
+            data.append("package ${renderFqName(descriptor)}")
+        }
+
+        override fun visitVariableDescriptor(variable: VariableDescriptor, data: StringBuilder) {
+            renderVariable(variable, true, data, true)
+        }
+
+        override fun visitFunctionDescriptor(function: FunctionDescriptor, data: StringBuilder) {
+            data.append("fun ")
+
+            visitTypeParameters(function.typeParameters, data)
+            if (function.typeParameters.isNotEmpty()) data.append(" ")
+
+            //render receiver
+            val receiver = renderReceiver(function, data)
+
+            //render name
+            data.append(renderName(function, receiver != null))
+
+            //render value parameters
+            visitValueParameters(function.valueParameters, data)
+
+            //render return type
+            val returnType = function.returnType
+            data.append(": ").append(if (returnType == null) "[NULL]" else renderType(returnType))
+
+            renderWhereSuffix(function.typeParameters, data)
+        }
+
+        private fun renderWhereSuffix(typeParameters: List<TypeParameterDescriptor>, data: StringBuilder) {
+            val upperBoundStrings = ArrayList<String>(0)
+
+            for (typeParameter in typeParameters) {
+                typeParameter.upperBounds
+                    .drop(1) // first parameter is rendered by renderTypeParameter
+                    .mapTo(upperBoundStrings) { renderName(typeParameter) + " : " + renderType(it) }
+            }
+
+            if (upperBoundStrings.isNotEmpty()) {
+                data.append(" where ")
+                upperBoundStrings.joinTo(data, ", ")
+            }
+        }
+
+        private fun visitTypeParameters(typeParameters: List<TypeParameterDescriptor>, data: StringBuilder) {
+            if (typeParameters.isNotEmpty()) {
+                data.append("<")
+                val iterator = typeParameters.iterator()
+                while (iterator.hasNext()) {
+                    val typeParameterDescriptor = iterator.next()
+                    visitTypeParameterDescriptor(typeParameterDescriptor, data)
+                    if (iterator.hasNext()) {
+                        data.append(", ")
+                    }
+                }
+                data.append(">")
+            }
+        }
+
+        override fun visitTypeParameterDescriptor(typeParameter: TypeParameterDescriptor, data: StringBuilder) {
+            data.append(renderName(typeParameter, true))
+            val upperBoundsCount = typeParameter.upperBounds.size
+            if (upperBoundsCount >= 1) {
+                val upperBound = typeParameter.upperBounds.iterator().next()
+                if (!KotlinBuiltIns.isDefaultBound(upperBound)) {
+                    data.append(" : ").append(renderType(upperBound))
+                }
+            }
+        }
+
+        override fun visitClassDescriptor(klass: ClassDescriptor, data: StringBuilder) {
+            data.append(DescriptorRenderer.getClassifierKindPrefix(klass)).append(" ")
+
+            //render name
+            data.append(renderName(klass))
+
+            if (klass.kind == ClassKind.ENUM_ENTRY) return
+
+            visitTypeParameters(klass.declaredTypeParameters, data)
+
+            renderSuperTypes(klass, data)
+            renderWhereSuffix(klass.declaredTypeParameters, data)
+        }
+
+        override fun visitTypeAliasDescriptor(typeAlias: TypeAliasDescriptor, data: StringBuilder) {
+            data.append("typealias").append(" ")
+            data.append(renderName(typeAlias))
+
+            visitTypeParameters(typeAlias.declaredTypeParameters, data)
+
+            data.append(" = ").append(renderType(typeAlias.underlyingType))
+        }
+
+        override fun visitModuleDeclaration(descriptor: ModuleDescriptor, data: StringBuilder) {
+            data.append(renderName(descriptor))
+        }
+
+        override fun visitConstructorDescriptor(constructor: ConstructorDescriptor, data: StringBuilder) {
+            data.append("constructor").append(" ")
+
+            val classDescriptor = constructor.containingDeclaration
+
+            data.append(renderName(classDescriptor))
+            visitTypeParameters(constructor.typeParameters, data)
+
+            visitValueParameters(constructor.valueParameters, data)
+
+            renderWhereSuffix(constructor.typeParameters, data)
+        }
+
+        override fun visitScriptDescriptor(scriptDescriptor: ScriptDescriptor, data: StringBuilder) {
+            visitClassDescriptor(scriptDescriptor, data)
+        }
+
+        override fun visitPropertyDescriptor(property: PropertyDescriptor, data: StringBuilder) {
+            data.append(if (property.isVar) "var" else "val").append(" ")
+            visitTypeParameters(property.typeParameters, data)
+            if (property.typeParameters.isNotEmpty()) data.append(" ")
+
+            //render receiver
+            val receiver = renderReceiver(property, data)
+
+            //render name
+            data.append(renderName(property, receiver != null))
+
+            //render return type
+            data.append(": ").append(renderType(property.type))
+
+            renderWhereSuffix(property.typeParameters, data)
+        }
+
+        private fun visitValueParameters(parameters: List<ValueParameterDescriptor>, data: StringBuilder) {
+            data.append("(")
+            for ((index, parameter) in parameters.withIndex()) {
+                renderValueParameter(parameter, false, data, false)
+                if (index != parameters.size - 1) {
+                    data.append(", ")
+                }
+            }
+            data.append(")")
+        }
+
+        override fun visitValueParameterDescriptor(parameter: ValueParameterDescriptor, data: StringBuilder) {
+            renderValueParameter(parameter, true, data, true)
+        }
+
+        override fun visitPropertyGetterDescriptor(descriptor: PropertyGetterDescriptor, data: StringBuilder) {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+
+        override fun visitPropertySetterDescriptor(descriptor: PropertySetterDescriptor, data: StringBuilder) {
+            TODO("not implemented") //To change body of created functions use File | Settings | File Templates.
+        }
+
+        override fun visitReceiverParameterDescriptor(descriptor: ReceiverParameterDescriptor, data: StringBuilder) {
+            data.append(renderType(descriptor.type))
+        }
     }
 }