[FIR LT] Add proper source for when with subject

#KT-56982 Fixed

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirLightTreeSteppingTestGenerated.java

@@ -535,6 +535,12 @@ public class FirLightTreeSteppingTestGenerated extends AbstractFirLightTreeStepp
         runTest("compiler/testData/debug/stepping/whenConstant.kt");
     }
 
+    @Test
+    @TestMetadata("whenExpr.kt")
+    public void testWhenExpr() throws Exception {
+        runTest("compiler/testData/debug/stepping/whenExpr.kt");
+    }
+
     @Test
     @TestMetadata("whenIsChecks.kt")
     public void testWhenIsChecks() throws Exception {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirPsiSteppingTestGenerated.java

@@ -535,6 +535,12 @@ public class FirPsiSteppingTestGenerated extends AbstractFirPsiSteppingTest {
         runTest("compiler/testData/debug/stepping/whenConstant.kt");
     }
 
+    @Test
+    @TestMetadata("whenExpr.kt")
+    public void testWhenExpr() throws Exception {
+        runTest("compiler/testData/debug/stepping/whenExpr.kt");
+    }
+
     @Test
     @TestMetadata("whenIsChecks.kt")
     public void testWhenIsChecks() throws Exception {

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -835,23 +835,25 @@ class ExpressionsConverter(
         }
 
         val calculatedFirExpression = firExpression ?: buildErrorExpression(
-            null,
+            source = null,
             ConeSimpleDiagnostic("No expression in condition with expression", DiagnosticKind.Syntax)
         )
 
-        return if (whenRefWithSubject != null) {
+        if (whenRefWithSubject == null) {
-            buildEqualityOperatorCall {
+            return calculatedFirExpression
+        }
+
+        return buildEqualityOperatorCall {
             source = whenCondition.toFirSourceElement(KtFakeSourceElementKind.WhenCondition)
             operation = FirOperation.EQ
             argumentList = buildBinaryArgumentList(
-                    buildWhenSubjectExpression {
+                left = buildWhenSubjectExpression {
+                    source = whenCondition.toFirSourceElement()
                     whenRef = whenRefWithSubject
-                    }, calculatedFirExpression
+                },
+                right = calculatedFirExpression
             )
         }
-        } else {
-            calculatedFirExpression
-        }
     }
 
     private fun convertWhenConditionInRange(
@@ -918,6 +920,7 @@ class ExpressionsConverter(
 
         val subjectExpression = if (whenRefWithSubject != null) {
             buildWhenSubjectExpression {
+                source = whenCondition.toFirSourceElement()
                 whenRef = whenRefWithSubject
             }
         } else {

compiler/testData/debug/stepping/whenExpr.kt

@@ -0,0 +1,30 @@
+// FILE: test.kt
+
+fun box() {
+    val a: Int? = 3
+    when (a) {
+        1 -> {
+            1
+        }
+        2 -> {
+            2
+        }
+        else -> {
+            0
+        }
+    }
+}
+
+// EXPECTATIONS JVM JVM_IR
+// test.kt:4 box
+// test.kt:5 box
+// test.kt:6 box
+// test.kt:9 box
+// test.kt:13 box
+// test.kt:16 box
+
+// EXPECTATIONS JS_IR
+// test.kt:4 box
+// test.kt:5 box
+// test.kt:5 box
+// test.kt:16 box

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrSteppingWithBytecodeInlinerTestGenerated.java

@@ -535,6 +535,12 @@ public class IrSteppingWithBytecodeInlinerTestGenerated extends AbstractIrSteppi
         runTest("compiler/testData/debug/stepping/whenConstant.kt");
     }
 
+    @Test
+    @TestMetadata("whenExpr.kt")
+    public void testWhenExpr() throws Exception {
+        runTest("compiler/testData/debug/stepping/whenExpr.kt");
+    }
+
     @Test
     @TestMetadata("whenIsChecks.kt")
     public void testWhenIsChecks() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrSteppingWithIrInlinerTestGenerated.java

@@ -535,6 +535,12 @@ public class IrSteppingWithIrInlinerTestGenerated extends AbstractIrSteppingWith
         runTest("compiler/testData/debug/stepping/whenConstant.kt");
     }
 
+    @Test
+    @TestMetadata("whenExpr.kt")
+    public void testWhenExpr() throws Exception {
+        runTest("compiler/testData/debug/stepping/whenExpr.kt");
+    }
+
     @Test
     @TestMetadata("whenIsChecks.kt")
     public void testWhenIsChecks() throws Exception {

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/SteppingTestGenerated.java

@@ -535,6 +535,12 @@ public class SteppingTestGenerated extends AbstractSteppingTest {
         runTest("compiler/testData/debug/stepping/whenConstant.kt");
     }
 
+    @Test
+    @TestMetadata("whenExpr.kt")
+    public void testWhenExpr() throws Exception {
+        runTest("compiler/testData/debug/stepping/whenExpr.kt");
+    }
+
     @Test
     @TestMetadata("whenIsChecks.kt")
     public void testWhenIsChecks() throws Exception {

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/IrJsSteppingTestGenerated.java

@@ -535,6 +535,12 @@ public class IrJsSteppingTestGenerated extends AbstractIrJsSteppingTest {
         runTest("compiler/testData/debug/stepping/whenConstant.kt");
     }
 
+    @Test
+    @TestMetadata("whenExpr.kt")
+    public void testWhenExpr() throws Exception {
+        runTest("compiler/testData/debug/stepping/whenExpr.kt");
+    }
+
     @Test
     @TestMetadata("whenIsChecks.kt")
     public void testWhenIsChecks() throws Exception {