From 9b903b5e13ea50abde1fdbc15b4c75320d4f67cd Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Fri, 10 Apr 2020 17:08:18 +0300
Subject: [PATCH] [FIR] Process implicit receivers in calls with contract

---
 .../kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt        | 4 ++--
 .../dfa/contracts/ConeConditionalEffectToFirVisitor.kt   | 9 +++------
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
index 6ee262e9eba..be0534c8437 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
@@ -728,8 +728,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 else -> throw IllegalArgumentException("Unsupported constant reference: $value")
             }
         }
-        graphBuilder.exitContract(functionCall).mergeIncomingFlow()
-        contractDescriptionVisitingMode = true
+        graphBuilder.exitContract(functionCall).mergeIncomingFlow(updateReceivers = true)
+        contractDescriptionVisitingMode = false
     }
 
     fun exitConstExpresion(constExpression: FirConstExpression<*>) {
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/contracts/ConeConditionalEffectToFirVisitor.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/contracts/ConeConditionalEffectToFirVisitor.kt
index 475b2bee598..861aadedd68 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/contracts/ConeConditionalEffectToFirVisitor.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/contracts/ConeConditionalEffectToFirVisitor.kt
@@ -8,9 +8,8 @@ package org.jetbrains.kotlin.fir.resolve.dfa.contracts
 import org.jetbrains.kotlin.fir.contracts.description.*
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.expressions.*
-import org.jetbrains.kotlin.fir.references.impl.FirSimpleNamedReference
-import org.jetbrains.kotlin.fir.expressions.FirConstKind
 import org.jetbrains.kotlin.fir.expressions.builder.*
+import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.references.builder.buildSimpleNamedReference
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.util.OperatorNameConventions
@@ -89,10 +88,8 @@ fun ConeConditionalEffectDeclaration.buildContractFir(argumentMapping: Map<Int,
 fun createArgumentsMapping(functionCall: FirFunctionCall): Map<Int, FirExpression>? {
     val function = functionCall.toResolvedCallableSymbol()?.fir as? FirSimpleFunction ?: return null
     val argumentsMapping = mutableMapOf<Int, FirExpression>()
-    // TODO: process implicit receiver
-    // TODO: change to mapping from candidate
-    //  problem: we have already resolved reference without candidate
-    functionCall.explicitReceiver?.let { argumentsMapping[-1] = it }
+    functionCall.extensionReceiver.takeIf { it != FirNoReceiverExpression }?.let { argumentsMapping[-1] = it }
+        ?: functionCall.dispatchReceiver.takeIf { it != FirNoReceiverExpression }?.let { argumentsMapping[-1] = it }
 
 
     val nameToIndex = function.valueParameters.mapIndexed { index, parameter -> parameter.name to index }.toMap()