FIR Java: make method annotations lazy

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

@@ -178,7 +178,6 @@ class JavaSymbolProvider(
                 javaTypeParameterStack.addStack(parentStack)
             }
         }
-        val methodMap = mutableMapOf<JavaMethod, FirJavaMethod>()
         val firJavaClass = buildJavaClass {
             source = (javaClass as? JavaElementImpl<*>)?.psi?.toFirPsiSourceElement()
             session = this@JavaSymbolProvider.session
@@ -201,7 +200,7 @@ class JavaSymbolProvider(
                 }
             }
 
-            val dispatchReceiver = classId.defaultType(typeParameters.map { it.symbol } )
+            val dispatchReceiver = classId.defaultType(typeParameters.map { it.symbol })
 
             status = FirResolvedDeclarationStatusImpl(
                 javaClass.visibility,
@@ -230,9 +229,7 @@ class JavaSymbolProvider(
                     classIsAnnotation,
                     valueParametersForAnnotationConstructor,
                     dispatchReceiver
-                ).apply {
+                )
-                    methodMap[javaMethod] = this
-                }
             }
             val javaClassDeclaredConstructors = javaClass.constructors
             val constructorId = CallableId(classId.packageFqName, classId.relativeClassName, classId.shortClassName)
@@ -283,10 +280,6 @@ class JavaSymbolProvider(
             }
         )
         firJavaClass.addAnnotationsFrom(this@JavaSymbolProvider.session, javaClass, javaTypeParameterStack)
-        // NB: this is done here to unbind possible annotation cycle
-        for ((javaMethod, firJavaMethod) in methodMap) {
-            firJavaMethod.annotations.addAnnotationsFrom(session, javaMethod, javaTypeParameterStack)
-        }
         return firJavaClass
     }
 
@@ -340,7 +333,7 @@ class JavaSymbolProvider(
                 returnTypeRef = returnType.toFirJavaTypeRef(this@JavaSymbolProvider.session, javaTypeParameterStack)
                 isVar = !javaField.isFinal
                 isStatic = javaField.isStatic
-                annotationBuilder = { javaField.annotations.map { it.toFirAnnotationCall(session, javaTypeParameterStack) }}
+                annotationBuilder = { javaField.annotations.map { it.toFirAnnotationCall(session, javaTypeParameterStack) } }
                 initializer = convertJavaInitializerToFir(javaField.initializerValue)
 
                 if (!javaField.isStatic) {
@@ -386,6 +379,7 @@ class JavaSymbolProvider(
                     this@JavaSymbolProvider.session, index, javaTypeParameterStack,
                 )
             }
+            annotationBuilder = { javaMethod.annotations.map { it.toFirAnnotationCall(session, javaTypeParameterStack) } }
             status = FirResolvedDeclarationStatusImpl(
                 javaMethod.visibility,
                 javaMethod.modality

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

@@ -10,17 +10,24 @@ import org.jetbrains.kotlin.descriptors.Visibility
 import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSourceElement
+import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
 import org.jetbrains.kotlin.fir.builder.FirBuilderDsl
+import org.jetbrains.kotlin.fir.contracts.FirContractDescription
 import org.jetbrains.kotlin.fir.contracts.impl.FirEmptyContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.builder.FirSimpleFunctionBuilder
+import org.jetbrains.kotlin.fir.declarations.builder.FirFunctionBuilder
-import org.jetbrains.kotlin.fir.declarations.impl.FirSimpleFunctionImpl
+import org.jetbrains.kotlin.fir.declarations.builder.FirTypeParametersOwnerBuilder
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.references.FirControlFlowGraphReference
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.jvm.FirJavaTypeRef
+import org.jetbrains.kotlin.fir.visitors.FirTransformer
+import org.jetbrains.kotlin.fir.visitors.FirVisitor
+import org.jetbrains.kotlin.fir.visitors.transformInplace
+import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedContainerSource
 import org.jetbrains.kotlin.util.OperatorNameConventions.ASSIGNMENT_OPERATIONS
@@ -40,51 +47,153 @@ import kotlin.properties.Delegates
 
 @OptIn(FirImplementationDetail::class)
 class FirJavaMethod @FirImplementationDetail constructor(
-    source: FirSourceElement?,
+    override val source: FirSourceElement?,
-    session: FirSession,
+    override val session: FirSession,
-    resolvePhase: FirResolvePhase,
+    override var resolvePhase: FirResolvePhase,
-    attributes: FirDeclarationAttributes,
+    override val attributes: FirDeclarationAttributes,
-    returnTypeRef: FirTypeRef,
+    override var returnTypeRef: FirTypeRef,
-    receiverTypeRef: FirTypeRef?,
+    override val typeParameters: MutableList<FirTypeParameter>,
-    typeParameters: MutableList<FirTypeParameter>,
+    override val valueParameters: MutableList<FirValueParameter>,
-    valueParameters: MutableList<FirValueParameter>,
+    override val name: Name,
-    body: FirBlock?,
+    override var status: FirDeclarationStatus,
-    name: Name,
+    override val symbol: FirNamedFunctionSymbol,
-    status: FirDeclarationStatus,
+    annotationBuilder: () -> List<FirAnnotationCall>,
-    containerSource: DeserializedContainerSource?,
+    override val dispatchReceiverType: ConeKotlinType?,
-    symbol: FirNamedFunctionSymbol,
+) : FirSimpleFunction() {
-    annotations: MutableList<FirAnnotationCall>,
+    init {
-    dispatchReceiverType: ConeKotlinType?,
+        symbol.bind(this)
-) : FirSimpleFunctionImpl(
+    }
-    source,
+
-    session,
+    override val receiverTypeRef: FirTypeRef?
-    resolvePhase,
+        get() = null
-    FirDeclarationOrigin.Java,
+
-    attributes,
+    override val body: FirBlock?
-    returnTypeRef,
+        get() = null
-    receiverTypeRef,
+
-    valueParameters,
+    override val containerSource: DeserializedContainerSource?
-    body,
+        get() = null
-    status,
+
-    containerSource,
+    override val origin: FirDeclarationOrigin
-    dispatchReceiverType = dispatchReceiverType,
+        get() = FirDeclarationOrigin.Java
-    contractDescription = FirEmptyContractDescription,
+
-    name,
+    override val contractDescription: FirContractDescription
-    symbol,
+        get() = FirEmptyContractDescription
-    annotations,
+
-    typeParameters,
+    override var controlFlowGraphReference: FirControlFlowGraphReference? = null
-)
+
+    override val annotations: List<FirAnnotationCall> by lazy { annotationBuilder() }
+
+    override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
+        returnTypeRef.accept(visitor, data)
+        receiverTypeRef?.accept(visitor, data)
+        controlFlowGraphReference?.accept(visitor, data)
+        valueParameters.forEach { it.accept(visitor, data) }
+        body?.accept(visitor, data)
+        status.accept(visitor, data)
+        contractDescription.accept(visitor, data)
+        annotations.forEach { it.accept(visitor, data) }
+        typeParameters.forEach { it.accept(visitor, data) }
+    }
+
+    override fun <D> transformChildren(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        transformReturnTypeRef(transformer, data)
+        transformReceiverTypeRef(transformer, data)
+        controlFlowGraphReference = controlFlowGraphReference?.transformSingle(transformer, data)
+        transformValueParameters(transformer, data)
+        transformBody(transformer, data)
+        transformStatus(transformer, data)
+        transformContractDescription(transformer, data)
+        transformAnnotations(transformer, data)
+        transformTypeParameters(transformer, data)
+        return this
+    }
+
+    override fun <D> transformReturnTypeRef(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        returnTypeRef = returnTypeRef.transformSingle(transformer, data)
+        return this
+    }
+
+    override fun <D> transformReceiverTypeRef(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        return this
+    }
+
+    override fun <D> transformValueParameters(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        valueParameters.transformInplace(transformer, data)
+        return this
+    }
+
+    override fun <D> transformBody(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        return this
+    }
+
+    override fun <D> transformStatus(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        status = status.transformSingle(transformer, data)
+        return this
+    }
+
+    override fun <D> transformContractDescription(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        return this
+    }
+
+    override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        return this
+    }
+
+    override fun <D> transformTypeParameters(transformer: FirTransformer<D>, data: D): FirSimpleFunction {
+        typeParameters.transformInplace(transformer, data)
+        return this
+    }
+
+    override fun replaceResolvePhase(newResolvePhase: FirResolvePhase) {
+        resolvePhase = newResolvePhase
+    }
+
+    override fun replaceReturnTypeRef(newReturnTypeRef: FirTypeRef) {
+        returnTypeRef = newReturnTypeRef
+    }
+
+    override fun replaceReceiverTypeRef(newReceiverTypeRef: FirTypeRef?) {
+    }
+
+    override fun replaceControlFlowGraphReference(newControlFlowGraphReference: FirControlFlowGraphReference?) {
+        controlFlowGraphReference = newControlFlowGraphReference
+    }
+
+    override fun replaceValueParameters(newValueParameters: List<FirValueParameter>) {
+        valueParameters.clear()
+        valueParameters.addAll(newValueParameters)
+    }
+
+    override fun replaceBody(newBody: FirBlock?) {
+    }
+
+    override fun replaceContractDescription(newContractDescription: FirContractDescription) {
+    }
+}
 
 val ALL_JAVA_OPERATION_NAMES =
     UNARY_OPERATION_NAMES + BINARY_OPERATION_NAMES + ASSIGNMENT_OPERATIONS + DELEGATED_PROPERTY_OPERATORS +
             EQUALS + COMPARE_TO + CONTAINS + INVOKE + ITERATOR + GET + SET + NEXT + HAS_NEXT
 
 @FirBuilderDsl
-class FirJavaMethodBuilder : FirSimpleFunctionBuilder() {
+class FirJavaMethodBuilder : FirFunctionBuilder, FirTypeParametersOwnerBuilder, FirAnnotationContainerBuilder {
+    override var source: FirSourceElement? = null
+    override lateinit var session: FirSession
+    override var attributes: FirDeclarationAttributes = FirDeclarationAttributes()
+    override lateinit var returnTypeRef: FirTypeRef
+    override val valueParameters: MutableList<FirValueParameter> = mutableListOf()
+    override var body: FirBlock? = null
+    lateinit var status: FirDeclarationStatus
+    var dispatchReceiverType: ConeKotlinType? = null
+    lateinit var name: Name
+    lateinit var symbol: FirNamedFunctionSymbol
+    override val annotations: MutableList<FirAnnotationCall> = mutableListOf()
+    override val typeParameters: MutableList<FirTypeParameter> = mutableListOf()
     lateinit var visibility: Visibility
     var modality: Modality? = null
     var isStatic: Boolean by Delegates.notNull()
-    override var resolvePhase: FirResolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
+    var resolvePhase: FirResolvePhase = FirResolvePhase.ANALYZED_DEPENDENCIES
+    lateinit var annotationBuilder: () -> List<FirAnnotationCall>
 
     @Deprecated("Modification of 'origin' has no impact for FirJavaFunctionBuilder", level = DeprecationLevel.HIDDEN)
     override var origin: FirDeclarationOrigin
@@ -101,15 +210,12 @@ class FirJavaMethodBuilder : FirSimpleFunctionBuilder() {
             resolvePhase,
             attributes,
             returnTypeRef as FirJavaTypeRef,
-            receiverTypeRef = null,
             typeParameters,
             valueParameters,
-            body,
             name,
             status,
-            containerSource,
             symbol,
-            annotations,
+            annotationBuilder,
             dispatchReceiverType,
         )
     }

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/builder/FirSimpleFunctionBuilder.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.fir.declarations.builder
 
 import kotlin.contracts.*
-import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.builder.FirAnnotationContainerBuilder
@@ -58,7 +57,6 @@ open class FirSimpleFunctionBuilder : FirFunctionBuilder, FirTypeParametersOwner
     override val annotations: MutableList<FirAnnotationCall> = mutableListOf()
     override val typeParameters: MutableList<FirTypeParameter> = mutableListOf()
 
-    @OptIn(FirImplementationDetail::class)
     override fun build(): FirSimpleFunction {
         return FirSimpleFunctionImpl(
             source,

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/declarations/impl/FirSimpleFunctionImpl.kt

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.declarations.impl
 
-import org.jetbrains.kotlin.fir.FirImplementationDetail
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.FirSourceElement
 import org.jetbrains.kotlin.fir.contracts.FirContractDescription
@@ -31,7 +30,7 @@ import org.jetbrains.kotlin.fir.visitors.*
  * DO NOT MODIFY IT MANUALLY
  */
 
-open class FirSimpleFunctionImpl @FirImplementationDetail constructor(
+internal class FirSimpleFunctionImpl(
     override val source: FirSourceElement?,
     override val session: FirSession,
     override var resolvePhase: FirResolvePhase,

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/ImplementationConfigurator.kt

@@ -449,9 +449,7 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
             default("name", "Name.identifier(\"value\")")
         }
 
-        impl(simpleFunction) {
+        impl(simpleFunction)
-            kind = OpenClass
-        }
 
         impl(safeCallExpression) {
             useTypes(safeCallCheckedSubjectType)