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

#KT-57802

compiler/fir/checkers/checkers.jvm/src/org/jetbrains/kotlin/fir/analysis/jvm/checkers/declaration/FirJvmNameChecker.kt

@@ -18,8 +18,8 @@ import org.jetbrains.kotlin.fir.declarations.utils.isOverridable
 import org.jetbrains.kotlin.fir.declarations.utils.isOverride
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirConstExpression
+import org.jetbrains.kotlin.fir.java.findJvmNameAnnotation
 import org.jetbrains.kotlin.fir.resolve.getContainingClass
-import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.resolvedType
 import org.jetbrains.kotlin.name.JvmStandardClassIds
@@ -57,12 +57,6 @@ object FirJvmNameChecker : FirBasicDeclarationChecker() {
         }
     }
 
-    private fun FirDeclaration.findJvmNameAnnotation(): FirAnnotation? {
-        return annotations.firstOrNull {
-            it.annotationTypeRef.coneType.classId == JvmStandardClassIds.Annotations.JvmName
-        }
-    }
-
     private fun CheckerContext.isRenamableFunction(function: FirFunction): Boolean {
         val containingClass = function.getContainingClassSymbol(session)
         return containingClass != null || !function.symbol.callableId.isLocal

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirLightTreeBlackBoxCodegenTestGenerated.java

@@ -29261,6 +29261,12 @@ public class FirLightTreeBlackBoxCodegenTestGenerated extends AbstractFirLightTr
                 runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConst.kt");
             }
 
+            @Test
+            @TestMetadata("accessTopLevelConstWithCustomFileName.kt")
+            public void testAccessTopLevelConstWithCustomFileName() throws Exception {
+                runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConstWithCustomFileName.kt");
+            }
+
             @Test
             public void testAllFilesPresentInConstEvaluationFromJavaWorld() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirLightTreeBlackBoxCodegenWithIrFakeOverrideGeneratorTestGenerated.java

@@ -29261,6 +29261,12 @@ public class FirLightTreeBlackBoxCodegenWithIrFakeOverrideGeneratorTestGenerated
                 runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConst.kt");
             }
 
+            @Test
+            @TestMetadata("accessTopLevelConstWithCustomFileName.kt")
+            public void testAccessTopLevelConstWithCustomFileName() throws Exception {
+                runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConstWithCustomFileName.kt");
+            }
+
             @Test
             public void testAllFilesPresentInConstEvaluationFromJavaWorld() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirPsiBlackBoxCodegenTestGenerated.java

@@ -29261,6 +29261,12 @@ public class FirPsiBlackBoxCodegenTestGenerated extends AbstractFirPsiBlackBoxCo
                 runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConst.kt");
             }
 
+            @Test
+            @TestMetadata("accessTopLevelConstWithCustomFileName.kt")
+            public void testAccessTopLevelConstWithCustomFileName() throws Exception {
+                runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConstWithCustomFileName.kt");
+            }
+
             @Test
             public void testAllFilesPresentInConstEvaluationFromJavaWorld() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);

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.caches.FirCache
-import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.caches.firCachesFactory
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
+import org.jetbrains.kotlin.fir.caches.getValue
-import org.jetbrains.kotlin.fir.declarations.FirTypeParameter
+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) }
+            val firFile = fileCache.getValue(classId.packageFqName)[classId.relativeClassName.asString()]?.singleOrNull()
-                ?.singleOrNull { classId.relativeClassName.asString() == it.name.removeSuffix(".kt").capitalizeAsciiOnly() + "Kt" }
             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
+    }
+}

compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConstWithCustomFileName.kt

@@ -0,0 +1,21 @@
+// TARGET_BACKEND: JVM
+// IGNORE_LIGHT_ANALYSIS
+// WITH_STDLIB
+// FILE: Bar.java
+package one.two;
+
+public class Bar {
+    public static final int BAR = OtherKt.FOO + 1;
+}
+
+// FILE: Main.kt
+@file:JvmName(<!EVALUATED("OtherKt")!>"OtherKt"<!>)
+package one.two
+
+const val FOO = <!EVALUATED("1")!>1<!>
+
+const val BAZ = <!EVALUATED("3")!>Bar.BAR + 1<!>
+
+fun box(): String {
+    return "OK"
+}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/BlackBoxCodegenTestGenerated.java

@@ -27839,6 +27839,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConst.kt");
             }
 
+            @Test
+            @TestMetadata("accessTopLevelConstWithCustomFileName.kt")
+            public void testAccessTopLevelConstWithCustomFileName() throws Exception {
+                runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConstWithCustomFileName.kt");
+            }
+
             @Test
             public void testAllFilesPresentInConstEvaluationFromJavaWorld() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM, true);

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -29261,6 +29261,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConst.kt");
             }
 
+            @Test
+            @TestMetadata("accessTopLevelConstWithCustomFileName.kt")
+            public void testAccessTopLevelConstWithCustomFileName() throws Exception {
+                runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConstWithCustomFileName.kt");
+            }
+
             @Test
             public void testAllFilesPresentInConstEvaluationFromJavaWorld() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenWithIrInlinerTestGenerated.java

@@ -29261,6 +29261,12 @@ public class IrBlackBoxCodegenWithIrInlinerTestGenerated extends AbstractIrBlack
                 runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConst.kt");
             }
 
+            @Test
+            @TestMetadata("accessTopLevelConstWithCustomFileName.kt")
+            public void testAccessTopLevelConstWithCustomFileName() throws Exception {
+                runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConstWithCustomFileName.kt");
+            }
+
             @Test
             public void testAllFilesPresentInConstEvaluationFromJavaWorld() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -24708,6 +24708,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConst.kt");
             }
 
+            @TestMetadata("accessTopLevelConstWithCustomFileName.kt")
+            public void testAccessTopLevelConstWithCustomFileName() throws Exception {
+                runTest("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld/accessTopLevelConstWithCustomFileName.kt");
+            }
+
             public void testAllFilesPresentInConstEvaluationFromJavaWorld() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/codegen/box/involvesIrInterpreter/constEvaluationFromJavaWorld"), Pattern.compile("^(.+)\\.kt$"), null, TargetBackend.JVM_IR, true);
             }