[FIR] Add smartcasts on implicit receivers

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ImplicitReceiverStack.kt

@@ -9,31 +9,67 @@ import com.google.common.collect.LinkedHashMultimap
 import com.google.common.collect.SetMultimap
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitDispatchReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
+import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.name.Name
 
-class ImplicitReceiverStack {
+interface ImplicitReceiverStack {
+    fun add(name: Name, value: ImplicitReceiverValue<*>)
+    fun pop(name: Name)
+
+    operator fun get(name: String?): ImplicitReceiverValue<*>?
+
+    fun lastDispatchReceiver(): ImplicitDispatchReceiverValue?
+    fun receiversAsReversed(): List<ImplicitReceiverValue<*>>
+}
+
+class ImplicitReceiverStackImpl : ImplicitReceiverStack, Iterable<ImplicitReceiverValue<*>> {
     private val stack: MutableList<ImplicitReceiverValue<*>> = mutableListOf()
+    private val originalTypes: MutableList<ConeKotlinType> = mutableListOf()
     // This multi-map holds indexes of the stack ^
     private val indexesPerLabel: SetMultimap<Name, Int> = LinkedHashMultimap.create()
+    private val indexesPerSymbol: MutableMap<FirBasedSymbol<*>, Int> = mutableMapOf()
+    val size: Int get() = stack.size
 
-    fun add(name: Name, value: ImplicitReceiverValue<*>) {
+    override fun add(name: Name, value: ImplicitReceiverValue<*>) {
         stack += value
-        indexesPerLabel.put(name, stack.size - 1)
+        originalTypes += value.type
+        val index = stack.size - 1
+        indexesPerLabel.put(name, index)
+        indexesPerSymbol.put(value.boundSymbol, index)
     }
 
-    fun pop(name: Name) {
+    override fun pop(name: Name) {
-        indexesPerLabel.remove(name, stack.size - 1)
+        val index = stack.size - 1
-        stack.removeAt(stack.size - 1)
+        indexesPerLabel.remove(name, index)
+        originalTypes.removeAt(index)
+        val value = stack.removeAt(index)
+        indexesPerSymbol.remove(value.boundSymbol)
     }
 
-    operator fun get(name: String?): ImplicitReceiverValue<*>? {
+    override operator fun get(name: String?): ImplicitReceiverValue<*>? {
         if (name == null) return stack.lastOrNull()
         return indexesPerLabel[Name.identifier(name)].lastOrNull()?.let { stack[it] }
     }
 
-    fun lastDispatchReceiver(): ImplicitDispatchReceiverValue? {
+    override fun lastDispatchReceiver(): ImplicitDispatchReceiverValue? {
         return stack.filterIsInstance<ImplicitDispatchReceiverValue>().lastOrNull()
     }
 
-    fun receiversAsReversed(): List<ImplicitReceiverValue<*>> = stack.asReversed()
+    override fun receiversAsReversed(): List<ImplicitReceiverValue<*>> = stack.asReversed()
+
+    fun getReceiverIndex(symbol: FirBasedSymbol<*>): Int? = indexesPerSymbol[symbol]
+
+    fun getOriginalType(index: Int): ConeKotlinType {
+        return originalTypes[index]
+    }
+
+    fun replaceReceiverType(index: Int, type: ConeKotlinType) {
+        assert(index >= 0 && index < stack.size)
+        stack[index].replaceType(type)
+    }
+
+    override operator fun iterator(): Iterator<ImplicitReceiverValue<*>> {
+        return stack.iterator()
+    }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/FirReceivers.kt

@@ -60,15 +60,28 @@ class ExpressionReceiverValue(
 
 abstract class ImplicitReceiverValue<S : AbstractFirBasedSymbol<*>>(
     val boundSymbol: S,
-    final override val type: ConeKotlinType,
+    type: ConeKotlinType,
-    useSiteSession: FirSession,
+    private val useSiteSession: FirSession,
-    scopeSession: ScopeSession
+    private val scopeSession: ScopeSession
 ) : ReceiverValue {
-    val implicitScope: FirScope? = type.scope(useSiteSession, scopeSession)
+    final override var type: ConeKotlinType = type
+        private set
+
+    var implicitScope: FirScope? = type.scope(useSiteSession, scopeSession)
+        private set
 
     override fun scope(useSiteSession: FirSession, scopeSession: ScopeSession): FirScope? = implicitScope
 
     override val receiverExpression: FirExpression = receiverExpression(boundSymbol, type)
+
+    /*
+     * Should be called only in ImplicitReceiverStack
+     */
+    internal fun replaceType(type: ConeKotlinType) {
+        if (type == this.type) return
+        this.type = type
+        implicitScope = type.scope(useSiteSession, scopeSession)
+    }
 }
 
 class ImplicitDispatchReceiverValue(

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/DataFlowVariable.kt

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.FirThisReference
+import org.jetbrains.kotlin.fir.expressions.impl.FirThisReceiverExpressionImpl
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 
 /*
@@ -15,6 +17,7 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 sealed class DataFlowVariable(val index: Int, val fir: FirElement) {
     abstract val isSynthetic: Boolean
     abstract val real: DataFlowVariable
+    abstract val isThisReference: Boolean
 
     final override fun hashCode(): Int {
         return index
@@ -30,7 +33,7 @@ sealed class DataFlowVariable(val index: Int, val fir: FirElement) {
     }
 }
 
-private class RealDataFlowVariable(index: Int, fir: FirElement) : DataFlowVariable(index, fir) {
+private class RealDataFlowVariable(index: Int, fir: FirElement, override val isThisReference: Boolean) : DataFlowVariable(index, fir) {
     override val isSynthetic: Boolean get() = false
 
     override val real: DataFlowVariable get() = this
@@ -40,12 +43,16 @@ private class SyntheticDataFlowVariable(index: Int, fir: FirElement) : DataFlowV
     override val isSynthetic: Boolean get() = true
 
     override val real: DataFlowVariable get() = this
+
+    override val isThisReference: Boolean get() = false
 }
 
 private class AliasedDataFlowVariable(index: Int, fir: FirElement, var delegate: DataFlowVariable) : DataFlowVariable(index, fir) {
     override val isSynthetic: Boolean get() = delegate.isSynthetic
 
     override val real: DataFlowVariable get() = delegate.real
+
+    override val isThisReference: Boolean get() = false
 }
 
 
@@ -54,9 +61,17 @@ class DataFlowVariableStorage {
     private var counter: Int = 1
 
     fun getOrCreateNewRealVariable(symbol: FirBasedSymbol<*>): DataFlowVariable {
+        return getOrCreateNewRealVariableImpl(symbol, false)
+    }
+
+    fun getOrCreateNewThisRealVariable(symbol: FirBasedSymbol<*>): DataFlowVariable {
+        return getOrCreateNewRealVariableImpl(symbol, true)
+    }
+
+    private fun getOrCreateNewRealVariableImpl(symbol: FirBasedSymbol<*>, isThisReference: Boolean): DataFlowVariable {
         val fir = symbol.fir
         get(fir)?.let { return it }
-        return RealDataFlowVariable(counter++, fir).also { storeVariable(it, fir) }
+        return RealDataFlowVariable(counter++, fir, isThisReference).also { storeVariable(it, fir) }
     }
 
     fun getOrCreateNewSyntheticVariable(fir: FirElement): DataFlowVariable {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -11,13 +11,15 @@ import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
 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.expressions.impl.FirThisReceiverExpressionImpl
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
+import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStackImpl
 import org.jetbrains.kotlin.fir.resolve.dfa.Condition.*
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.resolve.transformers.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
+import org.jetbrains.kotlin.fir.symbols.FirSymbolOwner
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
@@ -33,6 +35,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     }
 
     private val context: DataFlowInferenceContext get() = inferenceComponents.ctx as DataFlowInferenceContext
+    private val receiverStack: ImplicitReceiverStackImpl = transformer.implicitReceiverStack
 
     private val graphBuilder = ControlFlowGraphBuilder()
     private val logicSystem = LogicSystem(context)
@@ -43,8 +46,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
      * If there is no types from smartcasts function returns null
      */
     fun getTypeUsingSmartcastInfo(qualifiedAccessExpression: FirQualifiedAccessExpression): Collection<ConeKotlinType>? {
-        val symbol: FirBasedSymbol<*> = qualifiedAccessExpression.resolvedSymbol ?: return null
+        val variable = qualifiedAccessExpression.variable?.real ?: return null
-        val variable = variableStorage[symbol]?.real ?: return null
         return graphBuilder.lastNode.flow.approvedFacts(variable)?.exactType ?: return null
     }
 
@@ -54,7 +56,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         graphBuilder.enterFunction(function).passFlow()
 
         for (valueParameter in function.valueParameters) {
-            getRealVariable(valueParameter.symbol)
+            getOrCreateRealVariable(valueParameter.symbol)
         }
     }
 
@@ -90,7 +92,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
         val previousNode = node.usefulPreviousNodes.singleOrNull() as? WhenBranchConditionExitNode
         if (previousNode != null) {
-            node.flow = logicSystem.approveFactsInsideFlow(previousNode.variable, EqTrue, node.flow)
+            node.flow = approveFactsAndUpdateImplicitReceivers(previousNode.variable, EqTrue, node.flow)
         }
     }
 
@@ -100,25 +102,17 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     // ----------------------------------- Operator call -----------------------------------
 
-    private fun FirExpression.getResolvedSymbol(): FirCallableSymbol<*>? {
-        val expression = (this as? FirWhenSubjectExpression)?.whenSubject?.whenExpression?.let {
-            it.subjectVariable?.symbol?.let { return it }
-            it.subject
-        } ?: this
-        return expression.toResolvedCallableSymbol() as? FirCallableSymbol<*>
-    }
-
     fun exitTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall) {
         val node = graphBuilder.exitTypeOperatorCall(typeOperatorCall).passFlow(false)
+
         if (typeOperatorCall.operation !in FirOperation.TYPES) return
-        val symbol: FirCallableSymbol<*> = typeOperatorCall.argument.getResolvedSymbol() ?: return
+        val type = typeOperatorCall.conversionTypeRef.coneTypeSafe<ConeKotlinType>() ?: return
-            val type = typeOperatorCall.conversionTypeRef.coneTypeSafe<ConeKotlinType>() ?: return
+        val varVariable = getOrCreateRealVariable(typeOperatorCall.argument)?.real ?: return
-        val varVariable = getRealVariable(symbol)
 
         var flow = node.flow
         when (typeOperatorCall.operation) {
             FirOperation.IS, FirOperation.NOT_IS -> {
-                val expressionVariable = getSyntheticVariable(typeOperatorCall)
+                val expressionVariable = getOrCreateSyntheticVariable(typeOperatorCall)
 
                 val trueInfo = FirDataFlowInfo(setOf(type), emptySet())
                 val falseInfo = FirDataFlowInfo(emptySet(), setOf(type))
@@ -142,11 +136,11 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             }
 
             FirOperation.AS -> {
-                flow = flow.addApprovedFact(varVariable, FirDataFlowInfo(setOf(type), emptySet()))
+                flow = addApprovedFact(flow, varVariable, FirDataFlowInfo(setOf(type), emptySet()))
             }
 
             FirOperation.SAFE_AS -> {
-                val expressionVariable = getSyntheticVariable(typeOperatorCall)
+                val expressionVariable = getOrCreateSyntheticVariable(typeOperatorCall)
                 flow = flow.addNotApprovedFact(
                     expressionVariable,
                     UnapprovedFirDataFlowInfo(
@@ -195,7 +189,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             else -> return
         }
 
-        val expressionVariable = getVariable(node.fir)
+        val expressionVariable = getOrCreateVariable(node.fir)
         var flow = node.flow
 
         // not null for comparisons with constants
@@ -239,7 +233,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     private fun processEqNull(node: OperatorCallNode, operand: FirExpression, operation: FirOperation) {
         var flow = node.flow
-        val expressionVariable = getVariable(node.fir)
+        val expressionVariable = getOrCreateVariable(node.fir)
 
         variableStorage[operand]?.let { operandVariable ->
             val condition = when (operation) {
@@ -312,7 +306,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         val node = graphBuilder.enterWhenBranchCondition(whenBranch).passFlow(false)
         val previousNode = node.previousNodes.single()
         if (previousNode is WhenBranchConditionExitNode) {
-            node.flow = logicSystem.approveFactsInsideFlow(previousNode.variable, EqFalse, node.flow)
+            node.flow = approveFactsAndUpdateImplicitReceivers(previousNode.variable, EqFalse, node.flow)
         }
         node.flow.freeze()
     }
@@ -320,19 +314,19 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     fun exitWhenBranchCondition(whenBranch: FirWhenBranch) {
         val node = graphBuilder.exitWhenBranchCondition(whenBranch).passFlow(true)
 
-        val conditionVariable = getVariable(whenBranch.condition)
+        val conditionVariable = getOrCreateVariable(whenBranch.condition)
         node.variable = conditionVariable
     }
 
     fun exitWhenBranchResult(whenBranch: FirWhenBranch) {
         val node = graphBuilder.exitWhenBranchResult(whenBranch).passFlow(false)
-        val conditionVariable = getVariable(whenBranch.condition)
+        val conditionVariable = getOrCreateVariable(whenBranch.condition)
         node.flow = node.flow.removeSyntheticVariable(conditionVariable).apply { freeze() }
     }
 
     fun exitWhenExpression(whenExpression: FirWhenExpression) {
-        val node = graphBuilder.exitWhenExpression(whenExpression)
+        val node = graphBuilder.exitWhenExpression(whenExpression).passFlow()
-        var flow = logicSystem.or(node.usefulPreviousNodes.map { it.flow })
+        var flow = node.flow
         val subjectSymbol = whenExpression.subjectVariable?.symbol
         if (subjectSymbol != null) {
             variableStorage[subjectSymbol]?.let { flow = flow.removeVariable(it) }
@@ -426,7 +420,17 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         }
     }
 
-    private val FirElement.resolvedSymbol: FirBasedSymbol<*>? get() = this.safeAs<FirResolvable>()?.calleeReference.safeAs<FirResolvedCallableReference>()?.coneSymbol as? FirBasedSymbol<*>
+    private val FirElement.resolvedSymbol: FirBasedSymbol<*>?
+        get() {
+            val expression = (this as? FirWhenSubjectExpression)?.whenSubject?.whenExpression?.let {
+                it.subjectVariable?.symbol?.let { symbol -> return symbol }
+                it.subject
+            } ?: this
+            return (expression as? FirResolvable)?.resolvedSymbol
+        }
+
+    private val FirResolvable.resolvedSymbol: FirBasedSymbol<*>?
+        get() = (calleeReference as? FirResolvedCallableReference)?.coneSymbol as? FirBasedSymbol<*>
 
     private fun FirFunctionCall.isBooleanNot(): Boolean {
         val symbol = calleeReference.safeAs<FirResolvedCallableReference>()?.coneSymbol as? FirNamedFunctionSymbol ?: return false
@@ -452,7 +456,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
          */
         variableStorage[initializer]?.let { initializerVariable ->
             assert(initializerVariable.isSynthetic)
-            val realVariable = getRealVariable(variable.symbol)
+            val realVariable = getOrCreateRealVariable(variable.symbol)
             node.flow = node.flow.copyNotApprovedFacts(initializerVariable, realVariable)
         }
 
@@ -483,8 +487,9 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         val (leftNode, rightNode) = graphBuilder.exitLeftBinaryAndArgument(binaryLogicExpression)
         leftNode.passFlow(true)
         rightNode.passFlow(false)
-        val leftOperandVariable = getVariable(leftNode.previousNodes.first().fir)
+        val leftOperandVariable = getOrCreateVariable(leftNode.previousNodes.first().fir)
-        rightNode.flow = logicSystem.approveFactsInsideFlow(leftOperandVariable, EqTrue, rightNode.flow).also { it.freeze() }
+        val flow = approveFactsAndUpdateImplicitReceivers(leftOperandVariable, EqTrue, rightNode.flow)
+        rightNode.flow = flow.also { it.freeze() }
     }
 
     fun exitBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression) {
@@ -495,7 +500,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         val flowFromRight = node.rightOperandNode.flow
         val flow = node.flow
 
-        val andVariable = getVariable(binaryLogicExpression)
+        val andVariable = getOrCreateVariable(binaryLogicExpression)
 
         val leftIsTrue = approveFact(leftVariable, EqTrue, flowFromRight)
         val leftIsFalse = approveFact(leftVariable, EqFalse, flowFromRight)
@@ -528,8 +533,9 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         val (leftNode, rightNode) = graphBuilder.exitLeftBinaryOrArgument(binaryLogicExpression)
         leftNode.passFlow(true)
         rightNode.passFlow(false)
-        val leftOperandVariable = getVariable(leftNode.previousNodes.first().fir)
+        val leftOperandVariable = getOrCreateVariable(leftNode.previousNodes.first().fir)
-        rightNode.flow = logicSystem.approveFactsInsideFlow(leftOperandVariable, EqFalse, rightNode.flow).also { it.freeze() }
+        val flow = approveFactsAndUpdateImplicitReceivers(leftOperandVariable, EqFalse, rightNode.flow)
+        rightNode.flow = flow.also { it.freeze() }
     }
 
     fun exitBinaryOr(binaryLogicExpression: FirBinaryLogicExpression) {
@@ -540,7 +546,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         val flowFromRight = node.rightOperandNode.flow
         val flow = node.flow
 
-        val orVariable = getVariable(binaryLogicExpression)
+        val orVariable = getOrCreateVariable(binaryLogicExpression)
 
         val leftIsTrue = approveFact(leftVariable, EqTrue, flowFromLeft)
         val leftIsFalse = approveFact(leftVariable, EqFalse, flowFromLeft)
@@ -561,8 +567,8 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     }
 
     private fun exitBooleanNot(functionCall: FirFunctionCall, node: FunctionCallNode) {
-        val booleanExpressionVariable = getVariable(node.previousNodes.first().fir)
+        val booleanExpressionVariable = getOrCreateVariable(node.previousNodes.first().fir)
-        val variable = getVariable(functionCall)
+        val variable = getOrCreateVariable(functionCall)
         node.flow = node.flow.copyNotApprovedFacts(booleanExpressionVariable, variable) { it.invert() }
     }
 
@@ -592,7 +598,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         logicSystem.approveFact(variable, condition, flow)
 
     private fun FirBinaryLogicExpression.getVariables(): Pair<DataFlowVariable, DataFlowVariable> =
-        getVariable(leftOperand) to getVariable(rightOperand)
+        getOrCreateVariable(leftOperand) to getOrCreateVariable(rightOperand)
 
     private var CFGNode<*>.flow: Flow
         get() = edges[this] ?: Flow.EMPTY
@@ -601,22 +607,58 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         }
 
     private fun <T : CFGNode<*>> T.passFlow(shouldFreeze: Boolean = false): T = this.also { node ->
-        node.flow = logicSystem.or(node.usefulPreviousNodes.map { it.flow }).also {
+        val previousFlows = node.usefulPreviousNodes.map { it.flow }
+        val flow = logicSystem.or(previousFlows).also {
             if (shouldFreeze) it.freeze()
         }
+        if (previousFlows.size > 1) {
+            receiverStack.forEachIndexed { index, receiver ->
+                val variable = variableStorage[receiver.boundSymbol] ?: return@forEachIndexed
+                val approvedFacts = flow.approvedFacts(variable)
+                if (approvedFacts == null) {
+                    receiver.replaceType(receiverStack.getOriginalType(index))
+                } else {
+                    updateReceiverType(index, approvedFacts)
+                }
+            }
+        }
+        node.flow = flow
     }
 
-    private fun getSyntheticVariable(fir: FirElement): DataFlowVariable = variableStorage.getOrCreateNewSyntheticVariable(fir)
+    // -------------------------------- get or create variable --------------------------------
-    private fun getRealVariable(symbol: FirBasedSymbol<*>): DataFlowVariable = variableStorage.getOrCreateNewRealVariable(symbol).real
+    private fun getOrCreateSyntheticVariable(fir: FirElement): DataFlowVariable = variableStorage.getOrCreateNewSyntheticVariable(fir)
 
-    private fun getVariable(fir: FirElement): DataFlowVariable {
+    private fun getOrCreateRealVariable(fir: FirElement): DataFlowVariable? {
+        if (fir is FirThisReceiverExpressionImpl) {
+            return variableStorage.getOrCreateNewThisRealVariable(fir.calleeReference.boundSymbol ?: return null)
+        }
+        val symbol: FirBasedSymbol<*> = fir.resolvedSymbol ?: return null
+        return variableStorage.getOrCreateNewRealVariable(symbol)
+    }
+
+    private fun getOrCreateRealVariable(symbol: FirBasedSymbol<*>): DataFlowVariable =
+        variableStorage.getOrCreateNewRealVariable(symbol).real
+
+    private fun getOrCreateVariable(fir: FirElement): DataFlowVariable {
         val symbol = fir.resolvedSymbol
         return if (symbol == null)
-            getSyntheticVariable(fir)
+            getOrCreateSyntheticVariable(fir)
         else
-            getRealVariable(symbol)
+            getOrCreateRealVariable(symbol)
     }
 
+    // -------------------------------- get variable --------------------------------
+
+    private val FirElement.variable: DataFlowVariable?
+        get() {
+            val symbol: FirBasedSymbol<*> = if (this is FirThisReceiverExpressionImpl) {
+                calleeReference.boundSymbol
+            } else {
+                resolvedSymbol
+            } ?: return null
+            return variableStorage[symbol]
+        }
+
     private fun getRealVariablesForSafeCallChain(call: FirExpression): Collection<DataFlowVariable> {
         val result = mutableListOf<DataFlowVariable>()
 
@@ -630,12 +672,13 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
                     }
                     ((call.calleeReference as? FirResolvedCallableReference)?.coneSymbol)?.let { symbol ->
                         if (symbol is FirVariableSymbol<*> || symbol is FirPropertySymbol) {
-                            result += getRealVariable(symbol as FirBasedSymbol<*>)
+                            result += getOrCreateRealVariable(symbol as FirBasedSymbol<*>)
                         }
                     }
                 }
                 is FirWhenSubjectExpression -> {
-                    call.whenSubject.whenExpression.subjectVariable?.let { result += getRealVariable(it.symbol) }
+                    // TODO: check
+                    call.whenSubject.whenExpression.subjectVariable?.let { result += getOrCreateRealVariable(it.symbol) }
                     call.whenSubject.whenExpression.subject?.let { collect(it) }
                 }
             }
@@ -645,6 +688,37 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         return result
     }
 
+    private fun approveFactsAndUpdateImplicitReceivers(variable: DataFlowVariable, condition: Condition, flow: Flow): Flow {
+        val (newFlow, variables) = logicSystem.approveFactsInsideFlow(variable, condition, flow)
+        for (variable in variables) {
+            updateReceiverType(newFlow, variable)
+        }
+        return newFlow
+    }
+
+    private fun updateReceiverType(flow: Flow, variable: DataFlowVariable) {
+        val symbol = (variable.fir as? FirSymbolOwner<*>)?.symbol ?: return
+        val index = receiverStack.getReceiverIndex(symbol) ?: return
+        val info = flow.approvedFacts[variable] ?: return
+        updateReceiverType(index, info)
+    }
+
+    private fun updateReceiverType(index: Int, info: FirDataFlowInfo) {
+        val types = info.exactType.toMutableList().also {
+            it += receiverStack.getOriginalType(index)
+        }
+        receiverStack.replaceReceiverType(index, context.myIntersectTypes(types)!!)
+
+    }
+
+    private fun addApprovedFact(flow: Flow, variable: DataFlowVariable, info: FirDataFlowInfo): Flow {
+        return flow.addApprovedFact(variable, info).also {
+            if (variable.isThisReference) {
+                updateReceiverType(flow, variable)
+            }
+        }
+    }
+
     private fun Flow.removeVariable(variable: DataFlowVariable): Flow {
         variableStorage.removeVariable(variable)
         return removeVariableFromFlow(variable)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/LogicSystem.kt

@@ -69,10 +69,10 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
         return map
     }
 
-    fun approveFactsInsideFlow(variable: DataFlowVariable, condition: Condition, flow: Flow): Flow {
+    fun approveFactsInsideFlow(variable: DataFlowVariable, condition: Condition, flow: Flow): Pair<Flow, Collection<DataFlowVariable>> {
         val notApprovedFacts: Set<UnapprovedFirDataFlowInfo> = flow.notApprovedFacts[variable]
         if (notApprovedFacts.isEmpty()) {
-            return flow
+            return flow to emptyList()
         }
         @Suppress("NAME_SHADOWING")
         val flow = flow.copyForBuilding()
@@ -82,6 +82,8 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
                 newFacts.put(it.variable, it.info)
             }
         }
+        val updatedReceivers = mutableSetOf<DataFlowVariable>()
+
         newFacts.asMap().forEach { (variable, infos) ->
             @Suppress("NAME_SHADOWING")
             val infos = ArrayList(infos)
@@ -89,8 +91,11 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
                 infos.add(it)
             }
             flow.approvedFacts[variable] = context.and(infos)
+            if (variable.isThisReference) {
+                updatedReceivers += variable
+            }
         }
-        return flow
+        return flow to updatedReceivers
     }
 
     fun approveFact(variable: DataFlowVariable, condition: Condition, flow: Flow): MutableMap<DataFlowVariable, FirDataFlowInfo> {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirBodyResolveTransformer.kt

@@ -72,7 +72,7 @@ open class FirBodyResolveTransformer(
 
     private val localScopes = mutableListOf<FirLocalScope>()
     private val topLevelScopes = mutableListOf<FirScope>()
-    final override val implicitReceiverStack: ImplicitReceiverStack = ImplicitReceiverStack()
+    final override val implicitReceiverStack: ImplicitReceiverStackImpl = ImplicitReceiverStackImpl()
     final override val inferenceComponents = inferenceComponents(session, returnTypeCalculator, scopeSession)
 
     private var primaryConstructorParametersScope: FirLocalScope? = null

compiler/fir/resolve/testData/resolve/smartcasts/implicitReceivers.dot

@@ -0,0 +1,309 @@
+digraph implicitReceivers_kt {
+  graph [splines=ortho nodesep=3]
+  node [shape=box penwidth=2]
+  edge [penwidth=2]
+
+  subgraph cluster_0 {
+    color=red
+    0 [label="Enter function <init>" style="filled" fillcolor=red];
+    1 [label="Exit function <init>" style="filled" fillcolor=red];
+  }
+
+  0 -> {1};
+
+  subgraph cluster_1 {
+    color=red
+    2 [label="Enter function foo" style="filled" fillcolor=red];
+    subgraph cluster_2 {
+      color=blue
+      3 [label="Enter block"];
+      4 [label="Exit block"];
+    }
+    5 [label="Exit function foo" style="filled" fillcolor=red];
+  }
+
+  2 -> {3};
+  3 -> {4};
+  4 -> {5};
+
+  subgraph cluster_3 {
+    color=red
+    6 [label="Enter function with" style="filled" fillcolor=red];
+    subgraph cluster_4 {
+      color=blue
+      7 [label="Enter block"];
+      8 [label="Exit block"];
+    }
+    9 [label="Exit function with" style="filled" fillcolor=red];
+  }
+
+  6 -> {7};
+  7 -> {8};
+  8 -> {9};
+
+  subgraph cluster_5 {
+    color=red
+    10 [label="Enter function test_1" style="filled" fillcolor=red];
+    subgraph cluster_6 {
+      color=blue
+      11 [label="Enter block"];
+      subgraph cluster_7 {
+        color=blue
+        12 [label="Enter when"];
+        subgraph cluster_8 {
+          color=blue
+          13 [label="Enter when branch condition "];
+          14 [label="Access variable this@R|/test_1|"];
+          15 [label="Type operator: this is A"];
+          16 [label="Exit when branch condition"];
+        }
+        subgraph cluster_9 {
+          color=blue
+          17 [label="Enter block"];
+          18 [label="Access variable this@R|/test_1|"];
+          19 [label="Function call: this@R|/test_1|.R|/A.foo|()"];
+          20 [label="Function call: this@R|/A|.R|/A.foo|()"];
+          21 [label="Exit block"];
+        }
+        22 [label="Exit when branch result"];
+        subgraph cluster_10 {
+          color=blue
+          23 [label="Enter when branch condition else"];
+          24 [label="Exit when branch condition"];
+        }
+        subgraph cluster_11 {
+          color=blue
+          25 [label="Enter block"];
+          26 [label="Access variable this@R|/test_1|"];
+          27 [label="Function call: this@R|/test_1|.<Unresolved name: foo>#()"];
+          28 [label="Function call: <Unresolved name: foo>#()"];
+          29 [label="Exit block"];
+        }
+        30 [label="Exit when branch result"];
+        31 [label="Exit when"];
+      }
+      32 [label="Access variable this@R|/test_1|"];
+      33 [label="Function call: this@R|/test_1|.<Unresolved name: foo>#()"];
+      34 [label="Function call: <Unresolved name: foo>#()"];
+      35 [label="Exit block"];
+    }
+    36 [label="Exit function test_1" style="filled" fillcolor=red];
+  }
+
+  10 -> {11};
+  11 -> {12};
+  12 -> {13};
+  13 -> {14};
+  14 -> {15};
+  15 -> {16};
+  16 -> {17 23};
+  17 -> {18};
+  18 -> {19};
+  19 -> {20};
+  20 -> {21};
+  21 -> {22};
+  22 -> {31};
+  23 -> {24};
+  24 -> {25};
+  25 -> {26};
+  26 -> {27};
+  27 -> {28};
+  28 -> {29};
+  29 -> {30};
+  30 -> {31};
+  31 -> {32};
+  32 -> {33};
+  33 -> {34};
+  34 -> {35};
+  35 -> {36};
+
+  subgraph cluster_12 {
+    color=red
+    37 [label="Enter function test_2" style="filled" fillcolor=red];
+    subgraph cluster_13 {
+      color=blue
+      38 [label="Enter block"];
+      subgraph cluster_14 {
+        color=blue
+        39 [label="Enter when"];
+        subgraph cluster_15 {
+          color=blue
+          40 [label="Enter when branch condition "];
+          41 [label="Access variable this@R|/test_2|"];
+          42 [label="Type operator: this !is A"];
+          43 [label="Exit when branch condition"];
+        }
+        subgraph cluster_16 {
+          color=blue
+          44 [label="Enter block"];
+          45 [label="Access variable this@R|/test_2|"];
+          46 [label="Function call: this@R|/test_2|.<Unresolved name: foo>#()"];
+          47 [label="Function call: <Unresolved name: foo>#()"];
+          48 [label="Exit block"];
+        }
+        49 [label="Exit when branch result"];
+        subgraph cluster_17 {
+          color=blue
+          50 [label="Enter when branch condition else"];
+          51 [label="Exit when branch condition"];
+        }
+        subgraph cluster_18 {
+          color=blue
+          52 [label="Enter block"];
+          53 [label="Access variable this@R|/test_2|"];
+          54 [label="Function call: this@R|/test_2|.R|/A.foo|()"];
+          55 [label="Function call: this@R|/A|.R|/A.foo|()"];
+          56 [label="Exit block"];
+        }
+        57 [label="Exit when branch result"];
+        58 [label="Exit when"];
+      }
+      59 [label="Access variable this@R|/test_2|"];
+      60 [label="Function call: this@R|/test_2|.<Unresolved name: foo>#()"];
+      61 [label="Function call: <Unresolved name: foo>#()"];
+      62 [label="Exit block"];
+    }
+    63 [label="Exit function test_2" style="filled" fillcolor=red];
+  }
+
+  37 -> {38};
+  38 -> {39};
+  39 -> {40};
+  40 -> {41};
+  41 -> {42};
+  42 -> {43};
+  43 -> {44 50};
+  44 -> {45};
+  45 -> {46};
+  46 -> {47};
+  47 -> {48};
+  48 -> {49};
+  49 -> {58};
+  50 -> {51};
+  51 -> {52};
+  52 -> {53};
+  53 -> {54};
+  54 -> {55};
+  55 -> {56};
+  56 -> {57};
+  57 -> {58};
+  58 -> {59};
+  59 -> {60};
+  60 -> {61};
+  61 -> {62};
+  62 -> {63};
+
+  subgraph cluster_19 {
+    color=red
+    64 [label="Enter function test_3" style="filled" fillcolor=red];
+    subgraph cluster_20 {
+      color=blue
+      65 [label="Enter block"];
+      66 [label="Access variable R|<local>/a|"];
+      subgraph cluster_21 {
+        color=blue
+        67 [label="Enter function anonymousFunction"];
+        subgraph cluster_22 {
+          color=blue
+          68 [label="Enter block"];
+          69 [label="Access variable R|<local>/b|"];
+          subgraph cluster_23 {
+            color=blue
+            70 [label="Enter function anonymousFunction"];
+            subgraph cluster_24 {
+              color=blue
+              71 [label="Enter block"];
+              72 [label="Access variable R|<local>/c|"];
+              subgraph cluster_25 {
+                color=blue
+                73 [label="Enter function anonymousFunction"];
+                subgraph cluster_26 {
+                  color=blue
+                  74 [label="Enter block"];
+                  75 [label="Access variable this@R|special/anonymous|"];
+                  76 [label="Type operator: this@wb as A"];
+                  77 [label="Access variable this@R|special/anonymous|"];
+                  78 [label="Function call: this@R|special/anonymous|.R|/A.foo|()"];
+                  79 [label="Function call: this@R|/A|.R|/A.foo|()"];
+                  80 [label="Exit block"];
+                }
+                81 [label="Exit function anonymousFunction"];
+              }
+              82 [label="Function call: R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/c|, <L> = wc@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+    (this@R|special/anonymous| as R|A|)
+    this@R|special/anonymous|.R|/A.foo|()
+    this@R|/A|.R|/A.foo|()
+}
+)"];
+              83 [label="Access variable this@R|special/anonymous|"];
+              84 [label="Function call: this@R|special/anonymous|.R|/A.foo|()"];
+              85 [label="Function call: this@R|/A|.R|/A.foo|()"];
+              86 [label="Exit block"];
+            }
+            87 [label="Exit function anonymousFunction"];
+          }
+          88 [label="Function call: R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/b|, <L> = wb@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+    R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/c|, <L> = wc@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+        (this@R|special/anonymous| as R|A|)
+        this@R|special/anonymous|.R|/A.foo|()
+        this@R|/A|.R|/A.foo|()
+    }
+    )
+    this@R|special/anonymous|.R|/A.foo|()
+    this@R|/A|.R|/A.foo|()
+}
+)"];
+          89 [label="Exit block"];
+        }
+        90 [label="Exit function anonymousFunction"];
+      }
+      91 [label="Function call: R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/a|, <L> = wa@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+    R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/b|, <L> = wb@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+        R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/c|, <L> = wc@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+            (this@R|special/anonymous| as R|A|)
+            this@R|special/anonymous|.R|/A.foo|()
+            this@R|/A|.R|/A.foo|()
+        }
+        )
+        this@R|special/anonymous|.R|/A.foo|()
+        this@R|/A|.R|/A.foo|()
+    }
+    )
+}
+)"];
+      92 [label="Exit block"];
+    }
+    93 [label="Exit function test_3" style="filled" fillcolor=red];
+  }
+
+  64 -> {65};
+  65 -> {66};
+  66 -> {67};
+  67 -> {68};
+  68 -> {69};
+  69 -> {70};
+  70 -> {71};
+  71 -> {72};
+  72 -> {73};
+  73 -> {74};
+  74 -> {75};
+  75 -> {76};
+  76 -> {77};
+  77 -> {78};
+  78 -> {79};
+  79 -> {80};
+  80 -> {81};
+  81 -> {82};
+  82 -> {83};
+  83 -> {84};
+  84 -> {85};
+  85 -> {86};
+  86 -> {87};
+  87 -> {88};
+  88 -> {89};
+  89 -> {90};
+  90 -> {91};
+  91 -> {92};
+  92 -> {93};
+
+}

compiler/fir/resolve/testData/resolve/smartcasts/implicitReceivers.kt

@@ -0,0 +1,44 @@
+class A {
+    fun foo() {}
+}
+
+fun <T> T.with(block: T.() -> Unit) {}
+
+fun Any?.test_1() {
+    if (this is A) {
+        this.foo()
+        foo()
+    } else {
+        this.foo()
+        foo()
+    }
+    this.foo()
+    foo()
+}
+
+fun Any?.test_2() {
+    if (this !is A) {
+        this.foo()
+        foo()
+    } else {
+        this.foo()
+        foo()
+    }
+    this.foo()
+    foo()
+}
+
+
+fun test_3(a: Any, b: Any, c: Any) {
+    with(a) wa@{
+        with(b) wb@{
+            with(c) wc@{
+                this@wb as A
+                this@wb.foo()
+                foo()
+            }
+            this.foo()
+            foo()
+        }
+    }
+}

compiler/fir/resolve/testData/resolve/smartcasts/implicitReceivers.txt

@@ -0,0 +1,58 @@
+FILE: implicitReceivers.kt
+    public final class A : R|kotlin/Any| {
+        public constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final fun foo(): R|kotlin/Unit| {
+        }
+
+    }
+    public final fun <T> R|T|.with(block: R|kotlin/Function1<T, kotlin/Unit>|): R|kotlin/Unit| {
+    }
+    public final fun R|kotlin/Any?|.test_1(): R|kotlin/Unit| {
+        when () {
+            (this@R|/test_1| is R|A|) ->  {
+                this@R|/test_1|.R|/A.foo|()
+                this@R|/A|.R|/A.foo|()
+            }
+            else ->  {
+                this@R|/test_1|.<Unresolved name: foo>#()
+                <Unresolved name: foo>#()
+            }
+        }
+
+        this@R|/test_1|.<Unresolved name: foo>#()
+        <Unresolved name: foo>#()
+    }
+    public final fun R|kotlin/Any?|.test_2(): R|kotlin/Unit| {
+        when () {
+            (this@R|/test_2| !is R|A|) ->  {
+                this@R|/test_2|.<Unresolved name: foo>#()
+                <Unresolved name: foo>#()
+            }
+            else ->  {
+                this@R|/test_2|.R|/A.foo|()
+                this@R|/A|.R|/A.foo|()
+            }
+        }
+
+        this@R|/test_2|.<Unresolved name: foo>#()
+        <Unresolved name: foo>#()
+    }
+    public final fun test_3(a: R|kotlin/Any|, b: R|kotlin/Any|, c: R|kotlin/Any|): R|kotlin/Unit| {
+        R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/a|, <L> = wa@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+            R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/b|, <L> = wb@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+                R|kotlin/with|<R|kotlin/Any|, R|kotlin/Unit|>(R|<local>/c|, <L> = wc@fun R|kotlin/Any|.<anonymous>(it: R|kotlin/Any|): R|kotlin/Unit| <kind=EXACTLY_ONCE>  {
+                    (this@R|special/anonymous| as R|A|)
+                    this@R|special/anonymous|.R|/A.foo|()
+                    this@R|/A|.R|/A.foo|()
+                }
+                )
+                this@R|special/anonymous|.R|/A.foo|()
+                this@R|/A|.R|/A.foo|()
+            }
+            )
+        }
+        )
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirCfgBuildingTestGenerated.java

@@ -134,6 +134,11 @@ public class FirCfgBuildingTestGenerated extends AbstractFirCfgBuildingTest {
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/equalsToBoolean.kt");
         }
 
+        @TestMetadata("implicitReceivers.kt")
+        public void testImplicitReceivers() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/smartcasts/implicitReceivers.kt");
+        }
+
         @TestMetadata("inPlaceLambdas.kt")
         public void testInPlaceLambdas() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/inPlaceLambdas.kt");