From 7f9ed58c7d0951e456aa626b4742419d86f5ca5d Mon Sep 17 00:00:00 2001
From: Mark Punzalan <markpun@google.com>
Date: Fri, 6 Aug 2021 23:50:21 +0000
Subject: [PATCH] FIR IDE: Expand KtCall API to include array access calls.

---
 .../kotlin/idea/frontend/api/components/KtCallResolver.kt  | 5 +++++
 .../idea/frontend/api/fir/components/KtFirCallResolver.kt  | 7 +++++++
 2 files changed, 12 insertions(+)

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 d256b413dbd..7e4aa9b80a5 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
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.idea.frontend.api.components
 
 import org.jetbrains.kotlin.idea.frontend.api.calls.KtCall
+import org.jetbrains.kotlin.psi.KtArrayAccessExpression
 import org.jetbrains.kotlin.psi.KtBinaryExpression
 import org.jetbrains.kotlin.psi.KtCallElement
 import org.jetbrains.kotlin.psi.KtUnaryExpression
@@ -14,6 +15,7 @@ public abstract class KtCallResolver : KtAnalysisSessionComponent() {
     public abstract fun resolveCall(call: KtCallElement): KtCall?
     public abstract fun resolveCall(call: KtBinaryExpression): KtCall?
     public abstract fun resolveCall(call: KtUnaryExpression): KtCall?
+    public abstract fun resolveCall(call: KtArrayAccessExpression): KtCall?
 }
 
 public interface KtCallResolverMixIn : KtAnalysisSessionMixIn {
@@ -25,4 +27,7 @@ public interface KtCallResolverMixIn : KtAnalysisSessionMixIn {
 
     public fun KtUnaryExpression.resolveCall(): KtCall? =
         analysisSession.callResolver.resolveCall(this)
+
+    public fun KtArrayAccessExpression.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 d8fb28ba63f..71d8c07bac8 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
@@ -78,6 +78,13 @@ internal class KtFirCallResolver(
         }
     }
 
+    override fun resolveCall(call: KtArrayAccessExpression): KtCall? = withValidityAssertion {
+        return when (val fir = call.getOrBuildFir(firResolveState)) {
+            is FirFunctionCall -> resolveCall(fir)
+            else -> null
+        }
+    }
+
     private fun resolveCall(firCall: FirFunctionCall): KtCall? {
         val session = firResolveState.rootModuleSession
         return when {