FirReturnsImpliesAnalyzer: find receiver/variable of accessor properly

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/bad/returnsImplies/propertyGetter.kt

@@ -11,8 +11,8 @@ fun Any?.isNotNull(): Boolean {
 @OptIn(ExperimentalContracts::class)
 val Any?.isNotNull: Boolean
     get() {
-        <!WRONG_IMPLIES_CONDITION!>contract {
+        contract {
             returns(true) implies (this@isNotNull != null)
-        }<!>
+        }
         return this@isNotNull != null
     }

compiler/fir/analysis-tests/testData/resolveWithStdlib/contracts/fromSource/good/returnsImplies/propertyAccessors.kt

@@ -6,16 +6,16 @@ interface A {
 
 var Any?.isNotNull: Boolean
     get() {
-        <!WRONG_IMPLIES_CONDITION!>contract {
+        contract {
             returns(true) implies (this@isNotNull != null)
-        }<!>
+        }
         return this != null
     }
     set(value) {
-        <!WRONG_IMPLIES_CONDITION!>contract {
+        contract {
             returns() implies (this@isNotNull != null)
             require(this != null)
-        }<!>
+        }
     }
 
 fun test_1(a: A?) {

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.contracts.coneEffects
 import org.jetbrains.kotlin.fir.contracts.description.*
 import org.jetbrains.kotlin.fir.declarations.FirContractDescriptionOwner
 import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirProperty
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.resolve.dfa.*
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.BlockExitNode
@@ -23,6 +24,7 @@ import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.ControlFlowGraph
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.JumpNode
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertyAccessorSymbol
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.coneType
@@ -61,7 +63,7 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
 
         effects.forEach { effect ->
             val wrongCondition = graph.exitNode.previousCfgNodes.any {
-                isWrongConditionOnNode(it, effect as ConeConditionalEffectDeclaration, function, logicSystem, dataFlowInfo)
+                isWrongConditionOnNode(it, effect as ConeConditionalEffectDeclaration, function, logicSystem, dataFlowInfo, context)
             }
 
             if (wrongCondition) {
@@ -77,7 +79,8 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         effectDeclaration: ConeConditionalEffectDeclaration,
         function: FirFunction<*>,
         logicSystem: LogicSystem<PersistentFlow>,
-        dataFlowInfo: DataFlowInfo
+        dataFlowInfo: DataFlowInfo,
+        context: CheckerContext
     ): Boolean {
         val effect = effectDeclaration.effect as ConeReturnsEffectDeclaration
         val builtinTypes = function.session.builtinTypes
@@ -92,7 +95,7 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
 
         if (isReturn && resultExpression is FirWhenExpression) {
             return node.collectBranchExits().any {
-                isWrongConditionOnNode(it, effectDeclaration, function, logicSystem, dataFlowInfo)
+                isWrongConditionOnNode(it, effectDeclaration, function, logicSystem, dataFlowInfo, context)
             }
         }
 
@@ -110,15 +113,16 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
             }
         }
 
-        val conditionStatements =
-            effectDeclaration.condition.buildTypeStatements(function, logicSystem, dataFlowInfo.variableStorage, flow) ?: return false
+        val conditionStatements = effectDeclaration.condition.buildTypeStatements(
+            function, logicSystem, dataFlowInfo.variableStorage, flow, context
+        ) ?: return false
 
         for ((realVar, requiredTypeStatement) in conditionStatements) {
             val fixedRealVar = typeStatements.keys.find { it.identifier == realVar.identifier } ?: realVar
             val resultTypeStatement = typeStatements[fixedRealVar]
 
             val resultType = mutableListOf<ConeKotlinType>().apply {
-                addIfNotNull(function.getParameterType(fixedRealVar.identifier.symbol))
+                addIfNotNull(function.getParameterType(fixedRealVar.identifier.symbol, context))
                 if (resultTypeStatement != null) addAll(resultTypeStatement.exactType)
             }.let { typeContext.intersectTypesOrNull(it) }
 
@@ -154,11 +158,12 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         function: FirFunction<*>,
         logicSystem: LogicSystem<*>,
         variableStorage: VariableStorage,
-        flow: Flow
+        flow: Flow,
+        context: CheckerContext
     ): MutableTypeStatements? = when (this) {
         is ConeBinaryLogicExpression -> {
-            val left = left.buildTypeStatements(function, logicSystem, variableStorage, flow)
-            val right = right.buildTypeStatements(function, logicSystem, variableStorage, flow)
+            val left = left.buildTypeStatements(function, logicSystem, variableStorage, flow, context)
+            val right = right.buildTypeStatements(function, logicSystem, variableStorage, flow, context)
             if (left != null && right != null) {
                 if (kind == LogicOperationKind.AND) {
                     left.apply { mergeTypeStatements(right) }
@@ -166,16 +171,16 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
             } else (left ?: right)
         }
         is ConeIsInstancePredicate -> {
-            val fir = function.getParameterSymbol(arg.parameterIndex).fir
+            val fir = function.getParameterSymbol(arg.parameterIndex, context).fir
             val realVar = variableStorage.getOrCreateRealVariable(flow, fir.symbol, fir)
             realVar?.to(simpleTypeStatement(realVar, !isNegated, type))?.let { mutableMapOf(it) }
         }
         is ConeIsNullPredicate -> {
-            val fir = function.getParameterSymbol(arg.parameterIndex).fir
+            val fir = function.getParameterSymbol(arg.parameterIndex, context).fir
             val realVar = variableStorage.getOrCreateRealVariable(flow, fir.symbol, fir)
             realVar?.to(simpleTypeStatement(realVar, isNegated, function.session.builtinTypes.anyType.type))?.let { mutableMapOf(it) }
         }
-        is ConeLogicalNot -> arg.buildTypeStatements(function, logicSystem, variableStorage, flow)
+        is ConeLogicalNot -> arg.buildTypeStatements(function, logicSystem, variableStorage, flow, context)
             ?.mapValuesTo(mutableMapOf()) { (_, value) -> value.invert() }
 
         else -> null
@@ -212,11 +217,31 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         return nodes
     }
 
-    private fun FirFunction<*>.getParameterType(symbol: AbstractFirBasedSymbol<*>): ConeKotlinType? {
-        return (if (this.symbol == symbol) receiverTypeRef else valueParameters.find { it.symbol == symbol }?.returnTypeRef)?.coneType
+    private val CheckerContext.containingProperty: FirProperty?
+        get() = (containingDeclarations.asReversed().firstOrNull { it is FirProperty } as? FirProperty)
+
+    private fun FirFunction<*>.getParameterType(symbol: AbstractFirBasedSymbol<*>, context: CheckerContext): ConeKotlinType? {
+        val typeRef = if (this.symbol == symbol) {
+            if (symbol is FirPropertyAccessorSymbol) {
+                context.containingProperty?.receiverTypeRef
+            } else {
+                receiverTypeRef
+            }
+        } else {
+            valueParameters.find { it.symbol == symbol }?.returnTypeRef
+        }
+        return typeRef?.coneType
     }
 
-    private fun FirFunction<*>.getParameterSymbol(index: Int): AbstractFirBasedSymbol<*> {
-        return if (index == -1) this.symbol else this.valueParameters[index].symbol
+    private fun FirFunction<*>.getParameterSymbol(index: Int, context: CheckerContext): AbstractFirBasedSymbol<*> {
+        return if (index == -1) {
+            if (symbol !is FirPropertyAccessorSymbol) {
+                symbol
+            } else {
+                context.containingProperty?.symbol ?: symbol
+            }
+        } else {
+            this.valueParameters[index].symbol
+        }
     }
 }