FIR: remove FirJavaEnhancementContext

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/enhancement/FirAnnotationTypeQualifierResolver.kt

@@ -10,12 +10,11 @@ import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.load.java.AbstractAnnotationTypeQualifierResolver
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames.DEFAULT_ANNOTATION_MEMBER_NAME
-import org.jetbrains.kotlin.load.java.JavaTypeEnhancementState
 import org.jetbrains.kotlin.name.FqName
 
 @Suppress("EXTENSION_SHADOWED_BY_MEMBER")
-class FirAnnotationTypeQualifierResolver(private val session: FirSession, javaTypeEnhancementState: JavaTypeEnhancementState) :
+class FirAnnotationTypeQualifierResolver(private val session: FirSession) :
-    AbstractAnnotationTypeQualifierResolver<FirAnnotationCall>(javaTypeEnhancementState)
+    AbstractAnnotationTypeQualifierResolver<FirAnnotationCall>(session.javaTypeEnhancementState)
 {
     override val FirAnnotationCall.annotations: Iterable<FirAnnotationCall>
         get() = coneClassLikeType?.lookupTag?.toSymbol(session)?.fir?.annotations.orEmpty()

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/enhancement/FirJavaEnhancementContext.kt

@@ -1,31 +0,0 @@
-/*
- * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.fir.java.enhancement
-
-import org.jetbrains.kotlin.fir.FirSession
-import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
-import org.jetbrains.kotlin.load.java.*
-
-class FirJavaEnhancementContext private constructor(
-    val session: FirSession,
-    delegateForDefaultTypeQualifiers: Lazy<JavaTypeQualifiersByElementType?>
-) {
-    constructor(session: FirSession, typeQualifiersComputation: () -> JavaTypeQualifiersByElementType?) :
-            this(session, lazy(LazyThreadSafetyMode.NONE, typeQualifiersComputation))
-
-    val defaultTypeQualifiers: JavaTypeQualifiersByElementType? by delegateForDefaultTypeQualifiers
-}
-
-fun FirJavaEnhancementContext.copyWithNewDefaultTypeQualifiers(
-    typeQualifierResolver: FirAnnotationTypeQualifierResolver,
-    additionalAnnotations: List<FirAnnotationCall>
-): FirJavaEnhancementContext =
-    when {
-        additionalAnnotations.isEmpty() -> this
-        else -> FirJavaEnhancementContext(session) {
-            typeQualifierResolver.extractAndMergeDefaultQualifiers(defaultTypeQualifiers, additionalAnnotations)
-        }
-    }

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/enhancement/SignatureEnhancement.kt

@@ -39,7 +39,6 @@ import org.jetbrains.kotlin.load.java.typeEnhancement.*
 import org.jetbrains.kotlin.load.kotlin.SignatureBuildingComponents
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.load.java.JavaTypeEnhancementState
 import org.jetbrains.kotlin.load.java.JavaTypeQualifiersByElementType
 import org.jetbrains.kotlin.name.FqNameUnsafe
 import org.jetbrains.kotlin.types.AbstractTypeChecker
@@ -58,17 +57,15 @@ class FirSignatureEnhancement(
      *   so owner is a only place where module data can be obtained. However it's guaranteed that `owner`
      *   was created for same session as one passed to constructor, so it's safe to use owners module data
      */
-    private val moduleData = owner.moduleData
+    private val moduleData get() = owner.moduleData
 
     private val javaTypeParameterStack: JavaTypeParameterStack =
         if (owner is FirJavaClass) owner.javaTypeParameterStack else JavaTypeParameterStack.EMPTY
 
-    private val jsr305State: JavaTypeEnhancementState = session.javaTypeEnhancementState
+    private val typeQualifierResolver = FirAnnotationTypeQualifierResolver(session)
 
-    private val typeQualifierResolver = FirAnnotationTypeQualifierResolver(session, jsr305State)
+    private val contextQualifiers: JavaTypeQualifiersByElementType? =
-
+        typeQualifierResolver.extractAndMergeDefaultQualifiers(null, owner.annotations)
-    private val context: FirJavaEnhancementContext =
-        FirJavaEnhancementContext(session) { null }.copyWithNewDefaultTypeQualifiers(typeQualifierResolver, owner.annotations)
 
     private val enhancements = mutableMapOf<FirCallableSymbol<*>, FirCallableSymbol<*>>()
 
@@ -89,6 +86,9 @@ class FirSignatureEnhancement(
         return enhancements.getOrPut(property) { enhance(property, name) } as FirVariableSymbol<*>
     }
 
+    private fun FirAnnotatedDeclaration.computeDefaultQualifiers() =
+        typeQualifierResolver.extractAndMergeDefaultQualifiers(contextQualifiers, annotations)
+
     private fun enhance(
         original: FirVariableSymbol<*>,
         name: Name
@@ -96,14 +96,13 @@ class FirSignatureEnhancement(
         when (val firElement = original.fir) {
             is FirEnumEntry -> {
                 if (firElement.returnTypeRef !is FirJavaTypeRef) return original
-                val memberContext = context.copyWithNewDefaultTypeQualifiers(typeQualifierResolver, firElement.annotations)
                 val predefinedInfo =
                     PredefinedFunctionEnhancementInfo(
                         TypeEnhancementInfo(0 to JavaTypeQualifiers(NullabilityQualifier.NOT_NULL, null, false)),
                         emptyList()
                     )
 
-                val newReturnTypeRef = enhanceReturnType(firElement, emptyList(), memberContext, predefinedInfo)
+                val newReturnTypeRef = enhanceReturnType(firElement, emptyList(), firElement.computeDefaultQualifiers(), predefinedInfo)
                 return firElement.symbol.apply {
                     this.fir.replaceReturnTypeRef(newReturnTypeRef)
                     session.lookupTracker?.recordTypeResolveAsLookup(newReturnTypeRef, this.fir.source, null)
@@ -111,8 +110,7 @@ class FirSignatureEnhancement(
             }
             is FirField -> {
                 if (firElement.returnTypeRef !is FirJavaTypeRef) return original
-                val memberContext = context.copyWithNewDefaultTypeQualifiers(typeQualifierResolver, firElement.annotations)
+                val newReturnTypeRef = enhanceReturnType(firElement, emptyList(), firElement.computeDefaultQualifiers(), null)
-                val newReturnTypeRef = enhanceReturnType(firElement, emptyList(), memberContext, null)
 
                 val symbol = FirFieldSymbol(original.callableId)
                 buildJavaField {
@@ -188,8 +186,6 @@ class FirSignatureEnhancement(
         methodId: CallableId,
         name: Name?
     ): FirFunctionSymbol<*> {
-        val memberContext = context.copyWithNewDefaultTypeQualifiers(typeQualifierResolver, firMethod.annotations)
-
         val predefinedEnhancementInfo =
             SignatureBuildingComponents.signature(
                 owner.symbol.classId,
@@ -204,16 +200,17 @@ class FirSignatureEnhancement(
             }
         }
 
+        val defaultQualifiers = firMethod.computeDefaultQualifiers()
         val overriddenMembers = (firMethod as? FirSimpleFunction)?.overridden().orEmpty()
         val hasReceiver = overriddenMembers.any { it.receiverTypeRef != null }
 
         val newReceiverTypeRef = if (firMethod is FirJavaMethod && hasReceiver) {
-            enhanceReceiverType(firMethod, overriddenMembers, memberContext)
+            enhanceReceiverType(firMethod, overriddenMembers, defaultQualifiers)
         } else null
-        val newReturnTypeRef = if (firMethod !is FirJavaMethod) {
+        val newReturnTypeRef = if (firMethod is FirJavaMethod) {
-            firMethod.returnTypeRef
+            enhanceReturnType(firMethod, overriddenMembers, defaultQualifiers, predefinedEnhancementInfo)
         } else {
-            enhanceReturnType(firMethod, overriddenMembers, memberContext, predefinedEnhancementInfo)
+            firMethod.returnTypeRef
         }
 
         val enhancedValueParameterTypes = mutableListOf<FirResolvedTypeRef>()
@@ -222,7 +219,7 @@ class FirSignatureEnhancement(
             if (hasReceiver && index == 0) continue
             enhancedValueParameterTypes += enhanceValueParameterType(
                 firMethod, overriddenMembers, hasReceiver,
-                memberContext, predefinedEnhancementInfo, valueParameter as FirJavaValueParameter,
+                defaultQualifiers, predefinedEnhancementInfo, valueParameter as FirJavaValueParameter,
                 if (hasReceiver) index - 1 else index
             )
         }
@@ -349,13 +346,13 @@ class FirSignatureEnhancement(
     private fun enhanceTypeParameterBound(typeParameter: FirTypeParameter, bound: FirTypeRef, forceOnlyHeadTypeConstructor: Boolean) =
         EnhancementSignatureParts(
             session, typeQualifierResolver, typeParameter, isCovariant = false, forceOnlyHeadTypeConstructor,
-            AnnotationQualifierApplicabilityType.TYPE_PARAMETER_BOUNDS, context.defaultTypeQualifiers
+            AnnotationQualifierApplicabilityType.TYPE_PARAMETER_BOUNDS, contextQualifiers
         ).enhance(bound, emptyList(), FirJavaTypeConversionMode.TYPE_PARAMETER_BOUND)
 
     fun enhanceSuperType(type: FirTypeRef): FirTypeRef =
         EnhancementSignatureParts(
             session, typeQualifierResolver, null, isCovariant = false, forceOnlyHeadTypeConstructor = false,
-            AnnotationQualifierApplicabilityType.TYPE_USE, context.defaultTypeQualifiers
+            AnnotationQualifierApplicabilityType.TYPE_USE, contextQualifiers
         ).enhance(type, emptyList(), FirJavaTypeConversionMode.SUPERTYPE)
 
     // ================================================================================================
@@ -363,15 +360,13 @@ class FirSignatureEnhancement(
     private fun enhanceReceiverType(
         ownerFunction: FirJavaMethod,
         overriddenMembers: List<FirCallableDeclaration>,
-        memberContext: FirJavaEnhancementContext
+        defaultQualifiers: JavaTypeQualifiersByElementType?
     ): FirResolvedTypeRef {
         return ownerFunction.enhanceValueParameter(
-            typeQualifierResolver,
             overriddenMembers,
-            // TODO: check me
+            ownerFunction,
-            parameterContainer = ownerFunction,
+            defaultQualifiers,
-            methodContext = memberContext,
+            TypeInSignature.Receiver,
-            typeInSignature = TypeInSignature.Receiver,
             predefined = null,
             forAnnotationMember = false
         )
@@ -381,7 +376,7 @@ class FirSignatureEnhancement(
         ownerFunction: FirFunction,
         overriddenMembers: List<FirCallableDeclaration>,
         hasReceiver: Boolean,
-        memberContext: FirJavaEnhancementContext,
+        defaultQualifiers: JavaTypeQualifiersByElementType?,
         predefinedEnhancementInfo: PredefinedFunctionEnhancementInfo?,
         ownerParameter: FirJavaValueParameter,
         index: Int
@@ -390,12 +385,11 @@ class FirSignatureEnhancement(
             return ownerParameter.returnTypeRef as FirResolvedTypeRef
         }
         return ownerFunction.enhanceValueParameter(
-            typeQualifierResolver,
             overriddenMembers,
-            parameterContainer = ownerParameter,
+            ownerParameter,
-            methodContext = memberContext,
+            defaultQualifiers,
-            typeInSignature = TypeInSignature.ValueParameter(hasReceiver, index),
+            TypeInSignature.ValueParameter(hasReceiver, index),
-            predefined = predefinedEnhancementInfo?.parametersInfo?.getOrNull(index),
+            predefinedEnhancementInfo?.parametersInfo?.getOrNull(index),
             forAnnotationMember = owner.classKind == ClassKind.ANNOTATION_CLASS
         )
     }
@@ -403,18 +397,21 @@ class FirSignatureEnhancement(
     private fun enhanceReturnType(
         owner: FirCallableDeclaration,
         overriddenMembers: List<FirCallableDeclaration>,
-        memberContext: FirJavaEnhancementContext,
+        defaultQualifiers: JavaTypeQualifiersByElementType?,
         predefinedEnhancementInfo: PredefinedFunctionEnhancementInfo?
     ): FirResolvedTypeRef {
+        val containerApplicabilityType = if (owner is FirJavaField)
+            AnnotationQualifierApplicabilityType.FIELD
+        else
+            AnnotationQualifierApplicabilityType.METHOD_RETURN_TYPE
         return owner.enhance(
             overriddenMembers,
-            typeContainer = owner,
+            owner,
-            isCovariant = true, containerContext = memberContext,
+            isCovariant = true,
-            containerApplicabilityType =
+            defaultQualifiers,
-            if (owner is FirJavaField) AnnotationQualifierApplicabilityType.FIELD
+            containerApplicabilityType,
-            else AnnotationQualifierApplicabilityType.METHOD_RETURN_TYPE,
+            TypeInSignature.Return,
-            typeInSignature = TypeInSignature.Return,
+            predefinedEnhancementInfo?.returnTypeInfo,
-            predefined = predefinedEnhancementInfo?.returnTypeInfo,
             forAnnotationMember = this.owner.classKind == ClassKind.ANNOTATION_CLASS
         )
     }
@@ -444,20 +441,19 @@ class FirSignatureEnhancement(
     }
 
     private fun FirFunction.enhanceValueParameter(
-        typeQualifierResolver: FirAnnotationTypeQualifierResolver,
         overriddenMembers: List<FirCallableDeclaration>,
-        // TODO: investigate if it's really can be a null (check properties' with extension overrides in Java)
         parameterContainer: FirAnnotationContainer?,
-        methodContext: FirJavaEnhancementContext,
+        defaultQualifiers: JavaTypeQualifiersByElementType?,
         typeInSignature: TypeInSignature,
         predefined: TypeEnhancementInfo?,
         forAnnotationMember: Boolean
-    ): FirResolvedTypeRef = (this as FirCallableDeclaration).enhance(
+    ): FirResolvedTypeRef = enhance(
         overriddenMembers,
-        parameterContainer,
+        parameterContainer ?: this,
-        false, parameterContainer?.let {
+        isCovariant = false,
-            methodContext.copyWithNewDefaultTypeQualifiers(typeQualifierResolver, it.annotations)
+        parameterContainer?.let {
-        } ?: methodContext,
+            typeQualifierResolver.extractAndMergeDefaultQualifiers(defaultQualifiers, it.annotations)
+        } ?: defaultQualifiers,
         AnnotationQualifierApplicabilityType.VALUE_PARAMETER,
         typeInSignature,
         predefined,
@@ -468,7 +464,7 @@ class FirSignatureEnhancement(
         overriddenMembers: List<FirCallableDeclaration>,
         typeContainer: FirAnnotationContainer?,
         isCovariant: Boolean,
-        containerContext: FirJavaEnhancementContext,
+        containerQualifiers: JavaTypeQualifiersByElementType?,
         containerApplicabilityType: AnnotationQualifierApplicabilityType,
         typeInSignature: TypeInSignature,
         predefined: TypeEnhancementInfo?,
@@ -479,7 +475,7 @@ class FirSignatureEnhancement(
         val mode = if (forAnnotationMember) FirJavaTypeConversionMode.ANNOTATION_MEMBER else FirJavaTypeConversionMode.DEFAULT
         return EnhancementSignatureParts(
             session, typeQualifierResolver, typeContainer, isCovariant, forceOnlyHeadTypeConstructor = false,
-            containerApplicabilityType, containerContext.defaultTypeQualifiers
+            containerApplicabilityType, containerQualifiers
         ).enhance(typeRef, typeRefsFromOverridden, mode, predefined)
     }