[FIR] Properly handle when subject expression in DFA of equality calls

^KT-50785 Fixed

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerFirTestdataTestGenerated.java

@@ -3539,6 +3539,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50788.kt");
             }
 
+            @Test
+            @TestMetadata("literalInWhenSubject.kt")
+            public void testLiteralInWhenSubject() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/literalInWhenSubject.kt");
+            }
+
             @Test
             @TestMetadata("multipleCasts.kt")
             public void testMultipleCasts() throws Exception {

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -3129,6 +3129,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50788.kt");
         }
 
+        @TestMetadata("literalInWhenSubject.kt")
+        public void testLiteralInWhenSubject() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/literalInWhenSubject.kt");
+        }
+
         @TestMetadata("multipleCasts.kt")
         public void testMultipleCasts() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/multipleCasts.kt");

compiler/fir/analysis-tests/testData/resolve/smartcasts/literalInWhenSubject.fir.txt

@@ -0,0 +1,31 @@
+FILE: literalInWhenSubject.kt
+    public final fun test_1(s: R|kotlin/String?|): R|kotlin/Unit| {
+        when (Boolean(true)) {
+            ==($subj$, !=(R|<local>/s|, Null(null))) ->  {
+                R|<local>/s|.R|kotlin/String.length|
+            }
+            else ->  {
+                Null(null)
+            }
+        }
+
+    }
+    public final fun test_2(s: R|kotlin/String?|): R|kotlin/Unit| {
+        when (!=(R|<local>/s|, Null(null))) {
+            ==($subj$, Boolean(true)) ->  {
+                R|<local>/s|.R|kotlin/String.length|
+            }
+            else ->  {
+                Null(null)
+            }
+        }
+
+    }
+    public final fun test_3(s: R|kotlin/String?|): R|kotlin/Unit| {
+        when () {
+            ==(Boolean(true), !=(R|<local>/s|, Null(null))) ->  {
+                R|<local>/s|.R|kotlin/String.length|
+            }
+        }
+
+    }

compiler/fir/analysis-tests/testData/resolve/smartcasts/literalInWhenSubject.kt

@@ -0,0 +1,19 @@
+// ISSUE: KT-50785
+
+fun test_1(s: String?) {
+    when (true) {
+        (s != null) -> s.length
+        else -> null
+    }
+}
+
+fun test_2(s: String?) {
+    when (s != null) {
+        true -> s.length
+        else -> null
+    }
+}
+
+fun test_3(s: String?) {
+    if (true == (s != null)) s.length
+}

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java

@@ -3539,6 +3539,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50788.kt");
             }
 
+            @Test
+            @TestMetadata("literalInWhenSubject.kt")
+            public void testLiteralInWhenSubject() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/literalInWhenSubject.kt");
+            }
+
             @Test
             @TestMetadata("multipleCasts.kt")
             public void testMultipleCasts() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -3539,6 +3539,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50788.kt");
             }
 
+            @Test
+            @TestMetadata("literalInWhenSubject.kt")
+            public void testLiteralInWhenSubject() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/literalInWhenSubject.kt");
+            }
+
             @Test
             @TestMetadata("multipleCasts.kt")
             public void testMultipleCasts() throws Exception {

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

@@ -478,7 +478,20 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val leftOperand = equalityOperatorCall.arguments[0]
         val rightOperand = equalityOperatorCall.arguments[1]
 
-        val leftConst = leftOperand as? FirConstExpression<*>
+        /*
+         * This unwrapping is needed for cases like
+         * when (true) {
+         *     s != null -> s.length
+         * }
+         *
+         * FirWhenSubjectExpression may be only in the lhs of equality operator call
+         *   by how is FIR for when branches is built, so there is no need to unwrap
+         *   right argument
+         */
+        val leftConst = when (leftOperand) {
+            is FirWhenSubjectExpression -> leftOperand.whenRef.value.subject
+            else -> leftOperand
+        } as? FirConstExpression<*>
         val rightConst = rightOperand as? FirConstExpression<*>
         val leftIsNullConst = leftConst?.kind == ConstantValueKind.Null
         val rightIsNullConst = rightConst?.kind == ConstantValueKind.Null