[VISUALIZER] Add annotation at Array.set method

compiler/visualizer/render-fir/src/org/jetbrains/kotlin/compiler/visualizer/FirVisualizer.kt

@@ -151,7 +151,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
         override fun visitNamedFunction(function: KtNamedFunction) {
             stack.push((function.name ?: "<no name provided>"))
             if (function.equalsToken != null) {
-                function.bodyExpression!!.firstOfTypeWithRender<FirExpression>(function.equalsToken) { this.typeRef }
+                function.bodyExpression!!.firstOfTypeWithRender<FirReturnExpression>(function.equalsToken) { this.result.typeRef }
                     ?: function.firstOfTypeWithRender<FirTypedDeclaration>(function.equalsToken) { this.returnTypeRef }
             }
             super.visitNamedFunction(function)
@@ -280,6 +280,13 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
         private val filePackage = firFile.packageFqName.toString().replace(".", "/")
         private val symbolProvider = firFile.session.symbolProvider
 
+        private fun FirTypeRef.renderWithNativeRenderer(): String {
+            return buildString {
+                val nativeRenderer = org.jetbrains.kotlin.fir.FirRenderer(this)
+                this@renderWithNativeRenderer.accept(nativeRenderer)
+            }.replace("R|", "").replace("|", "")
+        }
+
         private fun removeCurrentFilePackage(fqName: String): String {
             return if (fqName.startsWith(filePackage) && !fqName.substring(filePackage.length + 1).contains("/")) {
                 fqName.replaceFirst("$filePackage/", "")
@@ -479,8 +486,8 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
         }
 
         override fun visitResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef, data: StringBuilder) {
+            data.append(resolvedTypeRef.renderWithNativeRenderer())
             val coneType = resolvedTypeRef.type
-            data.append(removeCurrentFilePackage(coneType.render()))
             if (coneType is ConeClassLikeType) {
                 val original = coneType.directExpansionType(session)
                 original?.let { data.append(" /* = ${it.render()} */") }
@@ -528,11 +535,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
                     val callableName = symbol.callableId.callableName
                     val receiverType = symbol.fir.receiverTypeRef
                     if (receiverType == null) {
-                        if (symbol.callableId.className == null) {
-                            data.append(id)
-                        } else {
-                            data.append("($id).$callableName")
-                        }
+                        symbol.callableId.className?.let { data.append("($id).$callableName") } ?: data.append(id)
                     } else {
                         data.append("${receiverType.render()}.$callableName")
                     }
@@ -557,7 +560,8 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
 
                 }
                 is FirVariableSymbol<*> -> {
-                    renderVariable(symbol.fir)                }
+                    renderVariable(symbol.fir)
+                }
                 is FirConstructorSymbol -> {
                     data.append("constructor ")
                     val packageName = symbol.callableId.className

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

@@ -12,7 +12,6 @@ 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.descriptors.impl.LazyPackageViewDescriptorImpl
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
@@ -21,7 +20,6 @@ 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
@@ -30,10 +28,9 @@ 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.lazy.descriptors.LazyPackageDescriptor
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.types.*
-import java.util.ArrayList
+import java.util.*
 
 class PsiVisualizer(private val file: KtFile, analysisResult: AnalysisResult) : BaseRenderer() {
     private val bindingContext = analysisResult.bindingContext
@@ -136,7 +133,7 @@ class PsiVisualizer(private val file: KtFile, analysisResult: AnalysisResult) :
             addAnnotation(renderType(expression), expression)
         }
 
-        private fun renderCall(expression: KtExpression): ResolvedCall<out CallableDescriptor>? {
+        private fun renderCall(expression: KtExpression, renderOn: PsiElement = expression): ResolvedCall<out CallableDescriptor>? {
             val call = expression.getCall(bindingContext)
             val resolvedCall = expression.getResolvedCall(bindingContext)
             if (call == null) {
@@ -147,11 +144,11 @@ class PsiVisualizer(private val file: KtFile, analysisResult: AnalysisResult) :
             }
 
             val descriptor = resolvedCall.resultingDescriptor
-            val typeArguments = if (resolvedCall.typeArguments.isNotEmpty()) {
-                buildString { resolvedCall.typeArguments.values.joinTo(this, ", ", "<", ">") { renderType(it) } }
-            } else ""
+            val typeArguments = resolvedCall.typeArguments
+                .takeIf { it.isNotEmpty() }
+                ?.values?.joinToString(", ", "<", ">") { renderType(it) } ?: ""
             val annotation = descriptorRenderer.render(descriptor).replace(argumentsLabel, typeArguments)
-            addAnnotation(annotation, expression, deleteDuplicate = false)
+            addAnnotation(annotation, renderOn, deleteDuplicate = false)
 
             fun addReceiverAnnotation(receiver: ReceiverValue?, receiverKind: ExplicitReceiverKind) {
                 if (receiver != null && resolvedCall.explicitReceiverKind != receiverKind) {
@@ -177,6 +174,14 @@ class PsiVisualizer(private val file: KtFile, analysisResult: AnalysisResult) :
             }
         }
 
+        override fun visitBinaryExpression(expression: KtBinaryExpression) {
+            val opName = expression.operationReference.getReferencedName()
+            expression.left?.takeIf { it.node.elementType == KtNodeTypes.ARRAY_ACCESS_EXPRESSION && opName == "=" }?.let {
+                renderCall(it, expression.operationReference)
+            }
+            super.visitBinaryExpression(expression)
+        }
+
         override fun visitIfExpression(expression: KtIfExpression) {
             addAnnotation(renderType(expression), expression.ifKeyword)
             super.visitIfExpression(expression)

compiler/visualizer/testData/rawBuilder/expressions/arrayAssignment.kt

@@ -7,8 +7,10 @@ fun test() {
 //      │   │          │  │  │
     val x = intArrayOf(1, 2, 3)
 //  val test.x: IntArray
-//  │ Int  Int
-//  │ │    │
+//  │ Int
+//  │ │  fun (IntArray).set(Int, Int): Unit
+//  │ │  │ Int
+//  │ │  │ │
     x[1] = 0
 }
 
@@ -23,7 +25,8 @@ fun test2() {
 //  │          │  Int
 //  │          │  │  Int
 //  │          │  │  │  fun foo(): Int
-//  │          │  │  │  │        Int
-//  │          │  │  │  │        │
+//  │          │  │  │  │      fun (IntArray).set(Int, Int): Unit
+//  │          │  │  │  │      │ Int
+//  │          │  │  │  │      │ │
     intArrayOf(1, 2, 3)[foo()] = 1
 }