FIR: set PSI elements in JavaSymbolProvider / enhancement scope

Related to KT-24351

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/ConversionUtils.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.backend
 
+import com.intellij.psi.PsiCompiledElement
 import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassifierSymbol
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.types.Variance
 internal fun <T : IrElement> FirElement.convertWithOffsets(
     f: (startOffset: Int, endOffset: Int) -> T
 ): T {
+    if (psi is PsiCompiledElement) return f(-1, -1)
     val startOffset = psi?.startOffsetSkippingComments ?: -1
     val endOffset = psi?.endOffset ?: -1
     return f(startOffset, endOffset)

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.java
 
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Visibility
@@ -32,6 +33,7 @@ import org.jetbrains.kotlin.load.java.JavaClassFinder
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames
 import org.jetbrains.kotlin.load.java.structure.JavaClass
 import org.jetbrains.kotlin.load.java.structure.JavaTypeParameter
+import org.jetbrains.kotlin.load.java.structure.impl.JavaElementImpl
 import org.jetbrains.kotlin.load.kotlin.KotlinClassFinder
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -157,7 +159,8 @@ class JavaSymbolProvider(
                     }
                 }
                 FirJavaClass(
-                    session, firSymbol as FirClassSymbol, javaClass.name,
+                    session, (javaClass as? JavaElementImpl<*>)?.psi,
+                    firSymbol as FirClassSymbol, javaClass.name,
                     javaClass.visibility, javaClass.modality,
                     javaClass.classKind, isTopLevel = isTopLevel,
                     isStatic = javaClass.isStatic,
@@ -176,7 +179,8 @@ class JavaSymbolProvider(
                         val fieldSymbol = FirFieldSymbol(fieldId)
                         val returnType = javaField.type
                         val firJavaField = FirJavaField(
-                            this@JavaSymbolProvider.session, fieldSymbol, fieldName,
+                            this@JavaSymbolProvider.session, (javaField as? JavaElementImpl<*>)?.psi,
+                            fieldSymbol, fieldName,
                             javaField.visibility, javaField.modality,
                             returnTypeRef = returnType.toFirJavaTypeRef(this@JavaSymbolProvider.session, javaTypeParameterStack),
                             isVar = !javaField.isFinal,
@@ -192,7 +196,8 @@ class JavaSymbolProvider(
                         val methodSymbol = FirNamedFunctionSymbol(methodId)
                         val returnType = javaMethod.returnType
                         val firJavaMethod = FirJavaMethod(
-                            this@JavaSymbolProvider.session, methodSymbol, methodName,
+                            this@JavaSymbolProvider.session, (javaMethod as? JavaElementImpl<*>)?.psi,
+                            methodSymbol, methodName,
                             javaMethod.visibility, javaMethod.modality,
                             returnTypeRef = returnType.toFirJavaTypeRef(this@JavaSymbolProvider.session, javaTypeParameterStack),
                             isStatic = javaMethod.isStatic
@@ -212,12 +217,14 @@ class JavaSymbolProvider(
 
                     fun addJavaConstructor(
                         visibility: Visibility = this.visibility,
+                        psi: PsiElement? = null,
                         isPrimary: Boolean = false
                     ): FirJavaConstructor {
                         val constructorSymbol = FirConstructorSymbol(constructorId)
                         val classTypeParameters = javaClass.typeParameters.convertTypeParameters(javaTypeParameterStack)
                         val firJavaConstructor = FirJavaConstructor(
-                            this@JavaSymbolProvider.session, constructorSymbol, visibility, isPrimary,
+                            this@JavaSymbolProvider.session, psi,
+                            constructorSymbol, visibility, isPrimary,
                             FirResolvedTypeRefImpl(
                                 null,
                                 firSymbol.constructType(
@@ -236,7 +243,9 @@ class JavaSymbolProvider(
                         addJavaConstructor(isPrimary = true)
                     }
                     for (javaConstructor in javaClassDeclaredConstructors) {
-                        addJavaConstructor(javaConstructor.visibility).apply {
+                        addJavaConstructor(
+                            visibility = javaConstructor.visibility, psi = (javaConstructor as? JavaElementImpl<*>)?.psi
+                        ).apply {
                             this.typeParameters += javaConstructor.typeParameters.convertTypeParameters(javaTypeParameterStack)
                             addAnnotationsFrom(this@JavaSymbolProvider.session, javaConstructor, javaTypeParameterStack)
                             for (valueParameter in javaConstructor.valueParameters) {

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

@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.fir.types.impl.ConeTypeParameterTypeImpl
 import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.load.java.structure.*
+import org.jetbrains.kotlin.load.java.structure.impl.JavaElementImpl
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.Variance.*
@@ -195,7 +196,7 @@ internal fun JavaValueParameter.toFirValueParameters(
     session: FirSession, javaTypeParameterStack: JavaTypeParameterStack
 ): FirValueParameter {
     return FirJavaValueParameter(
-        session, name ?: Name.special("<anonymous Java parameter>"),
+        session, (this as? JavaElementImpl<*>)?.psi, name ?: Name.special("<anonymous Java parameter>"),
         returnTypeRef = type.toFirJavaTypeRef(session, javaTypeParameterStack),
         isVararg = isVararg
     ).apply {

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/declarations/FirJavaClass.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.java.declarations
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
@@ -21,6 +22,7 @@ import org.jetbrains.kotlin.name.Name
 
 class FirJavaClass internal constructor(
     session: FirSession,
+    psi: PsiElement?,
     override val symbol: FirClassSymbol,
     name: Name,
     visibility: Visibility,
@@ -30,8 +32,8 @@ class FirJavaClass internal constructor(
     isStatic: Boolean,
     internal val javaTypeParameterStack: JavaTypeParameterStack
 ) : FirAbstractMemberDeclaration(
-    session, psi = null, name = name,
+    session, psi, name,
-    visibility = visibility, modality = modality,
+    visibility, modality,
     isExpect = false, isActual = false
 ), FirRegularClass, FirModifiableClass {
     init {

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/declarations/FirJavaConstructor.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.java.declarations
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.fir.FirSession
@@ -20,13 +21,14 @@ import org.jetbrains.kotlin.fir.types.FirTypeRef
 
 class FirJavaConstructor(
     session: FirSession,
+    psi: PsiElement?,
     override val symbol: FirConstructorSymbol,
     visibility: Visibility,
     override val isPrimary: Boolean,
     delegatedSelfTypeRef: FirTypeRef
 ) : FirAbstractCallableMember<FirConstructor>(
     session,
-    psi = null,
+    psi,
     name = NAME,
     visibility = visibility,
     modality = Modality.FINAL,

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/declarations/FirJavaField.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.java.declarations
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.fir.FirSession
@@ -20,6 +21,7 @@ import org.jetbrains.kotlin.name.Name
 
 class FirJavaField(
     session: FirSession,
+    psi: PsiElement?,
     override val symbol: FirFieldSymbol,
     name: Name,
     visibility: Visibility,
@@ -28,8 +30,8 @@ class FirJavaField(
     override val isVar: Boolean,
     isStatic: Boolean
 ) : FirAbstractCallableMember<FirField>(
-    session, psi = null, name = name,
+    session, psi, name,
-    visibility = visibility, modality = modality,
+    visibility, modality,
     isExpect = false, isActual = false, isOverride = false,
     receiverTypeRef = null, returnTypeRef = returnTypeRef
 ), FirField {

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/declarations/FirJavaMethod.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.java.declarations
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.fir.FirSession
@@ -16,6 +17,7 @@ import org.jetbrains.kotlin.name.Name
 
 class FirJavaMethod(
     session: FirSession,
+    psi: PsiElement?,
     symbol: FirNamedFunctionSymbol,
     name: Name,
     visibility: Visibility,
@@ -23,7 +25,7 @@ class FirJavaMethod(
     returnTypeRef: FirJavaTypeRef,
     isStatic: Boolean
 ) : FirMemberFunctionImpl(
-    session, null, symbol, name,
+    session, psi, symbol, name,
     visibility, modality,
     false, isActual = false,
     isOverride = false,

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/declarations/FirJavaValueParameter.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.java.declarations
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.impl.FirValueParameterImpl
@@ -13,11 +14,12 @@ import org.jetbrains.kotlin.name.Name
 
 class FirJavaValueParameter(
     session: FirSession,
+    psi: PsiElement?,
     name: Name,
     returnTypeRef: FirJavaTypeRef,
     isVararg: Boolean
 ) : FirValueParameterImpl(
-    session, psi = null, name = name, returnTypeRef = returnTypeRef,
+    session, psi, name, returnTypeRef,
     defaultValue = null, isCrossinline = false, isNoinline = false, isVararg = isVararg
 ) {
     init {

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

@@ -86,6 +86,7 @@ class JavaClassEnhancementScope(
                 with(firElement) {
                     FirJavaField(
                         this@JavaClassEnhancementScope.session,
+                        firElement.psi,
                         symbol,
                         name,
                         visibility,
@@ -185,7 +186,7 @@ class JavaClassEnhancementScope(
                 val symbol = FirConstructorSymbol(methodId)
                 if (firMethod.isPrimary) {
                     FirPrimaryConstructorImpl(
-                        this@JavaClassEnhancementScope.session, null, symbol,
+                        this@JavaClassEnhancementScope.session, firMethod.psi, symbol,
                         firMethod.visibility,
                         isExpect = false,
                         isActual = false,
@@ -194,7 +195,7 @@ class JavaClassEnhancementScope(
                     )
                 } else {
                     FirConstructorImpl(
-                        this@JavaClassEnhancementScope.session, null, symbol,
+                        this@JavaClassEnhancementScope.session, firMethod.psi, symbol,
                         newReceiverTypeRef, newReturnTypeRef
                     )
                 }.apply {
@@ -204,7 +205,7 @@ class JavaClassEnhancementScope(
                 }
             }
             else -> FirMemberFunctionImpl(
-                this@JavaClassEnhancementScope.session, null,
+                this@JavaClassEnhancementScope.session, firMethod.psi,
                 if (!isAccessor) FirNamedFunctionSymbol(methodId)
                 else FirAccessorSymbol(callableId = propertyId!!, accessorId = methodId),
                 name, newReceiverTypeRef, newReturnTypeRef