Use BodyResolveContext.with* functions in FirContractResolveTransformer

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -270,9 +270,15 @@ class BodyResolveContext(
     inline fun <T> withRegularClass(
         regularClass: FirRegularClass,
         holder: SessionHolder,
+        forContracts: Boolean = false,
         crossinline f: () -> T
     ): T {
         storeClassIfNotNested(regularClass)
+        if (forContracts) {
+            return withTypeParametersOf(regularClass) {
+                withContainer(regularClass, f)
+            }
+        }
         return withTowerModeCleanup {
             if (!regularClass.isInner && containerIfAny is FirRegularClass) {
                 towerDataMode = if (regularClass.isCompanion) {
@@ -463,6 +469,7 @@ class BodyResolveContext(
         property: FirProperty,
         accessor: FirPropertyAccessor,
         holder: SessionHolder,
+        forContracts: Boolean = false,
         crossinline f: () -> T
     ): T {
         if (accessor is FirDefaultPropertyAccessor || accessor.body == null) {
@@ -471,7 +478,7 @@ class BodyResolveContext(
         return withTowerDataCleanup {
             val receiverTypeRef = property.receiverTypeRef
             addLocalScope(FirLocalScope())
-            if (receiverTypeRef == null && property.returnTypeRef !is FirImplicitTypeRef &&
+            if (!forContracts && receiverTypeRef == null && property.returnTypeRef !is FirImplicitTypeRef &&
                 !property.isLocal && property.delegate == null
             ) {
                 storeBackingField(property)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/ConeEffectExtractor.kt

@@ -8,8 +8,10 @@ package org.jetbrains.kotlin.fir.resolve.transformers.contracts
 import org.jetbrains.kotlin.contracts.description.EventOccurrencesRange
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSymbolOwner
 import org.jetbrains.kotlin.fir.contracts.description.*
 import org.jetbrains.kotlin.fir.declarations.FirContractDescriptionOwner
+import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.declarations.FirValueParameter
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
@@ -141,12 +143,16 @@ class ConeEffectExtractor(
         }
     }
 
+    private fun FirContractDescriptionOwner.isAccessorOf(declaration: FirSymbolOwner<*>): Boolean {
+        return declaration is FirProperty && (declaration.getter == this || declaration.setter == this)
+    }
+
     override fun visitThisReceiverExpression(
         thisReceiverExpression: FirThisReceiverExpression,
         data: Nothing?
     ): ConeContractDescriptionElement? {
         val declaration = thisReceiverExpression.calleeReference.boundSymbol?.fir ?: return null
-        return if (declaration == owner) {
+        return if (declaration == owner || owner.isAccessorOf(declaration)) {
             val type = thisReceiverExpression.typeRef.coneType
             toValueParameterReference(type, -1, "this")
         } else {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/contracts/FirContractResolveTransformer.kt

@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirBodyResolve
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirDeclarationsResolveTransformer
 import org.jetbrains.kotlin.fir.symbols.impl.FirAnonymousFunctionSymbol
 import org.jetbrains.kotlin.fir.types.builder.buildImplicitTypeRef
-import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.visitors.*
 import org.jetbrains.kotlin.name.Name
 
@@ -67,15 +66,8 @@ open class FirContractResolveTransformer(
                 return simpleFunction.compose()
             }
             @Suppress("UNCHECKED_CAST")
-            return withTypeParametersOf(simpleFunction) {
-                val receiverTypeRef = simpleFunction.receiverTypeRef
-                if (receiverTypeRef != null) {
-                    withLabelAndReceiverType(simpleFunction.name, simpleFunction, receiverTypeRef.coneType) {
-                        transformContractDescriptionOwner(simpleFunction)
-                    }
-                } else {
-                    transformContractDescriptionOwner(simpleFunction)
-                }
+            return context.withSimpleFunction(simpleFunction, components) {
+                transformContractDescriptionOwner(simpleFunction)
             }
         }
 
@@ -91,13 +83,9 @@ open class FirContractResolveTransformer(
                 transformSimpleFunction(property.getter.delegate, data)
                 return property.compose()
             }
-            withTypeParametersOf(property) {
-                withLocalScopeCleanup {
-                    context.withContainer(property) {
-                        property.getter?.let { transformPropertyAccessor(it, property) }
-                        property.setter?.let { transformPropertyAccessor(it, property) }
-                    }
-                }
+            context.withProperty(property) {
+                property.getter?.let { transformPropertyAccessor(it, property) }
+                property.setter?.let { transformPropertyAccessor(it, property) }
             }
             return property.compose()
         }
@@ -115,12 +103,7 @@ open class FirContractResolveTransformer(
             if (!propertyAccessor.hasContractToResolve) {
                 return propertyAccessor.compose()
             }
-            val receiverTypeRef = owner.receiverTypeRef
-            return if (receiverTypeRef != null) {
-                withLabelAndReceiverType(owner.name, propertyAccessor, receiverTypeRef.coneType) {
-                    transformContractDescriptionOwner(propertyAccessor)
-                }
-            } else {
+            return context.withPropertyAccessor(owner, propertyAccessor, components, forContracts = true) {
                 transformContractDescriptionOwner(propertyAccessor)
             }
         }
@@ -141,14 +124,10 @@ open class FirContractResolveTransformer(
             contractDescription: FirLegacyRawContractDescription
         ): CompositeTransformResult<T> {
             val valueParameters = owner.valueParameters
-            val contractCall = withNewLocalScope {
-                for (valueParameter in valueParameters) {
-                    context.storeVariable(valueParameter)
-                }
-                context.withContainer(owner as FirDeclaration) {
-                    contractDescription.contractCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
-                }
+            for (valueParameter in valueParameters) {
+                context.storeVariable(valueParameter)
             }
+            val contractCall = contractDescription.contractCall.transformSingle(transformer, ResolutionMode.ContextIndependent)
             val resolvedId = contractCall.toResolvedCallableSymbol()?.callableId ?: return transformOwnerWithUnresolvedContract(owner)
             if (resolvedId != FirContractsDslNames.CONTRACT) return transformOwnerWithUnresolvedContract(owner)
             if (contractCall.arguments.size != 1) return transformOwnerOfErrorContract(owner)
@@ -233,12 +212,9 @@ open class FirContractResolveTransformer(
 
         override fun transformRegularClass(regularClass: FirRegularClass, data: ResolutionMode): CompositeTransformResult<FirStatement> {
             regularClass.updatePhase()
-            context.storeClassIfNotNested(regularClass)
             regularClass.transformCompanionObject(this, data)
-            withTypeParametersOf(regularClass) {
-                context.withContainer(regularClass) {
-                    regularClass.transformDeclarations(this, data)
-                }
+            context.withRegularClass(regularClass, components, forContracts = true) {
+                regularClass.transformDeclarations(this, data)
             }
             return regularClass.compose()
         }
@@ -248,7 +224,7 @@ open class FirContractResolveTransformer(
             data: ResolutionMode
         ): CompositeTransformResult<FirStatement> {
             anonymousObject.updatePhase()
-            context.withContainer(anonymousObject) {
+            context.withAnonymousObject(anonymousObject, components) {
                 anonymousObject.transformDeclarations(this, data)
             }
             return anonymousObject.compose()