FIR IDE: resolve arrayOf calls in annotations

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirCallResolver.kt

@@ -19,10 +19,7 @@ import org.jetbrains.kotlin.fir.references.impl.FirSimpleNamedReference
 import org.jetbrains.kotlin.fir.resolve.calls.FirErrorReferenceWithCandidate
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.ConeKotlinErrorType
-import org.jetbrains.kotlin.fir.types.classId
-import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.analysis.api.fir.getCandidateSymbols
 import org.jetbrains.kotlin.analysis.api.fir.isImplicitFunctionCall
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
@@ -32,6 +29,10 @@ import org.jetbrains.kotlin.analysis.api.components.KtCallResolver
 import org.jetbrains.kotlin.analysis.api.diagnostics.KtNonBoundToPsiErrorDiagnostic
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.buildSymbol
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayOf
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayOfSymbol
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayTypeToArrayOfCall
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirFunctionSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.analysis.api.symbols.markers.KtSymbolWithMembers
 import org.jetbrains.kotlin.analysis.api.tokens.ValidityToken
@@ -106,6 +107,7 @@ internal class KtFirCallResolver(
 
     override fun resolveCall(call: KtCallElement): KtCall? = withValidityAssertion {
         return when (val fir = call.getOrBuildFir(firResolveState)) {
+            is FirArrayOfCall -> resolveArrayOfCall(fir)
             is FirFunctionCall -> resolveCall(fir)
             is FirAnnotationCall -> fir.asAnnotationCall()
             is FirDelegatedConstructorCall -> fir.asDelegatedConstructorCall()
@@ -122,6 +124,25 @@ internal class KtFirCallResolver(
         }
     }
 
+    private fun resolveArrayOfCall(arrayOfCall: FirArrayOfCall): KtCall? {
+        val arrayOfSymbol = with(analysisSession) {
+            val type = arrayOfCall.typeRef.coneTypeSafe<ConeClassLikeType>()
+                ?: return run {
+                    val defaultArrayOfSymbol = arrayOfSymbol(arrayOf) ?: return null
+                    KtFunctionCall(
+                        arrayOfCall.createArgumentMapping(defaultArrayOfSymbol),
+                        KtErrorCallTarget(
+                            listOf(defaultArrayOfSymbol),
+                            KtNonBoundToPsiErrorDiagnostic(factoryName = null, "type of arrayOf call is not resolved", token)
+                        )
+                    )
+                }
+            val call = arrayTypeToArrayOfCall[type.lookupTag.classId] ?: arrayOf
+            arrayOfSymbol(call)
+        } ?: return null
+        return KtFunctionCall(arrayOfCall.createArgumentMapping(arrayOfSymbol), KtSuccessCallTarget(arrayOfSymbol))
+    }
+
     private fun resolveCall(firCall: FirFunctionCall): KtCall? {
         val session = firResolveState.rootModuleSession
         return when {
@@ -210,35 +231,56 @@ internal class KtFirCallResolver(
         }
     }
 
+    private fun FirExpression.findSourceKtExpressionForCallArgument(): KtExpression? {
+        // For spread, named, and lambda arguments, the source is the KtValueArgument.
+        // For other arguments (including array indices), the source is the KtExpression.
+        return when (this) {
+            is FirNamedArgumentExpression, is FirSpreadArgumentExpression, is FirLambdaArgumentExpression ->
+                realPsi.safeAs<KtValueArgument>()?.getArgumentExpression()
+            else -> realPsi as? KtExpression
+        }
+    }
+
+    private fun mapArgumentExpressionToParameter(
+        argumentExpression: FirExpression,
+        parameterSymbol: KtValueParameterSymbol,
+        argumentMapping: LinkedHashMap<KtExpression, KtValueParameterSymbol>
+    ) {
+        if (argumentExpression is FirVarargArgumentsExpression) {
+            for (varargArgument in argumentExpression.arguments) {
+                val valueArgument = varargArgument.findSourceKtExpressionForCallArgument() ?: return
+                argumentMapping[valueArgument] = parameterSymbol
+            }
+        } else {
+            val valueArgument = argumentExpression.findSourceKtExpressionForCallArgument() ?: return
+            argumentMapping[valueArgument] = parameterSymbol
+        }
+    }
+
     private fun FirCall.createArgumentMapping(): LinkedHashMap<KtExpression, KtValueParameterSymbol> {
         val ktArgumentMapping = LinkedHashMap<KtExpression, KtValueParameterSymbol>()
         argumentMapping?.let {
-            fun FirExpression.findKtExpression(): KtExpression? {
-                // For spread, named, and lambda arguments, the source is the KtValueArgument.
-                // For other arguments (including array indices), the source is the KtExpression.
-                return when (this) {
-                    is FirNamedArgumentExpression, is FirSpreadArgumentExpression, is FirLambdaArgumentExpression ->
-                        realPsi.safeAs<KtValueArgument>()?.getArgumentExpression()
-                    else -> realPsi as? KtExpression
-                }
-            }
-
             for ((firExpression, firValueParameter) in it.entries) {
                 val parameterSymbol = firValueParameter.buildSymbol(firSymbolBuilder) as KtValueParameterSymbol
-                if (firExpression is FirVarargArgumentsExpression) {
+                mapArgumentExpressionToParameter(firExpression, parameterSymbol, ktArgumentMapping)
-                    for (varargArgument in firExpression.arguments) {
-                        val valueArgument = varargArgument.findKtExpression() ?: continue
-                        ktArgumentMapping[valueArgument] = parameterSymbol
-                    }
-                } else {
-                    val valueArgument = firExpression.findKtExpression() ?: continue
-                    ktArgumentMapping[valueArgument] = parameterSymbol
-                }
             }
         }
         return ktArgumentMapping
     }
 
+    private fun FirArrayOfCall.createArgumentMapping(
+        arrayOfCallSymbol: KtFirFunctionSymbol
+    ): LinkedHashMap<KtExpression, KtValueParameterSymbol> {
+        val ktArgumentMapping = LinkedHashMap<KtExpression, KtValueParameterSymbol>()
+        val parameterSymbol = arrayOfCallSymbol.firRef.withFir {
+            it.valueParameters.single().buildSymbol(firSymbolBuilder) as KtValueParameterSymbol
+        }
+        for (firExpression in argumentList.arguments) {
+            mapArgumentExpressionToParameter(firExpression, parameterSymbol, ktArgumentMapping)
+        }
+        return ktArgumentMapping
+    }
+
     private fun FirErrorNamedReference.createErrorCallTarget(qualifiedAccessSource: FirSourceElement?): KtErrorCallTarget =
         KtErrorCallTarget(
             getCandidateSymbols().mapNotNull { it.fir.buildSymbol(firSymbolBuilder) as? KtFunctionLikeSymbol },

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirCollectionLiteralReference.kt

@@ -5,20 +5,17 @@
 
 package org.jetbrains.kotlin.idea.references
 
-import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
-import org.jetbrains.kotlin.fir.expressions.FirArrayOfCall
-import org.jetbrains.kotlin.fir.resolve.symbolProvider
-import org.jetbrains.kotlin.name.StandardClassIds
-import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
-import org.jetbrains.kotlin.fir.types.ConeClassLikeType
-import org.jetbrains.kotlin.fir.types.coneTypeSafe
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayOf
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayOfSymbol
+import org.jetbrains.kotlin.analysis.api.fir.symbols.KtFirArrayOfSymbolProvider.arrayTypeToArrayOfCall
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
-import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirSafe
-import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.fir.expressions.FirArrayOfCall
-import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.fir.types.ConeClassLikeType
+import org.jetbrains.kotlin.fir.types.coneTypeSafe
+
 import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
 
 class KtFirCollectionLiteralReference(
@@ -31,26 +28,4 @@ class KtFirCollectionLiteralReference(
         val call = arrayTypeToArrayOfCall[type.lookupTag.classId] ?: arrayOf
         return listOfNotNull(arrayOfSymbol(call))
     }
-
-    private fun KtFirAnalysisSession.arrayOfSymbol(identifier: Name): KtSymbol? {
-        val fir = firResolveState.rootModuleSession.symbolProvider.getTopLevelCallableSymbols(kotlinPackage, identifier).firstOrNull {
-            /* choose (for byte array)
-             * public fun byteArrayOf(vararg elements: kotlin.Byte): kotlin.ByteArray
-             */
-            (it as? FirFunctionSymbol<*>)?.fir?.valueParameters?.singleOrNull()?.isVararg == true
-        }?.fir as? FirSimpleFunction ?: return null
-        return firSymbolBuilder.functionLikeBuilder.buildFunctionSymbol(fir)
-    }
-
-    companion object {
-        private val kotlinPackage = FqName("kotlin")
-        private val arrayOf = Name.identifier("arrayOf")
-        private val arrayTypeToArrayOfCall = run {
-            StandardClassIds.primitiveArrayTypeByElementType.values + StandardClassIds.unsignedArrayTypeByElementType.values
-        }.associateWith { it.correspondingArrayOfCallFqName() }
-
-        private fun ClassId.correspondingArrayOfCallFqName(): Name =
-            Name.identifier("${shortClassName.identifier.replaceFirstChar(Char::lowercaseChar)}Of")
-
-    }
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/symbols/KtFirArrayOfSymbolProvider.kt

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2010-2021 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.analysis.api.fir.symbols
+
+import org.jetbrains.kotlin.analysis.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
+import org.jetbrains.kotlin.fir.resolve.symbolProvider
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.name.StandardClassIds
+
+object KtFirArrayOfSymbolProvider {
+    internal fun KtFirAnalysisSession.arrayOfSymbol(identifier: Name): KtFirFunctionSymbol? {
+        val fir = firResolveState.rootModuleSession.symbolProvider.getTopLevelCallableSymbols(kotlinPackage, identifier).firstOrNull {
+            /* choose (for byte array)
+             * public fun byteArrayOf(vararg elements: kotlin.Byte): kotlin.ByteArray
+             */
+            (it as? FirFunctionSymbol<*>)?.fir?.valueParameters?.singleOrNull()?.isVararg == true
+        }?.fir as? FirSimpleFunction ?: return null
+        return firSymbolBuilder.functionLikeBuilder.buildFunctionSymbol(fir)
+    }
+
+    private val kotlinPackage = FqName("kotlin")
+    internal val arrayOf = Name.identifier("arrayOf")
+    internal val arrayTypeToArrayOfCall = run {
+        StandardClassIds.primitiveArrayTypeByElementType.values + StandardClassIds.unsignedArrayTypeByElementType.values
+    }.associateWith { it.correspondingArrayOfCallFqName() }
+
+    private fun ClassId.correspondingArrayOfCallFqName(): Name =
+        Name.identifier("${shortClassName.identifier.replaceFirstChar(Char::lowercaseChar)}Of")
+}

analysis/analysis-api-fir/tests/org/jetbrains/kotlin/analysis/api/fir/components/ResolveCallTestGenerated.java

@@ -24,6 +24,12 @@ public class ResolveCallTestGenerated extends AbstractResolveCallTest {
         KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("analysis/analysis-api/testData/analysisSession/resolveCall"), Pattern.compile("^(.+)\\.kt$"), null, true);
     }
 
+    @Test
+    @TestMetadata("arrayOfInAnnotation.kt")
+    public void testArrayOfInAnnotation() throws Exception {
+        runTest("analysis/analysis-api/testData/analysisSession/resolveCall/arrayOfInAnnotation.kt");
+    }
+
     @Test
     @TestMetadata("delegatedConstructorCall_super.kt")
     public void testDelegatedConstructorCall_super() throws Exception {
@@ -150,6 +156,12 @@ public class ResolveCallTestGenerated extends AbstractResolveCallTest {
         runTest("analysis/analysis-api/testData/analysisSession/resolveCall/indexedSetWithTooManyArgs.kt");
     }
 
+    @Test
+    @TestMetadata("intArrayOfInAnnotation.kt")
+    public void testIntArrayOfInAnnotation() throws Exception {
+        runTest("analysis/analysis-api/testData/analysisSession/resolveCall/intArrayOfInAnnotation.kt");
+    }
+
     @Test
     @TestMetadata("javaFunctionCall.kt")
     public void testJavaFunctionCall() throws Exception {

analysis/analysis-api/testData/analysisSession/resolveCall/arrayOfInAnnotation.kt

@@ -0,0 +1,4 @@
+annotation class RequiresPermission(val anyOf: IntArray)
+
+@RequiresPermission(anyOf = <expr>arrayOf(1, 2, 3)</expr>)
+fun foo(): Int = 5

analysis/analysis-api/testData/analysisSession/resolveCall/arrayOfInAnnotation.txt

@@ -0,0 +1,3 @@
+KtFunctionCall:
+argumentMapping = { 1 -> (vararg elements: T), 2 -> (vararg elements: T), 3 -> (vararg elements: T) }
+targetFunction = kotlin/arrayOf(vararg elements: T): kotlin.Array<T>

analysis/analysis-api/testData/analysisSession/resolveCall/intArrayOfInAnnotation.kt

@@ -0,0 +1,4 @@
+annotation class RequiresPermission(val anyOf: IntArray)
+
+@RequiresPermission(anyOf = <expr>intArrayOf(1, 2, 3)</expr>)
+fun foo(): Int = 5

analysis/analysis-api/testData/analysisSession/resolveCall/intArrayOfInAnnotation.txt

@@ -0,0 +1,3 @@
+KtFunctionCall:
+argumentMapping = { 1 -> (vararg elements: kotlin.Int), 2 -> (vararg elements: kotlin.Int), 3 -> (vararg elements: kotlin.Int) }
+targetFunction = kotlin/intArrayOf(vararg elements: kotlin.Int): kotlin.IntArray