[FIR] Fix smartcasts in when expressions with this subject

Also get rid of excess `getOrCreateRealVariable` function.

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

@@ -13,6 +13,7 @@ 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.references.FirResolvedCallableReference
+import org.jetbrains.kotlin.fir.references.impl.FirExplicitThisReference
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
 import org.jetbrains.kotlin.fir.resolve.ImplicitReceiverStackImpl
 import org.jetbrains.kotlin.fir.resolve.dfa.Condition.*
@@ -21,8 +22,6 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirBodyResolveTransformer
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.CallableId
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirVariableSymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.fir.types.coneTypeUnsafe
@@ -30,6 +29,7 @@ import org.jetbrains.kotlin.fir.types.isMarkedNullable
 import org.jetbrains.kotlin.ir.expressions.IrConstKind
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addIfNotNull
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveComponents by transformer {
@@ -458,15 +458,21 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     private val FirElement.resolvedSymbol: AbstractFirBasedSymbol<*>?
         get() {
-            val expression = (this as? FirWhenSubjectExpression)?.whenSubject?.whenExpression?.let {
+            return when (this) {
-                it.subjectVariable?.symbol?.let { symbol -> return symbol }
+                is FirResolvable -> resolvedSymbol
-                it.subject
+                is FirSymbolOwner<*> -> symbol
-            } ?: this
+                else -> null
-            return (expression as? FirResolvable)?.resolvedSymbol
+            }
         }
 
     private val FirResolvable.resolvedSymbol: AbstractFirBasedSymbol<*>?
-        get() = (calleeReference as? FirResolvedCallableReference)?.resolvedSymbol
+        get() = calleeReference.let {
+            if (it is FirExplicitThisReference) {
+                it.boundSymbol
+            } else {
+                (it as? FirResolvedCallableReference)?.resolvedSymbol
+            }
+        }
 
     private fun FirFunctionCall.isBooleanNot(): Boolean {
         val symbol = calleeReference.safeAs<FirResolvedCallableReference>()?.resolvedSymbol as? FirNamedFunctionSymbol ?: return false
@@ -492,12 +498,13 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
          */
         variableStorage[initializer]?.let { initializerVariable ->
             assert(initializerVariable.isSynthetic())
-            val realVariable = getOrCreateRealVariable(variable.symbol)
+            val realVariable = getOrCreateRealVariable(variable)
+            requireNotNull(realVariable)
             logicSystem.changeVariableForConditionFlow(node.flow, initializerVariable, realVariable)
         }
 
-        initializer.resolvedSymbol?.let { initializerSymbol: AbstractFirBasedSymbol<*> ->
+
-            val rhsVariable = getOrCreateRealVariable(initializerSymbol)
+        getOrCreateRealVariable(initializer)?.let { rhsVariable ->
             variableStorage.createAliasVariable(variable.symbol, rhsVariable)
         }
     }
@@ -658,30 +665,34 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     private fun getOrCreateSyntheticVariable(fir: FirElement): SyntheticDataFlowVariable =
         variableStorage.getOrCreateNewSyntheticVariable(fir)
 
+    private fun FirElement.unwrapWhenSubjectExpression(): FirElement = if (this is FirWhenSubjectExpression) {
+        val whenExpression = whenSubject.whenExpression
+        whenExpression.subjectVariable
+            ?: whenExpression.subject
+            ?: throw IllegalStateException("Subject or subject variable must be not null")
+    } else {
+        this
+    }
+
     private fun getOrCreateRealVariable(fir: FirElement): RealDataFlowVariable? {
+        @Suppress("NAME_SHADOWING")
+        val fir = fir.unwrapWhenSubjectExpression()
         if (fir is FirThisReceiverExpressionImpl) {
             return variableStorage.getOrCreateNewThisRealVariable(fir.calleeReference.boundSymbol ?: return null)
         }
-        val symbol = fir.resolvedSymbol ?: return null
+        val symbol: AbstractFirBasedSymbol<*> = fir.resolvedSymbol ?: return null
-        return variableStorage.getOrCreateNewRealVariable(symbol)
+        return variableStorage.getOrCreateNewRealVariable(symbol).variableUnderAlias
     }
 
-    private fun getOrCreateRealVariable(symbol: AbstractFirBasedSymbol<*>): RealDataFlowVariable =
-        variableStorage.getOrCreateNewRealVariable(symbol).variableUnderAlias
-
     private fun getOrCreateVariable(fir: FirElement): DataFlowVariable {
-        val symbol = fir.resolvedSymbol
+        return getOrCreateRealVariable(fir) ?: getOrCreateSyntheticVariable(fir)
-        return if (symbol == null)
-            getOrCreateSyntheticVariable(fir)
-        else
-            getOrCreateRealVariable(symbol)
     }
 
     // -------------------------------- get variable --------------------------------
 
     private val FirElement.realVariable: RealDataFlowVariable?
         get() {
-            val symbol = if (this is FirThisReceiverExpressionImpl) {
+            val symbol: AbstractFirBasedSymbol<*> = if (this is FirThisReceiverExpressionImpl) {
                 calleeReference.boundSymbol
             } else {
                 resolvedSymbol
@@ -700,15 +711,11 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
                         require(explicitReceiver != null)
                         collect(explicitReceiver)
                     }
-                    ((call.calleeReference as? FirResolvedCallableReference)?.resolvedSymbol)?.let { symbol ->
+                    result.addIfNotNull(getOrCreateRealVariable(call))
-                        if (symbol is FirVariableSymbol<*> || symbol is FirPropertySymbol) {
-                            result += getOrCreateRealVariable(symbol)
-                        }
-                    }
                 }
                 is FirWhenSubjectExpression -> {
                     // TODO: check
-                    call.whenSubject.whenExpression.subjectVariable?.let { result += getOrCreateRealVariable(it.symbol) }
+                    call.whenSubject.whenExpression.subjectVariable?.let { result += getOrCreateRealVariable(it)!! }
                     call.whenSubject.whenExpression.subject?.let { collect(it) }
                 }
             }

compiler/fir/resolve/testData/resolve/recursiveCallOnWhenWithSealedClass.txt

@@ -27,13 +27,13 @@ FILE: recursiveCallOnWhenWithSealedClass.kt
 
         }
 
-        public final fun unwrap(): <ERROR TYPE REF: > {
+        public final fun unwrap(): R|kotlin/Any?| {
             ^unwrap when (this@R|/Maybe|) {
                 ($subj$ is R|Maybe.Nope|) ->  {
                     throw R|java/lang/Exception.Exception|(String())
                 }
                 ($subj$ is R|Maybe.Yeah|) ->  {
-                    <Unresolved name: meat>#
+                    this@R|/Maybe.Yeah|.R|/Maybe.Yeah.meat|
                 }
                 else ->  {
                 }

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

@@ -0,0 +1,230 @@
+digraph implicitReceiverAsWhenSubject_kt {
+  graph [splines=ortho nodesep=3]
+  node [shape=box penwidth=2]
+  edge [penwidth=2]
+
+  subgraph cluster_0 {
+    color=red
+    0 [label="Enter function test_1" style="filled" fillcolor=red];
+    subgraph cluster_1 {
+      color=blue
+      1 [label="Enter block"];
+      subgraph cluster_2 {
+        color=blue
+        2 [label="Enter when"];
+        3 [label="Access variable this@R|/test_1|"];
+        subgraph cluster_3 {
+          color=blue
+          4 [label="Enter when branch condition "];
+          5 [label="Type operator: List<*>"];
+          6 [label="Exit when branch condition"];
+        }
+        subgraph cluster_4 {
+          color=blue
+          7 [label="Enter when branch condition "];
+          8 [label="Type operator: String"];
+          9 [label="Exit when branch condition"];
+        }
+        subgraph cluster_5 {
+          color=blue
+          10 [label="Enter when branch condition else"];
+          11 [label="Exit when branch condition"];
+        }
+        12 [label="Enter when branch result"];
+        subgraph cluster_6 {
+          color=blue
+          13 [label="Enter block"];
+          14 [label="Const: Int(0)"];
+          15 [label="Exit block"];
+        }
+        16 [label="Exit when branch result"];
+        17 [label="Enter when branch result"];
+        subgraph cluster_7 {
+          color=blue
+          18 [label="Enter block"];
+          19 [label="Access variable R|kotlin/String.length|"];
+          20 [label="Exit block"];
+        }
+        21 [label="Exit when branch result"];
+        22 [label="Enter when branch result"];
+        subgraph cluster_8 {
+          color=blue
+          23 [label="Enter block"];
+          24 [label="Access variable this@R|/test_1|"];
+          25 [label="Access variable R|kotlin/collections/List.size|"];
+          26 [label="Exit block"];
+        }
+        27 [label="Exit when branch result"];
+        28 [label="Exit when"];
+      }
+      29 [label="Jump: ^test_1 when (this@R|/test_1|) {
+    ($subj$ is R|kotlin/collections/List<*>|) ->  {
+        this@R|/test_1|.R|kotlin/collections/List.size|
+    }
+    ($subj$ is R|kotlin/String|) ->  {
+        this@R|kotlin/String|.R|kotlin/String.length|
+    }
+    else ->  {
+        Int(0)
+    }
+}
+"];
+      30 [label="Stub" style="filled" fillcolor=gray];
+      31 [label="Exit block" style="filled" fillcolor=gray];
+    }
+    32 [label="Exit function test_1" style="filled" fillcolor=red];
+  }
+
+  0 -> {1};
+  1 -> {2};
+  2 -> {3};
+  3 -> {4};
+  4 -> {5};
+  5 -> {6};
+  6 -> {22 7};
+  7 -> {8};
+  8 -> {9};
+  9 -> {17 10};
+  10 -> {11};
+  11 -> {12};
+  12 -> {13};
+  13 -> {14};
+  14 -> {15};
+  15 -> {16};
+  16 -> {28};
+  17 -> {18};
+  18 -> {19};
+  19 -> {20};
+  20 -> {21};
+  21 -> {28};
+  22 -> {23};
+  23 -> {24};
+  24 -> {25};
+  25 -> {26};
+  26 -> {27};
+  27 -> {28};
+  28 -> {29};
+  29 -> {32};
+  29 -> {30} [style=dotted];
+  30 -> {31} [style=dotted];
+  31 -> {32} [style=dotted];
+
+  subgraph cluster_9 {
+    color=red
+    33 [label="Enter function test_2" style="filled" fillcolor=red];
+    subgraph cluster_10 {
+      color=blue
+      34 [label="Enter block"];
+      subgraph cluster_11 {
+        color=blue
+        35 [label="Enter when"];
+        36 [label="Access variable this@R|/test_2|"];
+        37 [label="Variable declaration: lval x: R|kotlin/Any|"];
+        subgraph cluster_12 {
+          color=blue
+          38 [label="Enter when branch condition "];
+          39 [label="Type operator: List<*>"];
+          40 [label="Exit when branch condition"];
+        }
+        subgraph cluster_13 {
+          color=blue
+          41 [label="Enter when branch condition "];
+          42 [label="Type operator: String"];
+          43 [label="Exit when branch condition"];
+        }
+        subgraph cluster_14 {
+          color=blue
+          44 [label="Enter when branch condition else"];
+          45 [label="Exit when branch condition"];
+        }
+        46 [label="Enter when branch result"];
+        subgraph cluster_15 {
+          color=blue
+          47 [label="Enter block"];
+          48 [label="Const: Int(0)"];
+          49 [label="Exit block"];
+        }
+        50 [label="Exit when branch result"];
+        51 [label="Enter when branch result"];
+        subgraph cluster_16 {
+          color=blue
+          52 [label="Enter block"];
+          53 [label="Access variable R|<local>/x|"];
+          54 [label="Access variable R|kotlin/String.length|"];
+          55 [label="Access variable R|kotlin/String.length|"];
+          56 [label="Exit block"];
+        }
+        57 [label="Exit when branch result"];
+        58 [label="Enter when branch result"];
+        subgraph cluster_17 {
+          color=blue
+          59 [label="Enter block"];
+          60 [label="Access variable R|<local>/x|"];
+          61 [label="Access variable R|kotlin/collections/List.size|"];
+          62 [label="Access variable this@R|/test_2|"];
+          63 [label="Access variable R|kotlin/collections/List.size|"];
+          64 [label="Exit block"];
+        }
+        65 [label="Exit when branch result"];
+        66 [label="Exit when"];
+      }
+      67 [label="Jump: ^test_2 when (lval x: R|kotlin/Any| = this@R|/test_2|) {
+    ($subj$ is R|kotlin/collections/List<*>|) ->  {
+        R|<local>/x|.R|kotlin/collections/List.size|
+        this@R|/test_2|.R|kotlin/collections/List.size|
+    }
+    ($subj$ is R|kotlin/String|) ->  {
+        R|<local>/x|.R|kotlin/String.length|
+        this@R|kotlin/String|.R|kotlin/String.length|
+    }
+    else ->  {
+        Int(0)
+    }
+}
+"];
+      68 [label="Stub" style="filled" fillcolor=gray];
+      69 [label="Exit block" style="filled" fillcolor=gray];
+    }
+    70 [label="Exit function test_2" style="filled" fillcolor=red];
+  }
+
+  33 -> {34};
+  34 -> {35};
+  35 -> {36};
+  36 -> {37};
+  37 -> {38};
+  38 -> {39};
+  39 -> {40};
+  40 -> {58 41};
+  41 -> {42};
+  42 -> {43};
+  43 -> {51 44};
+  44 -> {45};
+  45 -> {46};
+  46 -> {47};
+  47 -> {48};
+  48 -> {49};
+  49 -> {50};
+  50 -> {66};
+  51 -> {52};
+  52 -> {53};
+  53 -> {54};
+  54 -> {55};
+  55 -> {56};
+  56 -> {57};
+  57 -> {66};
+  58 -> {59};
+  59 -> {60};
+  60 -> {61};
+  61 -> {62};
+  62 -> {63};
+  63 -> {64};
+  64 -> {65};
+  65 -> {66};
+  66 -> {67};
+  67 -> {70};
+  67 -> {68} [style=dotted];
+  68 -> {69} [style=dotted];
+  69 -> {70} [style=dotted];
+
+}

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

@@ -0,0 +1,17 @@
+fun Any.test_1(): Int = when (this) {
+    is List<*> -> this.size
+    is String -> length
+    else -> 0
+}
+
+fun Any.test_2(): Int = when (val x = this) {
+    is List<*> -> {
+        x.size
+        this.size
+    }
+    is String -> {
+        x.length
+        length
+    }
+    else -> 0
+}

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

@@ -0,0 +1,31 @@
+FILE: implicitReceiverAsWhenSubject.kt
+    public final fun R|kotlin/Any|.test_1(): R|kotlin/Int| {
+        ^test_1 when (this@R|/test_1|) {
+            ($subj$ is R|kotlin/collections/List<*>|) ->  {
+                this@R|/test_1|.R|kotlin/collections/List.size|
+            }
+            ($subj$ is R|kotlin/String|) ->  {
+                this@R|kotlin/String|.R|kotlin/String.length|
+            }
+            else ->  {
+                Int(0)
+            }
+        }
+
+    }
+    public final fun R|kotlin/Any|.test_2(): R|kotlin/Int| {
+        ^test_2 when (lval x: R|kotlin/Any| = this@R|/test_2|) {
+            ($subj$ is R|kotlin/collections/List<*>|) ->  {
+                R|<local>/x|.R|kotlin/collections/List.size|
+                this@R|/test_2|.R|kotlin/collections/List.size|
+            }
+            ($subj$ is R|kotlin/String|) ->  {
+                R|<local>/x|.R|kotlin/String.length|
+                this@R|kotlin/String|.R|kotlin/String.length|
+            }
+            else ->  {
+                Int(0)
+            }
+        }
+
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/AbstractFirCfgBuildingTest.kt

@@ -40,8 +40,8 @@ abstract class AbstractFirCfgBuildingTest : AbstractFirResolveTestCase() {
 
     override fun doTest(path: String) {
         val firFiles = processInputFile(path)
-        checkCfg(path, firFiles)
         checkFir(path, firFiles)
+        checkCfg(path, firFiles)
     }
 
     fun checkCfg(path: String, firFiles: List<FirFile>) {

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

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