[FIR] Save info that when was used as expression

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

@@ -499,6 +499,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
         runTest("compiler/fir/analysis-tests/testData/resolve/whenInference.kt");
     }
 
+    @TestMetadata("whenWithWhenAsStatement.kt")
+    public void testWhenWithWhenAsStatement() throws Exception {
+        runTest("compiler/fir/analysis-tests/testData/resolve/whenWithWhenAsStatement.kt");
+    }
+
     @TestMetadata("compiler/fir/analysis-tests/testData/resolve/arguments")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)

compiler/fir/analysis-tests/testData/resolve/whenWithWhenAsStatement.fir.txt

@@ -0,0 +1,16 @@
+FILE: whenWithWhenAsStatement.kt
+    public final fun test(value: R|kotlin/Int|): R|kotlin/Unit| {
+        when (R|<local>/value|) {
+            ==($subj$, Int(0)) ->  {
+            }
+            ==($subj$, Int(1)) ->  {
+                when (R|<local>/value|) {
+                    ==($subj$, Int(2)) ->  {
+                        Boolean(false)
+                    }
+                }
+
+            }
+        }
+
+    }

compiler/fir/analysis-tests/testData/resolve/whenWithWhenAsStatement.kt

@@ -0,0 +1,8 @@
+fun test(value: Int) {
+    when (value) {
+        0 -> {}
+        1 -> when (value) {
+            2 -> false
+        }
+    }
+}

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

@@ -590,6 +590,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/whenInference.kt");
         }
 
+        @Test
+        @TestMetadata("whenWithWhenAsStatement.kt")
+        public void testWhenWithWhenAsStatement() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/whenWithWhenAsStatement.kt");
+        }
+
         @Nested
         @TestMetadata("compiler/fir/analysis-tests/testData/resolve/arguments")
         @TestDataPath("$PROJECT_ROOT")

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

@@ -593,6 +593,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/whenInference.kt");
         }
 
+        @Test
+        @TestMetadata("whenWithWhenAsStatement.kt")
+        public void testWhenWithWhenAsStatement() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/whenWithWhenAsStatement.kt");
+        }
+
         @Nested
         @TestMetadata("compiler/fir/analysis-tests/testData/resolve/arguments")
         @TestDataPath("$PROJECT_ROOT")

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

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.fir.lightTree.converter
 import com.intellij.lang.LighterASTNode
 import com.intellij.psi.TokenType
 import com.intellij.util.diff.FlyweightCapableTreeStructure
+import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.KtNodeTypes.*
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.Visibilities
@@ -664,6 +665,7 @@ class ExpressionsConverter(
             source = whenExpression.toFirSourceElement()
             this.subject = subjectExpression
             this.subjectVariable = subjectVariable
+            usedAsExpression = whenExpression.usedAsExpression
             for (entry in whenEntries) {
                 val branch = entry.firBlock
                 branches += if (!entry.isElse) {
@@ -1095,11 +1097,24 @@ class ExpressionsConverter(
                     condition = buildElseIfTrueCondition()
                     result = elseBranch
                 }
-
             }
+            usedAsExpression = ifExpression.usedAsExpression
         }
     }
 
+    private val LighterASTNode.usedAsExpression: Boolean
+        get() {
+            val parent = getParent() ?: return true
+            val parentTokenType = parent.tokenType
+            if (parentTokenType == BLOCK) return false
+            if (parentTokenType == ELSE || parentTokenType == WHEN_ENTRY) {
+                return parent.getParent()?.usedAsExpression ?: true
+            }
+            if (parentTokenType != BODY) return true
+            val type = parent.getParent()?.tokenType ?: return true
+            return !(type == FOR || type == WHILE || type == DO_WHILE)
+        }
+
     /**
      * @see org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseJump
      * @see org.jetbrains.kotlin.fir.builder.RawFirBuilder.Visitor.visitBreakExpression

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1546,6 +1546,7 @@ class RawFirBuilder(
                         result = expression.`else`.toFirBlock()
                     }
                 }
+                usedAsExpression = expression.usedAsExpression
             }
         }
 
@@ -1583,6 +1584,7 @@ class RawFirBuilder(
                 source = expression.toFirSourceElement()
                 this.subject = subjectExpression
                 this.subjectVariable = subjectVariable
+                usedAsExpression = expression.usedAsExpression
 
                 for (entry in expression.entries) {
                     val entrySource = entry.toFirSourceElement()
@@ -1621,6 +1623,20 @@ class RawFirBuilder(
             }
         }
 
+        private val KtExpression.usedAsExpression: Boolean
+            get() {
+                if (parent is KtBlockExpression) return false
+                when (parent.elementType) {
+                    KtNodeTypes.ELSE, KtNodeTypes.WHEN_ENTRY -> {
+                        return (parent.parent as? KtExpression)?.usedAsExpression ?: true
+                    }
+                }
+                // Here we check that when used is a single statement of a loop
+                if (parent !is KtContainerNodeForControlStructureBody) return true
+                val type = parent.parent.elementType
+                return !(type == KtNodeTypes.FOR || type == KtNodeTypes.WHILE || type == KtNodeTypes.DO_WHILE)
+            }
+
         override fun visitDoWhileExpression(expression: KtDoWhileExpression, data: Unit): FirElement {
             return FirDoWhileLoopBuilder().apply {
                 source = expression.toFirSourceElement()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/CopyUtils.kt

@@ -137,6 +137,8 @@ fun FirWhenExpression.copy(
     branches += this@copy.branches
     typeRef = resultType
     this.annotations += annotations
+    usedAsExpression = this@copy.usedAsExpression
+    isExhaustive = this@copy.isExhaustive
 }
 
 fun FirTryExpression.copy(

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirWhenExpression.kt

@@ -25,6 +25,7 @@ abstract class FirWhenExpression : FirExpression(), FirResolvable {
     abstract val subjectVariable: FirVariable<*>?
     abstract val branches: List<FirWhenBranch>
     abstract val isExhaustive: Boolean
+    abstract val usedAsExpression: Boolean
 
     override fun <R, D> accept(visitor: FirVisitor<R, D>, data: D): R = visitor.visitWhenExpression(this, data)
 

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/builder/FirWhenExpressionBuilder.kt

@@ -37,6 +37,7 @@ class FirWhenExpressionBuilder : FirAnnotationContainerBuilder, FirExpressionBui
     var subjectVariable: FirVariable<*>? = null
     val branches: MutableList<FirWhenBranch> = mutableListOf()
     var isExhaustive: Boolean = false
+    var usedAsExpression: Boolean by kotlin.properties.Delegates.notNull<Boolean>()
 
     override fun build(): FirWhenExpression {
         return FirWhenExpressionImpl(
@@ -48,13 +49,14 @@ class FirWhenExpressionBuilder : FirAnnotationContainerBuilder, FirExpressionBui
             subjectVariable,
             branches,
             isExhaustive,
+            usedAsExpression,
         )
     }
 
 }
 
 @OptIn(ExperimentalContracts::class)
-inline fun buildWhenExpression(init: FirWhenExpressionBuilder.() -> Unit = {}): FirWhenExpression {
+inline fun buildWhenExpression(init: FirWhenExpressionBuilder.() -> Unit): FirWhenExpression {
     contract {
         callsInPlace(init, kotlin.contracts.InvocationKind.EXACTLY_ONCE)
     }

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

@@ -29,6 +29,7 @@ internal class FirWhenExpressionImpl(
     override var subjectVariable: FirVariable<*>?,
     override val branches: MutableList<FirWhenBranch>,
     override var isExhaustive: Boolean,
+    override val usedAsExpression: Boolean,
 ) : FirWhenExpression() {
     override fun <R, D> acceptChildren(visitor: FirVisitor<R, D>, data: D) {
         typeRef.accept(visitor, data)

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

@@ -588,6 +588,7 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
             +field("subjectVariable", variable.withArgs("F" to "*"), nullable = true)
             +fieldList("branches", whenBranch).withTransform()
             +booleanField("isExhaustive", withReplace = true)
+            +booleanField("usedAsExpression")
             needTransformOtherChildren()
         }