[FIR] Fix determining type of when subject in exhaustiveness checker

#KT-37091 Fixed

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

@@ -42,12 +42,12 @@ class FirWhenExhaustivenessTransformer(private val bodyResolveComponents: BodyRe
             return whenExpression
         }
 
-        val typeRef = (whenExpression.subjectVariable?.returnTypeRef
+        val typeRef = whenExpression.subjectVariable?.returnTypeRef
-            ?: (whenExpression.subject as? FirQualifiedAccessExpression)?.typeRef) as? FirResolvedTypeRef
+            ?: whenExpression.subject?.typeRef
             ?: return null
 
         // TODO: add some report logic about flexible type (see WHEN_ENUM_CAN_BE_NULL_IN_JAVA diagnostic in old frontend)
-        val type = typeRef.type.lowerBoundIfFlexible()
+        val type = (typeRef as? FirResolvedTypeRef)?.type?.lowerBoundIfFlexible() ?: return null
         val lookupTag = (type as? ConeLookupTagBasedType)?.lookupTag ?: return null
         val nullable = typeRef.type.nullability == ConeNullability.NULLABLE
         val isExhaustive = when {

compiler/fir/resolve/testData/resolve/exhaustiveWhenAndDNNType.dot

@@ -0,0 +1,275 @@
+digraph exhaustiveWhenAndDNNType_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 values" style="filled" fillcolor=red];
+        3 [label="Exit function values" style="filled" fillcolor=red];
+    }
+
+    2 -> {3};
+
+    subgraph cluster_2 {
+        color=red
+        4 [label="Enter function valueOf" style="filled" fillcolor=red];
+        5 [label="Exit function valueOf" style="filled" fillcolor=red];
+    }
+
+    4 -> {5};
+
+    subgraph cluster_3 {
+        color=red
+        6 [label="Enter function <init>" style="filled" fillcolor=red];
+        7 [label="Exit function <init>" style="filled" fillcolor=red];
+    }
+
+    6 -> {7};
+
+    subgraph cluster_4 {
+        color=red
+        8 [label="Enter function takeB" style="filled" fillcolor=red];
+        9 [label="Exit function takeB" style="filled" fillcolor=red];
+    }
+
+    8 -> {9};
+
+    subgraph cluster_5 {
+        color=red
+        10 [label="Enter function test_1" style="filled" fillcolor=red];
+        11 [label="Access variable R|/SomeEnum.A1|"];
+        12 [label="Variable declaration: lval flag: R|SomeEnum|"];
+        subgraph cluster_6 {
+            color=blue
+            13 [label="Enter when"];
+            14 [label="Access variable R|<local>/flag|"];
+            15 [label="Check not null: R|<local>/flag|!!"];
+            subgraph cluster_7 {
+                color=blue
+                16 [label="Enter when branch condition "];
+                17 [label="Access variable R|/SomeEnum.A1|"];
+                18 [label="Operator =="];
+                19 [label="Exit when branch condition"];
+            }
+            subgraph cluster_8 {
+                color=blue
+                20 [label="Enter when branch condition "];
+                21 [label="Access variable R|/SomeEnum.A2|"];
+                22 [label="Operator =="];
+                23 [label="Exit when branch condition"];
+            }
+            24 [label="Enter when branch result"];
+            subgraph cluster_9 {
+                color=blue
+                25 [label="Enter block"];
+                26 [label="Function call: R|/B.B|()"];
+                27 [label="Exit block"];
+            }
+            28 [label="Exit when branch result"];
+            29 [label="Enter when branch result"];
+            subgraph cluster_10 {
+                color=blue
+                30 [label="Enter block"];
+                31 [label="Function call: R|/B.B|()"];
+                32 [label="Exit block"];
+            }
+            33 [label="Exit when branch result"];
+            34 [label="Exit when"];
+        }
+        35 [label="Variable declaration: lval b: R|B|"];
+        36 [label="Access variable R|<local>/b|"];
+        37 [label="Function call: R|/takeB|(R|<local>/b|)"];
+        38 [label="Exit function test_1" style="filled" fillcolor=red];
+    }
+
+    10 -> {11};
+    11 -> {12};
+    12 -> {13};
+    13 -> {14};
+    14 -> {15};
+    15 -> {16};
+    16 -> {17};
+    17 -> {18};
+    18 -> {19};
+    19 -> {29 20};
+    20 -> {21};
+    21 -> {22};
+    22 -> {23};
+    23 -> {24};
+    24 -> {25};
+    25 -> {26};
+    26 -> {27};
+    27 -> {28};
+    28 -> {34};
+    29 -> {30};
+    30 -> {31};
+    31 -> {32};
+    32 -> {33};
+    33 -> {34};
+    34 -> {35};
+    35 -> {36};
+    36 -> {37};
+    37 -> {38};
+
+    subgraph cluster_11 {
+        color=red
+        39 [label="Enter function test_2" style="filled" fillcolor=red];
+        40 [label="Access variable R|/SomeEnum.A1|"];
+        41 [label="Variable declaration: lval flag: R|SomeEnum|"];
+        subgraph cluster_12 {
+            color=blue
+            42 [label="Enter when"];
+            43 [label="Access variable R|<local>/flag|"];
+            44 [label="Check not null: R|<local>/flag|!!"];
+            subgraph cluster_13 {
+                color=blue
+                45 [label="Enter when branch condition "];
+                46 [label="Access variable R|/SomeEnum.A1|"];
+                47 [label="Operator =="];
+                48 [label="Exit when branch condition"];
+            }
+            subgraph cluster_14 {
+                color=blue
+                49 [label="Enter when branch condition "];
+                50 [label="Access variable R|/SomeEnum.A2|"];
+                51 [label="Operator =="];
+                52 [label="Exit when branch condition"];
+            }
+            53 [label="Enter when branch result"];
+            subgraph cluster_15 {
+                color=blue
+                54 [label="Enter block"];
+                55 [label="Function call: R|/B.B|()"];
+                56 [label="Exit block"];
+            }
+            57 [label="Exit when branch result"];
+            58 [label="Enter when branch result"];
+            subgraph cluster_16 {
+                color=blue
+                59 [label="Enter block"];
+                60 [label="Function call: R|/B.B|()"];
+                61 [label="Exit block"];
+            }
+            62 [label="Exit when branch result"];
+            63 [label="Exit when"];
+        }
+        64 [label="Variable declaration: lval b: R|B|"];
+        65 [label="Access variable R|<local>/b|"];
+        66 [label="Function call: R|/takeB|(R|<local>/b|)"];
+        67 [label="Exit function test_2" style="filled" fillcolor=red];
+    }
+
+    39 -> {40};
+    40 -> {41};
+    41 -> {42};
+    42 -> {43};
+    43 -> {44};
+    44 -> {45};
+    45 -> {46};
+    46 -> {47};
+    47 -> {48};
+    48 -> {58 49};
+    49 -> {50};
+    50 -> {51};
+    51 -> {52};
+    52 -> {53};
+    53 -> {54};
+    54 -> {55};
+    55 -> {56};
+    56 -> {57};
+    57 -> {63};
+    58 -> {59};
+    59 -> {60};
+    60 -> {61};
+    61 -> {62};
+    62 -> {63};
+    63 -> {64};
+    64 -> {65};
+    65 -> {66};
+    66 -> {67};
+
+    subgraph cluster_17 {
+        color=red
+        68 [label="Enter function test_3" style="filled" fillcolor=red];
+        69 [label="Access variable R|/SomeEnum.A1|"];
+        70 [label="Variable declaration: lval flag: R|SomeEnum|"];
+        subgraph cluster_18 {
+            color=blue
+            71 [label="Enter when"];
+            72 [label="Access variable R|<local>/flag|"];
+            subgraph cluster_19 {
+                color=blue
+                73 [label="Enter when branch condition "];
+                74 [label="Access variable R|/SomeEnum.A1|"];
+                75 [label="Operator =="];
+                76 [label="Exit when branch condition"];
+            }
+            subgraph cluster_20 {
+                color=blue
+                77 [label="Enter when branch condition "];
+                78 [label="Access variable R|/SomeEnum.A2|"];
+                79 [label="Operator =="];
+                80 [label="Exit when branch condition"];
+            }
+            81 [label="Enter when branch result"];
+            subgraph cluster_21 {
+                color=blue
+                82 [label="Enter block"];
+                83 [label="Function call: R|/B.B|()"];
+                84 [label="Exit block"];
+            }
+            85 [label="Exit when branch result"];
+            86 [label="Enter when branch result"];
+            subgraph cluster_22 {
+                color=blue
+                87 [label="Enter block"];
+                88 [label="Function call: R|/B.B|()"];
+                89 [label="Exit block"];
+            }
+            90 [label="Exit when branch result"];
+            91 [label="Exit when"];
+        }
+        92 [label="Variable declaration: lval b: R|B|"];
+        93 [label="Access variable R|<local>/b|"];
+        94 [label="Function call: R|/takeB|(R|<local>/b|)"];
+        95 [label="Exit function test_3" style="filled" fillcolor=red];
+    }
+
+    68 -> {69};
+    69 -> {70};
+    70 -> {71};
+    71 -> {72};
+    72 -> {73};
+    73 -> {74};
+    74 -> {75};
+    75 -> {76};
+    76 -> {86 77};
+    77 -> {78};
+    78 -> {79};
+    79 -> {80};
+    80 -> {81};
+    81 -> {82};
+    82 -> {83};
+    83 -> {84};
+    84 -> {85};
+    85 -> {91};
+    86 -> {87};
+    87 -> {88};
+    88 -> {89};
+    89 -> {90};
+    90 -> {91};
+    91 -> {92};
+    92 -> {93};
+    93 -> {94};
+    94 -> {95};
+
+}

compiler/fir/resolve/testData/resolve/exhaustiveWhenAndDNNType.kt

@@ -0,0 +1,38 @@
+// !DUMP_CFG
+// ISSUE: KT-37091
+
+enum class SomeEnum {
+    A1, A2;
+}
+
+class B()
+
+fun takeB(b: B) {}
+
+fun test_1() {
+    val flag = SomeEnum.A1
+    val b: B = when (flag!!) {
+        SomeEnum.A1 -> B()
+        SomeEnum.A2 -> B()
+    }
+    takeB(b) // should be OK
+}
+
+fun test_2() {
+    val flag = SomeEnum.A1
+
+    val b = when (flag!!) {
+        SomeEnum.A1 -> B()
+        SomeEnum.A2 -> B()
+    }
+    takeB(b) // should be OK
+}
+
+fun test_3() {
+    val flag = SomeEnum.A1
+    val b = when (flag) { //there is no null-assertion! , no explicit type
+        SomeEnum.A1 -> B()
+        SomeEnum.A2 -> B()
+    }
+    takeB(b) // should be OK
+}

compiler/fir/resolve/testData/resolve/exhaustiveWhenAndDNNType.txt

@@ -0,0 +1,62 @@
+FILE: exhaustiveWhenAndDNNType.kt
+    public final enum class SomeEnum : R|kotlin/Enum<SomeEnum>| {
+        private constructor(): R|SomeEnum| {
+            super<R|kotlin/Enum<SomeEnum>|>()
+        }
+
+        public final static enum entry A1: R|SomeEnum|
+        public final static enum entry A2: R|SomeEnum|
+        public final static fun values(): R|kotlin/Array<SomeEnum>| {
+        }
+
+        public final static fun valueOf(value: R|kotlin/String|): R|SomeEnum| {
+        }
+
+    }
+    public final class B : R|kotlin/Any| {
+        public constructor(): R|B| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final fun takeB(b: R|B|): R|kotlin/Unit| {
+    }
+    public final fun test_1(): R|kotlin/Unit| {
+        lval flag: R|SomeEnum| = Q|SomeEnum|.R|/SomeEnum.A1|
+        lval b: R|B| = when (R|<local>/flag|!!) {
+            ==($subj$, Q|SomeEnum|.R|/SomeEnum.A1|) ->  {
+                R|/B.B|()
+            }
+            ==($subj$, Q|SomeEnum|.R|/SomeEnum.A2|) ->  {
+                R|/B.B|()
+            }
+        }
+
+        R|/takeB|(R|<local>/b|)
+    }
+    public final fun test_2(): R|kotlin/Unit| {
+        lval flag: R|SomeEnum| = Q|SomeEnum|.R|/SomeEnum.A1|
+        lval b: R|B| = when (R|<local>/flag|!!) {
+            ==($subj$, Q|SomeEnum|.R|/SomeEnum.A1|) ->  {
+                R|/B.B|()
+            }
+            ==($subj$, Q|SomeEnum|.R|/SomeEnum.A2|) ->  {
+                R|/B.B|()
+            }
+        }
+
+        R|/takeB|(R|<local>/b|)
+    }
+    public final fun test_3(): R|kotlin/Unit| {
+        lval flag: R|SomeEnum| = Q|SomeEnum|.R|/SomeEnum.A1|
+        lval b: R|B| = when (R|<local>/flag|) {
+            ==($subj$, Q|SomeEnum|.R|/SomeEnum.A1|) ->  {
+                R|/B.B|()
+            }
+            ==($subj$, Q|SomeEnum|.R|/SomeEnum.A2|) ->  {
+                R|/B.B|()
+            }
+        }
+
+        R|/takeB|(R|<local>/b|)
+    }

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

@@ -108,6 +108,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         runTest("compiler/fir/resolve/testData/resolve/enumWithCompanion.kt");
     }
 
+    @TestMetadata("exhaustiveWhenAndDNNType.kt")
+    public void testExhaustiveWhenAndDNNType() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/exhaustiveWhenAndDNNType.kt");
+    }
+
     @TestMetadata("exhaustiveWhenAndFlexibleType.kt")
     public void testExhaustiveWhenAndFlexibleType() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/exhaustiveWhenAndFlexibleType.kt");

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

@@ -108,6 +108,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         runTest("compiler/fir/resolve/testData/resolve/enumWithCompanion.kt");
     }
 
+    @TestMetadata("exhaustiveWhenAndDNNType.kt")
+    public void testExhaustiveWhenAndDNNType() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/exhaustiveWhenAndDNNType.kt");
+    }
+
     @TestMetadata("exhaustiveWhenAndFlexibleType.kt")
     public void testExhaustiveWhenAndFlexibleType() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/exhaustiveWhenAndFlexibleType.kt");