FIR: cache properties in use-site scope properly

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

@@ -16408,6 +16408,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/fir/IrBuiltIns.kt");
         }
 
+        @Test
+        @TestMetadata("JKEnumConstant.kt")
+        public void testJKEnumConstant() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/JKEnumConstant.kt");
+        }
+
         @Test
         @TestMetadata("jvmFieldInLocalClass.kt")
         public void testJvmFieldInLocalClass() throws Exception {

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/scopes/JavaClassUseSiteMemberScope.kt

@@ -103,16 +103,17 @@ class JavaClassUseSiteMemberScope(
         return minOf(a, b)
     }
 
-    override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
+    override fun doProcessProperties(name: Name): Collection<FirVariableSymbol<*>> {
         val fields = mutableSetOf<FirCallableSymbol<*>>()
         val fieldNames = mutableSetOf<Name>()
+        val result = mutableSetOf<FirVariableSymbol<*>>()
 
         // fields
         declaredMemberScope.processPropertiesByName(name) processor@{ variableSymbol ->
             if (variableSymbol.isStatic) return@processor
             fields += variableSymbol
             fieldNames += variableSymbol.fir.name
-            processor(variableSymbol)
+            result += variableSymbol
         }
 
         val fromSupertypes = superTypesScope.getProperties(name)
@@ -120,7 +121,7 @@ class JavaClassUseSiteMemberScope(
         for (propertyFromSupertype in fromSupertypes) {
             if (propertyFromSupertype is FirFieldSymbol) {
                 if (propertyFromSupertype.fir.name !in fieldNames) {
-                    processor(propertyFromSupertype)
+                    result += propertyFromSupertype
                 }
                 continue
             }
@@ -132,11 +133,12 @@ class JavaClassUseSiteMemberScope(
                 overrideInClass != null -> {
                     directOverriddenProperties.getOrPut(overrideInClass) { mutableListOf() }.add(propertyFromSupertype)
                     overrideByBase[propertyFromSupertype] = overrideInClass
-                    processor(overrideInClass)
+                    result += overrideInClass
                 }
-                else -> processor(propertyFromSupertype)
+                else -> result += propertyFromSupertype
             }
         }
+        return result
     }
 
     private fun FirVariableSymbol<*>.createOverridePropertyIfExists(

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/AbstractFirUseSiteMemberScope.kt

@@ -22,6 +22,7 @@ abstract class AbstractFirUseSiteMemberScope(
 ) : AbstractFirOverrideScope(session, overrideChecker) {
 
     private val functions = hashMapOf<Name, Collection<FirNamedFunctionSymbol>>()
+    private val properties = hashMapOf<Name, Collection<FirVariableSymbol<*>>>()
     val directOverriddenFunctions = hashMapOf<FirNamedFunctionSymbol, Collection<FirNamedFunctionSymbol>>()
     protected val directOverriddenProperties = hashMapOf<FirPropertySymbol, MutableList<FirPropertySymbol>>()
 
@@ -57,6 +58,16 @@ abstract class AbstractFirUseSiteMemberScope(
         }
     }
 
+    final override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
+        properties.getOrPut(name) {
+            doProcessProperties(name)
+        }.forEach {
+            processor(it)
+        }
+    }
+
+    protected abstract fun doProcessProperties(name: Name): Collection<FirVariableSymbol<*>>
+
     private fun computeDirectOverridden(symbol: FirNamedFunctionSymbol): Collection<FirNamedFunctionSymbol> {
         val result = mutableListOf<FirNamedFunctionSymbol>()
         val firSimpleFunction = symbol.fir

compiler/fir/providers/src/org/jetbrains/kotlin/fir/scopes/impl/FirClassUseSiteMemberScope.kt

@@ -20,8 +20,9 @@ class FirClassUseSiteMemberScope(
     declaredMemberScope: FirContainingNamesAwareScope
 ) : AbstractFirUseSiteMemberScope(session, FirStandardOverrideChecker(session), superTypesScope, declaredMemberScope) {
 
-    override fun processPropertiesByName(name: Name, processor: (FirVariableSymbol<*>) -> Unit) {
+    override fun doProcessProperties(name: Name): Collection<FirVariableSymbol<*>> {
         val seen = mutableSetOf<FirVariableSymbol<*>>()
+        val result = mutableSetOf<FirVariableSymbol<*>>()
         declaredMemberScope.processPropertiesByName(name) l@{
             if (it.isStatic) return@l
             if (it is FirPropertySymbol) {
@@ -29,15 +30,16 @@ class FirClassUseSiteMemberScope(
                 this@FirClassUseSiteMemberScope.directOverriddenProperties[it] = directOverridden
             }
             seen += it
-            processor(it)
+            result += it
         }
 
         superTypesScope.processPropertiesByName(name) {
             val overriddenBy = it.getOverridden(seen)
             if (overriddenBy == null) {
-                processor(it)
+                result += it
             }
         }
+        return result
     }
 
     private fun computeDirectOverridden(property: FirProperty): MutableList<FirPropertySymbol> {

compiler/testData/codegen/box/fir/JKEnumConstant.kt

@@ -0,0 +1,27 @@
+// TARGET_BACKEND: JVM_IR
+// WITH_RUNTIME
+
+interface PsiOwner {
+    var psi: String?
+}
+
+class PsiOwnerImpl(override var psi: String? = null) : PsiOwner
+
+interface JKElement
+
+interface JKFormattingOwner
+
+abstract class JKTreeElement : JKElement, JKFormattingOwner
+
+abstract class JKDeclaration : JKTreeElement(), PsiOwner by PsiOwnerImpl()
+
+interface JKAnnotationListOwner : JKFormattingOwner
+
+open class JKVariable : JKDeclaration(), JKAnnotationListOwner
+
+class JKEnumConstant : JKVariable()
+
+fun box(): String {
+    val constant = JKEnumConstant().also { it.psi = "OK" }
+    return constant.psi!!
+}

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

@@ -16408,6 +16408,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/fir/IrBuiltIns.kt");
         }
 
+        @Test
+        @TestMetadata("JKEnumConstant.kt")
+        public void testJKEnumConstant() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/JKEnumConstant.kt");
+        }
+
         @Test
         @TestMetadata("jvmFieldInLocalClass.kt")
         public void testJvmFieldInLocalClass() throws Exception {