[VISUALIZER] Add implicit receiver rendering in fir

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.compiler.visualizer
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirLabel
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
@@ -24,13 +25,17 @@ import org.jetbrains.kotlin.fir.visitors.FirVisitor
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
 import org.jetbrains.kotlin.psi.psiUtil.getChildrenOfType
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.util.capitalizeDecapitalize.toLowerCaseAsciiOnly
 
 private typealias Stack = MutableList<Pair<String, MutableList<String>>>
 
 class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
+    private val implicitReceivers = mutableListOf<FirTypeRef>()
+
     private fun FirElement.render(): String = buildString { this@render.accept(FirRenderer(), this) }
 
     override fun addAnnotation(annotationText: String, element: PsiElement?, deleteDuplicate: Boolean) {
@@ -44,7 +49,8 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
     }
 
     inner class PsiVisitor(private val map: Map<PsiElement, MutableList<FirElement>>) : KtVisitorVoid() {
-        //private var innerLambdaCount = 0
+        private var lastCallWithLambda: String? = null
+
         private val stack = mutableListOf("" to mutableListOf<String>())
 
         private fun Stack.push(
@@ -246,17 +252,26 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
 
         override fun visitCallExpression(expression: KtCallExpression) {
             expression.firstOfTypeWithLocalReplace<FirFunctionCall> { this.calleeReference.name.asString() }
-            expression.children.filter { it.node.elementType != KtNodeTypes.REFERENCE_EXPRESSION }.forEach { it.accept(this) }
+            expression.children.filter { it.node.elementType != KtNodeTypes.REFERENCE_EXPRESSION }.forEach { psi ->
+                when (psi) {
+                    is KtLambdaArgument -> {
+                        val firLambda = psi.firstOfType<FirLambdaArgumentExpression>()?.expression as FirAnonymousFunction
+                        firLambda.receiverTypeRef?.let {
+                            lastCallWithLambda = psi.getLambdaExpression()?.firstOfType<FirLabel>()?.name
+                            implicitReceivers += it
+                            psi.accept(this)
+                            implicitReceivers -= it
+                        } ?: psi.accept(this)
+                    }
+                    else -> psi.accept(this)
+                }
+            }
         }
 
         override fun visitLambdaExpression(lambdaExpression: KtLambdaExpression) {
-            //val firLabel = lambdaExpression.firstOfType<FirLabel>()
-
             stack.push("<anonymous>")
-            //firLabel?.let { addAnnotation("${it.name}@$innerLambdaCount", lambdaExpression) }
+            lastCallWithLambda?.let { addAnnotation("$it@${implicitReceivers.size - 1}", lambdaExpression) }
-            //innerLambdaCount++
             super.visitLambdaExpression(lambdaExpression)
-            //innerLambdaCount--
             stack.pop()
         }
 
@@ -327,6 +342,15 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
             }
         }
 
+        private fun renderImplicitReceiver(symbol: AbstractFirBasedSymbol<*>, psi: PsiElement?) {
+            val implicitReceiverIndex = (symbol.fir as? FirCallableMemberDeclaration<*>)?.dispatchReceiverType?.let {
+                implicitReceivers.map { (it as? FirResolvedTypeRef)?.coneType }.indexOf(it)
+            } ?: return
+            if (implicitReceiverIndex != -1) {
+                addAnnotation("this@$implicitReceiverIndex", psi)
+            }
+        }
+
         private fun renderConstructorSymbol(symbol: FirConstructorSymbol, data: StringBuilder) {
             data.append("constructor ")
             data.append(getSymbolId(symbol))
@@ -383,22 +407,18 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
                 return
             }
 
-            when (call.explicitReceiver) {
+            when {
-                null -> data.append(id)
+                call.dispatchReceiver !is FirNoReceiverExpression -> {
-                else -> {
+                    data.append("(")
-                    when {
+                    call.dispatchReceiver.typeRef.accept(this, data)
-                        call.dispatchReceiver !is FirNoReceiverExpression -> {
+                    data.append(").").append(symbol.callableId.callableName)
-                            data.append("(")
+                }
-                            call.dispatchReceiver.typeRef.accept(this, data)
+                call.extensionReceiver !is FirNoReceiverExpression -> {
-                            data.append(")")
+                    // render type from symbol because this way it will be consistent with psi render
-                        }
+                    symbol.fir.receiverTypeRef?.accept(this, data)
-                        call.extensionReceiver !is FirNoReceiverExpression -> {
-                            // render type from symbol because this way it will be consistent with psi render
-                            symbol.fir.receiverTypeRef?.accept(this, data)
-                        }
-                    }
                     data.append(".").append(symbol.callableId.callableName)
                 }
+                else -> data.append(id)
             }
 
             renderListInTriangles(call.typeArguments, data)
@@ -477,6 +497,7 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
 
         override fun visitResolvedNamedReference(resolvedNamedReference: FirResolvedNamedReference, data: StringBuilder) {
             val symbol = resolvedNamedReference.resolvedSymbol
+            renderImplicitReceiver(symbol, resolvedNamedReference.source.psi)
             when {
                 symbol is FirPropertySymbol && !symbol.fir.isLocal -> renderPropertySymbol(symbol, data)
                 symbol is FirNamedFunctionSymbol -> {
@@ -552,6 +573,9 @@ class FirVisualizer(private val firFile: FirFile) : BaseRenderer() {
         override fun visitFunctionCall(functionCall: FirFunctionCall, data: StringBuilder) {
             when (val callee = functionCall.calleeReference) {
                 is FirResolvedNamedReference -> {
+                    if (functionCall.explicitReceiver == null) {
+                        renderImplicitReceiver(callee.resolvedSymbol, functionCall.source.psi)
+                    }
                     when (callee.resolvedSymbol) {
                         is FirConstructorSymbol -> {
                             renderConstructorSymbol(callee.resolvedSymbol as FirConstructorSymbol, data)

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

@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.compiler.visualizer.Annotator.annotate
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getChildOfType
 import org.jetbrains.kotlin.renderer.ClassifierNamePolicy
@@ -149,13 +148,6 @@ class PsiVisualizer(private val file: KtFile, analysisResult: AnalysisResult) :
                 return null
             }
 
-            val descriptor = resolvedCall.candidateDescriptor
-            val typeArguments = resolvedCall.typeArguments
-                .takeIf { it.isNotEmpty() }
-                ?.values?.joinToString(", ", "<", ">") { renderType(it) } ?: ""
-            val annotation = descriptorRenderer.render(descriptor).replace(argumentsLabel, typeArguments)
-            addAnnotation(annotation, renderOn, deleteDuplicate = false)
-
             fun addReceiverAnnotation(receiver: ReceiverValue?, receiverKind: ExplicitReceiverKind) {
                 if (receiver != null && resolvedCall.explicitReceiverKind != receiverKind) {
                     val index = implicitReceivers.indexOf(receiver)
@@ -168,6 +160,13 @@ class PsiVisualizer(private val file: KtFile, analysisResult: AnalysisResult) :
             addReceiverAnnotation(resolvedCall.extensionReceiver, ExplicitReceiverKind.EXTENSION_RECEIVER)
             addReceiverAnnotation(resolvedCall.dispatchReceiver, ExplicitReceiverKind.DISPATCH_RECEIVER)
 
+            val descriptor = resolvedCall.candidateDescriptor
+            val typeArguments = resolvedCall.typeArguments
+                .takeIf { it.isNotEmpty() }
+                ?.values?.joinToString(", ", "<", ">") { renderType(it) } ?: ""
+            val annotation = descriptorRenderer.render(descriptor).replace(argumentsLabel, typeArguments)
+            addAnnotation(annotation, renderOn, deleteDuplicate = false)
+
             return resolvedCall
         }
 

compiler/visualizer/testData/rawBuilder/declarations/contracts/oldSyntax/contractDescription.kt

@@ -18,12 +18,12 @@ fun test_2() {
 //                   │        contract@0
 //                   │        │
     kotlin.contracts.contract {
-//      this@0
 //      fun <R> (contracts/ContractBuilder).callsInPlace<???>(Function<R>, contracts/InvocationKind = ...): contracts/CallsInPlace
+//      this@0
 //      │
         callsInPlace()
-//      this@0
 //      fun <R> (contracts/ContractBuilder).callsInPlace<???>(Function<R>, contracts/InvocationKind = ...): contracts/CallsInPlace
+//      this@0
 //      │
         callsInPlace()
     }
@@ -54,12 +54,12 @@ var test_3: Int = 1
 //                       │        contract@0
 //                       │        │
         kotlin.contracts.contract {
-//          this@0
 //          fun <R> (contracts/ContractBuilder).callsInPlace<???>(Function<R>, contracts/InvocationKind = ...): contracts/CallsInPlace
+//          this@0
 //          │
             callsInPlace()
-//          this@0
 //          fun <R> (contracts/ContractBuilder).callsInPlace<???>(Function<R>, contracts/InvocationKind = ...): contracts/CallsInPlace
+//          this@0
 //          │
             callsInPlace()
         }

compiler/visualizer/testData/uncommonCases/resultFiles/innerWith.kt

@@ -5,6 +5,7 @@ class A {
 //      Int     Int
 //      │       │
     val aProp = 10
+    fun call() {}
 }
 
 class B {
@@ -19,26 +20,30 @@ fun foo(a: Int, b: Int): Int {
 //  │    │    with@0
 //  │    │    │
     with(A()) {
-//      this@0
 //      val (A).aProp: Int
+//      this@0
 //      │
         aProp
+//      fun (A).call(): Unit
+//      this@0
+//      │
+        call()
 
 //      fun <T, R> with<B, Int>(T, T.() -> R): R
 //      │    constructor B()
 //      │    │    with@1
 //      │    │    │
         with(B()) {
-//          this@0
 //          val (A).aProp: Int
+//          this@0
 //          │
             aProp
-//          this@1
 //          val (B).bProp: Int
+//          this@1
 //          │
             bProp
-//          this@0
 //          val (A).aProp: Int
+//          this@0
 //          │
             aProp
         }
@@ -49,8 +54,8 @@ fun foo(a: Int, b: Int): Int {
 //  │    │    with@0
 //  │    │    │
     with(A()) {
-//      this@0
 //      val (A).aProp: Int
+//      this@0
 //      │
         aProp
 
@@ -59,12 +64,12 @@ fun foo(a: Int, b: Int): Int {
 //      │    │    with@1
 //      │    │    │
         with(B()) {
-//          this@0
 //          val (A).aProp: Int
+//          this@0
 //          │
             aProp
-//          this@1
 //          val (B).bProp: Int
+//          this@1
 //          │
             bProp
         }
@@ -74,12 +79,12 @@ fun foo(a: Int, b: Int): Int {
 //      │    │    with@1
 //      │    │    │
         with(B()) {
-//          this@0
 //          val (A).aProp: Int
+//          this@0
 //          │
             aProp
-//          this@1
 //          val (B).bProp: Int
+//          this@1
 //          │
             bProp
         }

compiler/visualizer/testData/uncommonCases/testFiles/innerWith.kt

@@ -2,6 +2,7 @@ package p
 
 class A {
     val aProp = 10
+    fun call() {}
 }
 
 class B {
@@ -11,6 +12,7 @@ class B {
 fun foo(a: Int, b: Int): Int {
     with(A()) {
         aProp
+        call()
 
         with(B()) {
             aProp