FIR IDE: introduce KtFirArrayAccessReference

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1790,16 +1790,17 @@ class RawFirBuilder(
                 if (getArgument != null) {
                     calleeReference = buildSimpleNamedReference {
                         source = expression.parent.toFirSourceElement()
-                        this.source = source
+                        this.source = source.fakeElement(FirFakeSourceElementKind.ArrayAccessNameReference)
                         name = OperatorNameConventions.SET
                     }
                 } else {
                     source = expression.toFirSourceElement()
                     calleeReference = buildSimpleNamedReference {
-                        this.source = source
+                        this.source = source.fakeElement(FirFakeSourceElementKind.ArrayAccessNameReference)
                         name = OperatorNameConventions.GET
                     }
                 }
+                this.source = source
                 explicitReceiver = arrayExpression.toFirExpression("No array expression")
                 argumentList = buildArgumentList {
                     for (indexExpression in expression.indexExpressions) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSourceElement.kt

@@ -139,6 +139,9 @@ sealed class FirFakeSourceElementKind : FirSourceElementKind() {
     // where `val $subj = a` has a fake source
     object WhenGeneratedSubject : FirFakeSourceElementKind()
 
+    // list[0] -> list.get(0) where name reference will have a fake source element
+    object ArrayAccessNameReference : FirFakeSourceElementKind()
+
     // super.foo() --> super<Supertype>.foo()
     // where `Supertype` has a fake source
     object SuperCallImplicitType : FirFakeSourceElementKind()

idea/idea-fir/tests/org/jetbrains/kotlin/idea/resolve/FirReferenceResolveTestGenerated.java

@@ -418,6 +418,29 @@ public class FirReferenceResolveTestGenerated extends AbstractFirReferenceResolv
         runTest("idea/testData/resolve/references/WrongNumberOfTypeArgumentsInSupertype.kt");
     }
 
+    @TestMetadata("idea/testData/resolve/references/arrayAccess")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ArrayAccess extends AbstractFirReferenceResolveTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInArrayAccess() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/resolve/references/arrayAccess"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("get.kt")
+        public void testGet() throws Exception {
+            runTest("idea/testData/resolve/references/arrayAccess/get.kt");
+        }
+
+        @TestMetadata("set.kt")
+        public void testSet() throws Exception {
+            runTest("idea/testData/resolve/references/arrayAccess/set.kt");
+        }
+    }
+
     @TestMetadata("idea/testData/resolve/references/delegatedPropertyAccessors")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirIdeDeserializedDeclarationSourceProvider.kt

@@ -116,7 +116,7 @@ object FirIdeDeserializedDeclarationSourceProvider {
     }
 
     private fun FirCallableDeclaration<*>.containingKtClass(project: Project): KtClassOrObject? =
-        symbol.callableId.classId?.let { classByClassId(it, scope(project), project) }
+        unrollFakeOverrides().symbol.callableId.classId?.let { classByClassId(it, scope(project), project) }
 
     private fun classByClassId(classId: ClassId, scope: GlobalSearchScope, project: Project): KtClassOrObject? {
         val fqName = classId.asStringForUsingInIndexes().let { classIdMapping[it] ?: it }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/fir/FirUtils.kt

@@ -63,7 +63,7 @@ fun FirFunctionCall.getCalleeSymbol(): FirBasedSymbol<*>? =
 fun FirReference.getResolvedSymbolOfNameReference(): FirBasedSymbol<*>? =
     (this as? FirResolvedNamedReference)?.resolvedSymbol
 
-inline fun <reified D> D.unrollFakeOverrides(): D where D : FirDeclaration, D : FirSymbolOwner<D> {
+inline fun <reified D> D.unrollFakeOverrides(): D where D : FirDeclaration, D : FirSymbolOwner<*> {
     val symbol = symbol
     if (symbol !is PossiblyFirFakeOverrideSymbol<*, *>) return this
     if (!symbol.isFakeOverride) return this

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KotlinFirReferenceContributor.kt

@@ -13,6 +13,7 @@ class KotlinFirReferenceContributor : KotlinReferenceProviderContributor {
             registerProvider(factory = ::KtFirInvokeFunctionReference)
             registerProvider(factory = ::KtFirPropertyDelegationMethodsReference)
             registerProvider(factory = ::KtFirDestructuringDeclarationReference)
+            registerProvider(factory = ::KtFirArrayAccessReference)
         }
     }
 }

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/KtFirArrayAccessReference.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.references
+
+import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
+import org.jetbrains.kotlin.idea.fir.getCalleeSymbol
+import org.jetbrains.kotlin.idea.fir.getOrBuildFirSafe
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.buildSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+import org.jetbrains.kotlin.psi.KtArrayAccessExpression
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtExpression
+
+class KtFirArrayAccessReference(
+    expression: KtArrayAccessExpression
+) : KtArrayAccessReference(expression), KtFirReference {
+    override fun resolveToSymbols(analysisSession: KtAnalysisSession): Collection<KtSymbol> {
+        check(analysisSession is KtFirAnalysisSession)
+        val fir = element.getOrBuildFirSafe<FirFunctionCall>(analysisSession.firResolveState) ?: return emptyList()
+        return listOfNotNull(fir.getCalleeSymbol()?.fir?.buildSymbol(analysisSession.firSymbolBuilder))
+    }
+
+    override fun moveFunctionLiteralOutsideParentheses(callExpression: KtCallExpression) {
+        TODO("Not yet implemented")
+    }
+
+    override fun canMoveLambdaOutsideParentheses(callExpression: KtCallExpression): Boolean {
+        TODO("Not yet implemented")
+    }
+
+    override fun doRenameImplicitConventionalCall(newName: String?): KtExpression {
+        TODO("Not yet implemented")
+    }
+}

idea/testData/resolve/references/arrayAccess/get.kt

@@ -0,0 +1,7 @@
+fun a () {
+    val list = mutableListOf(1)
+    println(list<caret>[0])
+}
+
+// REF: (in kotlin.collections.List).get(kotlin.Int)
+

idea/testData/resolve/references/arrayAccess/set.kt

@@ -0,0 +1,7 @@
+fun a () {
+    val list = mutableListOf(1)
+    list<caret>[0] = 1
+}
+
+// REF: (in kotlin.collections.MutableList).set(kotlin.Int, E)
+

idea/tests/org/jetbrains/kotlin/idea/resolve/ReferenceResolveTestGenerated.java

@@ -418,6 +418,29 @@ public class ReferenceResolveTestGenerated extends AbstractReferenceResolveTest
         runTest("idea/testData/resolve/references/WrongNumberOfTypeArgumentsInSupertype.kt");
     }
 
+    @TestMetadata("idea/testData/resolve/references/arrayAccess")
+    @TestDataPath("$PROJECT_ROOT")
+    @RunWith(JUnit3RunnerWithInners.class)
+    public static class ArrayAccess extends AbstractReferenceResolveTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        public void testAllFilesPresentInArrayAccess() throws Exception {
+            KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/testData/resolve/references/arrayAccess"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
+        }
+
+        @TestMetadata("get.kt")
+        public void testGet() throws Exception {
+            runTest("idea/testData/resolve/references/arrayAccess/get.kt");
+        }
+
+        @TestMetadata("set.kt")
+        public void testSet() throws Exception {
+            runTest("idea/testData/resolve/references/arrayAccess/set.kt");
+        }
+    }
+
     @TestMetadata("idea/testData/resolve/references/delegatedPropertyAccessors")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)