Remove unnecessary data from render result

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

@@ -8,18 +8,14 @@ package org.jetbrains.kotlin.compiler.visualizer
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analyzer.AnalysisResult
-import org.jetbrains.kotlin.builtins.*
 import org.jetbrains.kotlin.compiler.visualizer.Annotator.annotate
 import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 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.calls.callUtil.getCall
-import org.jetbrains.kotlin.resolve.calls.callUtil.getType
 import org.jetbrains.kotlin.types.*
 
 class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : BaseRenderer {
@@ -41,9 +37,21 @@ class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : Ba
         parameterNameRenderingPolicy = ParameterNameRenderingPolicy.NONE
     }
 
-    private fun addAnnotation(text: String, range: TextRange) {
+    private val unnecessaryData = mapOf(
-        annotations.removeIf { it.range.startOffset == range.startOffset }
+        "kotlin." to "",
-        annotations.add(Annotator.AnnotationInfo(text, range))
+        "kotlin.collections." to ""
+    )
+
+    private fun addAnnotation(annotationText: String, element: PsiElement) {
+        annotations.removeIf { it.range.startOffset == element.textRange.startOffset }
+
+        var textWithOutUnnecessaryData = annotationText
+        for ((key, value) in unnecessaryData) {
+            textWithOutUnnecessaryData = textWithOutUnnecessaryData.replace(key, value)
+        }
+        if (textWithOutUnnecessaryData != element.text) {
+            annotations.add(Annotator.AnnotationInfo(textWithOutUnnecessaryData, element.textRange))
+        }
     }
 
     override fun render(): String {
@@ -74,7 +82,7 @@ class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : Ba
 
         private fun renderVariableType(variable: KtVariableDeclaration) {
             val descriptor = bindingContext[VARIABLE, variable]
-            addAnnotation(renderType(descriptor), variable.nameIdentifier!!.textRange)
+            addAnnotation(renderType(descriptor), variable.nameIdentifier!!)
             variable.acceptChildren(this)
         }
 
@@ -86,12 +94,12 @@ class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : Ba
 
         override fun visitTypeReference(typeReference: KtTypeReference) {
             val type = bindingContext[TYPE, typeReference]
-            addAnnotation(renderType(type), typeReference.textRange)
+            addAnnotation(renderType(type), typeReference)
             super.visitTypeReference(typeReference)
         }
 
         override fun visitConstantExpression(expression: KtConstantExpression) {
-            addAnnotation(renderType(expression), expression.textRange)
+            addAnnotation(renderType(expression), expression)
         }
 
         private fun renderCall(call: Call): String {
@@ -109,30 +117,28 @@ class PsiRenderer(private val file: KtFile, analysisResult: AnalysisResult) : Ba
         override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
             val call = expression.getCall(bindingContext)
             if (call != null) {
-                addAnnotation(renderCall(call), expression.textRange)
+                addAnnotation(renderCall(call), expression)
                 super.visitSimpleNameExpression(expression)
             } else {
                 val qualifierDescriptor = bindingContext[QUALIFIER, expression]?.descriptor
                 if (qualifierDescriptor != null) {
-                    addAnnotation(descriptorRenderer.render(qualifierDescriptor), expression.textRange)
+                    addAnnotation(descriptorRenderer.render(qualifierDescriptor), expression)
                 }
             }
         }
 
         override fun visitIfExpression(expression: KtIfExpression) {
-            addAnnotation(renderType(expression), expression.ifKeyword.textRange)
+            addAnnotation(renderType(expression), expression.ifKeyword)
             super.visitIfExpression(expression)
         }
 
         override fun visitWhenExpression(expression: KtWhenExpression) {
-            addAnnotation(renderType(expression), expression.whenKeyword.textRange)
+            addAnnotation(renderType(expression), expression.whenKeyword)
             super.visitWhenExpression(expression)
         }
 
         override fun visitWhenEntry(jetWhenEntry: KtWhenEntry) {
-            val arrowTextRange = jetWhenEntry.expression?.textRange?.shiftLeft(3)
+            addAnnotation(renderType(jetWhenEntry.expression), jetWhenEntry.expression!!)
-                ?: jetWhenEntry.conditions.last().textRange.shiftRight(3)
-            addAnnotation(renderType(jetWhenEntry.expression), arrowTextRange)
             super.visitWhenEntry(jetWhenEntry)
         }
     }