[FIR] Fix transform of subject of when expressions with subject variable

compiler/fir/resolve/testData/resolve/invokeInWhenSubjectVariableInitializer.kt

@@ -0,0 +1,3 @@
+fun test(func: () -> String?) {
+    val x = func() ?: ""
+}

compiler/fir/resolve/testData/resolve/invokeInWhenSubjectVariableInitializer.txt

@@ -0,0 +1,12 @@
+FILE: invokeInWhenSubjectVariableInitializer.kt
+    public final fun test(func: R|() -> kotlin/String?|): R|kotlin/Unit| {
+        lval x: R|kotlin/String| = when (lval <elvis>: R|kotlin/String?| = R|<local>/func|.R|FakeOverride<kotlin/Function0.invoke: R|kotlin/String?|>|()) {
+            ==($subj$, Null(null)) ->  {
+                String()
+            }
+            else ->  {
+                R|<local>/<elvis>|
+            }
+        }
+
+    }

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

@@ -203,6 +203,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         runTest("compiler/fir/resolve/testData/resolve/intersectionTypes.kt");
     }
 
+    @TestMetadata("invokeInWhenSubjectVariableInitializer.kt")
+    public void testInvokeInWhenSubjectVariableInitializer() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/invokeInWhenSubjectVariableInitializer.kt");
+    }
+
     @TestMetadata("invokeOfLambdaWithReceiver.kt")
     public void testInvokeOfLambdaWithReceiver() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/invokeOfLambdaWithReceiver.kt");

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

@@ -203,6 +203,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         runTest("compiler/fir/resolve/testData/resolve/intersectionTypes.kt");
     }
 
+    @TestMetadata("invokeInWhenSubjectVariableInitializer.kt")
+    public void testInvokeInWhenSubjectVariableInitializer() throws Exception {
+        runTest("compiler/fir/resolve/testData/resolve/invokeInWhenSubjectVariableInitializer.kt");
+    }
+
     @TestMetadata("invokeOfLambdaWithReceiver.kt")
     public void testInvokeOfLambdaWithReceiver() throws Exception {
         runTest("compiler/fir/resolve/testData/resolve/invokeOfLambdaWithReceiver.kt");

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/impl/FirWhenExpressionImpl.kt

@@ -58,6 +58,7 @@ internal class FirWhenExpressionImpl(
     override fun <D> transformSubject(transformer: FirTransformer<D>, data: D): FirWhenExpressionImpl {
         if (subjectVariable != null) {
             subjectVariable = subjectVariable?.transformSingle(transformer, data)
+            subject = subjectVariable?.initializer
         } else {
             subject = subject?.transformSingle(transformer, data)
         }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/implementation.kt

@@ -266,6 +266,7 @@ fun SmartPrinter.printImplementation(implementation: Implementation) {
                                 """
                                 |if (subjectVariable != null) {
                                 |            subjectVariable = subjectVariable?.transformSingle(transformer, data)
+                                |            subject = subjectVariable?.initializer
                                 |        } else {
                                 |            subject = subject?.transformSingle(transformer, data)
                                 |        }