From ee023038169ad1b6f92a323b1283b3ce0fcdf046 Mon Sep 17 00:00:00 2001
From: Jinseong Jeon <jsjeon@google.com>
Date: Thu, 17 Jun 2021 21:54:41 +0200
Subject: [PATCH] FIR IDE/UAST: resolve call from KtUnaryExpression

---
 .../frontend/api/components/KtCallResolver.kt     |  6 +++++-
 .../api/fir/components/KtFirCallResolver.kt       | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtCallResolver.kt b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtCallResolver.kt
index ee27e0b8d4e..9a40fe43916 100644
--- a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtCallResolver.kt
+++ b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtCallResolver.kt
@@ -8,10 +8,12 @@ package org.jetbrains.kotlin.idea.frontend.api.components
 import org.jetbrains.kotlin.idea.frontend.api.calls.KtCall
 import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtUnaryExpression
 
 abstract class KtCallResolver : KtAnalysisSessionComponent() {
     abstract fun resolveCall(call: KtCallExpression): KtCall?
     abstract fun resolveCall(call: KtBinaryExpression): KtCall?
+    abstract fun resolveCall(call: KtUnaryExpression): KtCall?
 }
 
 interface KtCallResolverMixIn : KtAnalysisSessionMixIn {
@@ -21,4 +23,6 @@ interface KtCallResolverMixIn : KtAnalysisSessionMixIn {
     fun KtBinaryExpression.resolveCall(): KtCall? =
         analysisSession.callResolver.resolveCall(this)
 
-}
\ No newline at end of file
+    fun KtUnaryExpression.resolveCall(): KtCall? =
+        analysisSession.callResolver.resolveCall(this)
+}
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCallResolver.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCallResolver.kt
index 709ace88a17..25d7694954a 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCallResolver.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirCallResolver.kt
@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.idea.references.FirReferenceResolveHelper
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtUnaryExpression
 import org.jetbrains.kotlin.util.OperatorNameConventions
 
 internal class KtFirCallResolver(
@@ -46,6 +47,20 @@ internal class KtFirCallResolver(
         }
     }
 
+    override fun resolveCall(call: KtUnaryExpression): KtCall? = withValidityAssertion {
+        when (val fir = call.getOrBuildFir(firResolveState)) {
+            is FirFunctionCall -> resolveCall(fir)
+            is FirBlock -> {
+                // Desugared increment or decrement block. See [BaseFirBuilder#generateIncrementOrDecrementBlock]
+                // There would be corresponding inc()/dec() call that is assigned back to a temp variable.
+                val prefix = fir.statements.filterIsInstance<FirVariableAssignment>().find { it.rValue is FirFunctionCall }
+                (prefix?.rValue as? FirFunctionCall)?.let { resolveCall(it) }
+            }
+            is FirCheckNotNullCall -> null // TODO
+            else -> null
+        }
+    }
+
     override fun resolveCall(call: KtCallExpression): KtCall? = withValidityAssertion {
         val firCall = when (val fir = call.getOrBuildFir(firResolveState)) {
             is FirFunctionCall -> fir