[FIR] Find the correct FIR file in Java finder for file with JvmName

#KT-57802

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaElementFinder.kt

@@ -19,6 +19,7 @@ import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.stubs.StubBase
 import com.intellij.psi.stubs.StubElement
 import com.intellij.util.ArrayUtil
+import org.jetbrains.kotlin.builtins.StandardNames
 import org.jetbrains.kotlin.builtins.jvm.JavaToKotlinClassMap
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
@@ -26,11 +27,12 @@ import org.jetbrains.kotlin.descriptors.Visibilities
 import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSessionComponent
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
+import org.jetbrains.kotlin.fir.caches.FirCache
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.caches.getValue
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.*
+import org.jetbrains.kotlin.fir.expressions.FirConstExpression
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
@@ -65,6 +67,11 @@ class FirJavaElementFinder(
         session.collectAllDependentSourceSessions().mapTo(this) { it.firProvider }
     }
 
+    private val fileCache: FirCache<FqName, Map<String, List<FirFile>>, Nothing?> = session.firCachesFactory
+        .createCache { packageFqName, _ ->
+            firProviders.flatMap { it.getFirFilesByPackage(packageFqName) }.groupBy { it.jvmName() }
+        }
+
     override fun findPackage(qualifiedName: String): PsiPackage? {
         if (firProviders.none { it.symbolProvider.getPackage(FqName(qualifiedName)) != null }) return null
         return PsiPackageImpl(psiManager, qualifiedName)
@@ -91,8 +98,7 @@ class FirJavaElementFinder(
             val classId = ClassId.topLevel(topLevelClass)
 
             // 1. We could be asked to find class of kind "...MainKt" that was created from file "main.kt"
-            val firFile = firProviders.firstNotNullOfOrNull { it.getFirFilesByPackage(classId.packageFqName) }
-                ?.singleOrNull { classId.relativeClassName.asString() == it.name.removeSuffix(".kt").capitalizeAsciiOnly() + "Kt" }
+            val firFile = fileCache.getValue(classId.packageFqName)[classId.relativeClassName.asString()]?.singleOrNull()
             if (firFile != null) {
                 val fileStub = createJavaFileStub(classId.packageFqName, psiManager)
                 return buildFileAsClassStub(firFile, classId, fileStub).psi
@@ -112,6 +118,13 @@ class FirJavaElementFinder(
         return null
     }
 
+    private fun FirFile.jvmName(): String {
+        val jvmNameAnnotation = this.findJvmNameAnnotation()
+        val jvmName = jvmNameAnnotation?.findArgumentByName(StandardNames.NAME)
+        val jvmNameValue = (jvmName as? FirConstExpression<*>)?.value as? String
+        return jvmNameValue ?: (this.name.removeSuffix(".kt").capitalizeAsciiOnly() + "Kt")
+    }
+
     private fun buildFileAsClassStub(firFile: FirFile, classId: ClassId, parent: StubElement<*>): PsiClassStub<*> {
         val stub = PsiClassStubImpl<ClsClassImpl>(
             JavaStubElementTypes.CLASS, parent, classId.asSingleFqName().asString(), classId.relativeClassName.asString(), null,

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/JavaUtils.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.java
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.toAnnotationClassId
 import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
@@ -21,7 +22,6 @@ import org.jetbrains.kotlin.fir.expressions.builder.buildConstExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildErrorExpression
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.expectedConeType
 import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.kotlin.load.java.RXJAVA3_ANNOTATIONS
 import org.jetbrains.kotlin.load.java.structure.JavaAnnotation
@@ -142,3 +142,12 @@ fun FirProperty.hasJvmFieldAnnotation(session: FirSession): Boolean =
 
 fun FirAnnotation.isJvmFieldAnnotation(session: FirSession): Boolean =
     toAnnotationClassId(session) == JvmStandardClassIds.Annotations.JvmField
+
+fun FirDeclaration.findJvmNameAnnotation(): FirAnnotation? {
+    return annotations.firstOrNull {
+        // Access to type must be through `coneTypeOrNull`.
+        // Even if `JvmName` is in the list of annotations that must be resoled for compilation, we still could try to access some user
+        // annotations that could be not resolved.
+        it.annotationTypeRef.coneTypeOrNull?.classId == JvmStandardClassIds.Annotations.JvmName
+    }
+}