[FIR] Add wrapped expressions for when subject expression with smartcast

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

@@ -116,8 +116,10 @@ object FirTreeBuilder : AbstractFirTreeBuilder() {
     val componentCall by element(Expression, functionCall)
     val callableReferenceAccess by element(Expression, qualifiedAccessExpression)
     val thisReceiverExpression by element(Expression, qualifiedAccessExpression)
-    val expressionWithSmartcast by element(Expression, qualifiedAccessExpression)
-    val expressionWithSmartcastToNull by element(Expression, expressionWithSmartcast)
+    val wrappedExpressionWithSmartcast by element(Expression)
+    val wrappedExpressionWithSmartcastToNull by element(Expression, wrappedExpressionWithSmartcast)
+    val expressionWithSmartcast by element(Expression, qualifiedAccessExpression, wrappedExpressionWithSmartcast)
+    val expressionWithSmartcastToNull by element(Expression, expressionWithSmartcast, wrappedExpressionWithSmartcastToNull)
     val safeCallExpression by element(Expression, expression)
     val checkedSafeCallSubject by element(Expression, expression)
     val getClassCall by element(Expression, expression, call)
@@ -136,6 +138,8 @@ object FirTreeBuilder : AbstractFirTreeBuilder() {
     val throwExpression by element(Expression, expression)
     val variableAssignment by element(Expression, qualifiedAccess)
     val whenSubjectExpression by element(Expression, expression)
+    val whenSubjectExpressionWithSmartcast by element(Expression, whenSubjectExpression, wrappedExpressionWithSmartcast)
+    val whenSubjectExpressionWithSmartcastToNull by element(Expression, whenSubjectExpression, wrappedExpressionWithSmartcastToNull)
 
     val wrappedDelegateExpression by element(Expression, wrappedExpression)
 

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

@@ -193,6 +193,9 @@ object ImplementationConfigurator : AbstractFirTreeImplementationConfigurator()
         noImpl(expressionWithSmartcast)
         noImpl(expressionWithSmartcastToNull)
 
+        noImpl(whenSubjectExpressionWithSmartcast)
+        noImpl(whenSubjectExpressionWithSmartcastToNull)
+
         impl(getClassCall) {
             default("argument") {
                 value = "argumentList.arguments.first()"

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

@@ -476,8 +476,10 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             +intField("componentIndex")
         }
 
-        expressionWithSmartcast.configure {
-            +field("originalExpression", qualifiedAccessExpression)
+        wrappedExpressionWithSmartcast.configure {
+            withArg("E", expression)
+            +typeRefField
+            +field("originalExpression", "E", packageName = null)
             +field("typesFromSmartCast", "Collection<ConeKotlinType>", null, customType = coneKotlinTypeType)
             +field("originalType", typeRef)
             +field("smartcastType", typeRef)
@@ -485,13 +487,18 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             +smartcastStability
         }
 
-        expressionWithSmartcastToNull.configure {
-            +field("originalExpression", qualifiedAccessExpression)
-            +field("typesFromSmartCast", "Collection<ConeKotlinType>", null, customType = coneKotlinTypeType)
-            +field("originalType", typeRef)
-            +field("smartcastType", typeRef)
+        wrappedExpressionWithSmartcastToNull.configure {
+            withArg("E", expression)
+            parentArg(wrappedExpressionWithSmartcast, "E", "E")
             +field("smartcastTypeWithoutNullableNothing", typeRef)
-            +smartcastStability
+        }
+
+        expressionWithSmartcast.configure {
+            parentArg(wrappedExpressionWithSmartcast, "E", qualifiedAccessExpression)
+        }
+
+        expressionWithSmartcastToNull.configure {
+            parentArg(wrappedExpressionWithSmartcastToNull, "E", qualifiedAccessExpression)
         }
 
         safeCallExpression.configure {
@@ -560,6 +567,14 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             +field("whenRef", whenRefType)
         }
 
+        whenSubjectExpressionWithSmartcast.configure {
+            parentArg(wrappedExpressionWithSmartcast, "E", whenSubjectExpression)
+        }
+
+        whenSubjectExpressionWithSmartcastToNull.configure {
+            parentArg(wrappedExpressionWithSmartcastToNull, "E", whenSubjectExpression)
+        }
+
         wrappedExpression.configure {
             +field(expression).withReplace()
         }