K2: implement FirForLoopStatementAssignmentChecker

#KT-60006 Fixed

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/CommonExpressionCheckers.kt

@@ -106,6 +106,7 @@ object CommonExpressionCheckers : ExpressionCheckers() {
     override val loopExpressionCheckers: Set<FirLoopExpressionChecker>
         get() = setOf(
             FirLoopConditionChecker,
+            FirForLoopStatementAssignmentChecker
         )
 
     override val loopJumpCheckers: Set<FirLoopJumpChecker>
@@ -128,7 +129,7 @@ object CommonExpressionCheckers : ExpressionCheckers() {
         get() = setOf(
             FirForLoopChecker,
             FirConflictsExpressionChecker,
-            FirSingleNamedFunctionChecker
+            FirSingleNamedFunctionChecker,
         )
 
     override val checkNotNullCallCheckers: Set<FirCheckNotNullCallChecker>

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirForLoopStatementAssignmentChecker.kt

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.analysis.checkers.expression
+
+import org.jetbrains.kotlin.KtFakeSourceElementKind
+import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
+import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
+import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
+import org.jetbrains.kotlin.fir.declarations.FirDeclarationOrigin
+import org.jetbrains.kotlin.fir.declarations.FirProperty
+import org.jetbrains.kotlin.fir.declarations.utils.FirScriptCustomizationKind
+import org.jetbrains.kotlin.fir.expressions.FirBlock
+import org.jetbrains.kotlin.fir.expressions.FirErrorExpression
+import org.jetbrains.kotlin.fir.expressions.FirLoop
+import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
+
+object FirForLoopStatementAssignmentChecker : FirLoopExpressionChecker(MppCheckerKind.Common) {
+    override fun check(expression: FirLoop, context: CheckerContext, reporter: DiagnosticReporter) {
+        // Checks the pattern for desugared for loop.
+        val parent = if (context.containingElements.size >= 2) context.containingElements[context.containingElements.size - 2] else return
+        if (parent.source?.kind != KtFakeSourceElementKind.DesugaredForLoop) return
+
+        val grandParent = if (context.containingElements.size >= 3)
+            context.containingElements[context.containingElements.size - 3]
+        else
+            return
+
+        if (// It is used as a statement
+            grandParent is FirBlock
+            // It is used as a single statement in the method body
+            || (grandParent is FirReturnExpression && grandParent.source?.kind == KtFakeSourceElementKind.ImplicitReturn.FromLastStatement)
+            // It is used in a kotlin script as a last statement
+            || (grandParent is FirProperty && (grandParent.origin as? FirDeclarationOrigin.ScriptCustomization)?.kind == FirScriptCustomizationKind.RESULT_PROPERTY)
+            // There was a fail before (for example, using two labels before the for loop)
+            || (grandParent is FirErrorExpression)
+        )
+            return
+
+        reporter.reportOn(expression.source, FirErrors.EXPRESSION_EXPECTED, context)
+    }
+}

compiler/testData/diagnostics/tests/controlStructures/ForbidStatementAsDirectFunctionBody.fir.kt

@@ -1,11 +0,0 @@
-
-fun foo1() = <!EXPRESSION_EXPECTED!>while (b()) {}<!>
-
-fun foo2() = for (i in <!ITERATOR_MISSING!>10<!>) {}
-
-fun foo3() = when (b()) {
-    true -> 1
-    else -> 0
-}
-
-fun b(): Boolean = true

compiler/testData/diagnostics/tests/controlStructures/ForbidStatementAsDirectFunctionBody.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 fun foo1() = <!EXPRESSION_EXPECTED!>while (b()) {}<!>
 
 fun foo2() = <!EXPRESSION_EXPECTED!>for (i in <!ITERATOR_MISSING!>10<!>) {}<!>

compiler/testData/diagnostics/tests/controlStructures/kt770.kt351.kt735_StatementType.fir.kt

@@ -149,7 +149,7 @@ fun bar(a: Unit) {}
 fun testStatementInExpressionContext() {
     var z = 34
     val a1: Unit = <!ASSIGNMENT_IN_EXPRESSION_CONTEXT!>z = 334<!>
-    val f = for (i in 1..10) {}
+    val f = <!EXPRESSION_EXPECTED!>for (i in 1..10) {}<!>
     if (true) return <!ASSIGNMENT_IN_EXPRESSION_CONTEXT!>z = 34<!>
     return <!EXPRESSION_EXPECTED!>while (true) {}<!>
 }

compiler/tests-spec/testData/diagnostics/linked/statements/assignments/p-1/neg/2.1.fir.kt

@@ -16,8 +16,8 @@
  * NOTE: right-hand side of an assignment must be expression
  */
 fun case1() {
-    val x = <!ITERATOR_AMBIGUITY!>for (<!SYNTAX!><!>) { }<!>
-    val y = for (x in 1..2) { }
+    val x = <!EXPRESSION_EXPECTED, ITERATOR_AMBIGUITY!>for (<!SYNTAX!><!>) { }<!>
+    val y = <!EXPRESSION_EXPECTED!>for (x in 1..2) { }<!>
 
     val a = <!EXPRESSION_EXPECTED!>while (<!SYNTAX!><!>) { }<!>
     val b = <!EXPRESSION_EXPECTED!>while (false) { }<!>
@@ -29,8 +29,8 @@ fun case1() {
  * NOTE: right-hand side of an assignment must be expression
  */
 fun case2() {
-    var x = <!ITERATOR_AMBIGUITY!>for (<!SYNTAX!><!>) { }<!>
-    var y = for (x in 1..2) { }
+    var x = <!EXPRESSION_EXPECTED, ITERATOR_AMBIGUITY!>for (<!SYNTAX!><!>) { }<!>
+    var y = <!EXPRESSION_EXPECTED!>for (x in 1..2) { }<!>
 
     var a = <!EXPRESSION_EXPECTED!>while (<!SYNTAX!><!>) { }<!>
     var b = <!EXPRESSION_EXPECTED!>while (false) { }<!>
@@ -49,8 +49,8 @@ fun case3() {
     var b :Any?
     var c :Any?
 
-    x = <!ITERATOR_AMBIGUITY!>for (<!SYNTAX!><!>) { }<!>
-    y = for (x in 1..2) { }
+    x = <!EXPRESSION_EXPECTED, ITERATOR_AMBIGUITY!>for (<!SYNTAX!><!>) { }<!>
+    y = <!EXPRESSION_EXPECTED!>for (x in 1..2) { }<!>
 
     a = <!EXPRESSION_EXPECTED!>while (<!SYNTAX!><!>) { }<!>
     b = <!EXPRESSION_EXPECTED!>while (false) { }<!>
@@ -62,7 +62,7 @@ fun case3() {
  * NOTE: left-hand side of an assignment must be expression
  */
 fun case4() {
-    <!VARIABLE_EXPECTED!>for (x in 1..2) {}<!>  = TODO();
+    <!EXPRESSION_EXPECTED, VARIABLE_EXPECTED!>for (x in 1..2) {}<!>  = TODO();
 
     <!EXPRESSION_EXPECTED, VARIABLE_EXPECTED!>while (false) { }<!> = TODO()
 }