K2 LiftReturnOrAssignmentInspection: handles lift-return case

This inspection lifts return if it is possible. For example,

```
// Before:
fun foo(): Type {
    if (bar) return a
    else return b
}

// After:
fun foo(): Type {
    return if (bar) a
        else b
}
```

GitOrigin-RevId: 9e31be92ac476611d442e5342eb25cdcf2e6eb3a

plugins/kotlin/code-insight/inspections-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/inspections/LiftReturnOrAssignmentInspection.kt

@@ -12,9 +12,12 @@ import org.jetbrains.kotlin.analysis.api.analyze
 import org.jetbrains.kotlin.idea.base.psi.getLineCount
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.codeinsight.api.classic.inspections.AbstractKotlinInspection
+import org.jetbrains.kotlin.idea.codeinsight.utils.findExistingEditor
 import org.jetbrains.kotlin.idea.k2.codeinsight.inspections.branchedTransformations.BranchedFoldingUtils
+import org.jetbrains.kotlin.idea.k2.codeinsight.inspections.branchedTransformations.BranchedFoldingUtils.getFoldableReturnsFromBranches
 import org.jetbrains.kotlin.idea.k2.codeinsight.inspections.branchedTransformations.BranchedFoldingUtils.getNumberOfFoldableAssignmentsOrNull
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.anyDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 /**
@@ -54,10 +57,9 @@ private const val LINES_LIMIT = 15
  *       2 -> 3
  *       else -> 4
  *     }
- *
- * TODO: Handle the lift-return case.
  */
-class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
+class LiftReturnOrAssignmentInspection @JvmOverloads constructor(private val skipLongExpressions: Boolean = true) :
+    AbstractKotlinInspection() {
 
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean, session: LocalInspectionToolSession) =
         object : KtVisitorVoid() {
@@ -69,16 +71,27 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
 
                 // This inspection targets only return and assignment within expressions with branches.
                 // Their values must not be used by other expressions.
-                if (analyze(expression) { expression.isUsedAsExpression() }) return
+                if (expression.parent !is KtBlockExpression && analyze(expression) { expression.isUsedAsExpression() }) return
 
                 states.forEach { state ->
+                    val problemMessage = KotlinBundle.message(
+                        "0.1.be.lifted.out.of.2",
+                        when (state.liftType) {
+                            LiftType.LIFT_RETURN_OUT -> KotlinBundle.message("text.Return")
+                            LiftType.LIFT_ASSIGNMENT_OUT -> KotlinBundle.message("text.Assignment")
+                        },
+                        state.keyword.text,
+                    )
+
                     registerProblem(
                         expression,
                         state.keyword,
                         state.isSerious,
                         when (state.liftType) {
+                            LiftType.LIFT_RETURN_OUT -> LiftReturnOutFix(state.keyword.text)
                             LiftType.LIFT_ASSIGNMENT_OUT -> LiftAssignmentOutFix(state.keyword.text)
                         },
+                        problemMessage,
                         state.highlightElement,
                         state.highlightType,
                     )
@@ -90,36 +103,54 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
                 keyword: PsiElement,
                 isSerious: Boolean,
                 fix: LocalQuickFix,
+                message: String,
                 highlightElement: PsiElement = keyword,
                 highlightType: ProblemHighlightType = if (isSerious) GENERIC_ERROR_OR_WARNING else INFORMATION,
             ) {
-                val subject = KotlinBundle.message("text.Assignment")
                 holder.registerProblemWithoutOfflineInformation(
-                    expression,
-                    KotlinBundle.message("0.1.be.lifted.out.of.2", subject, keyword.text),
-                    isOnTheFly,
-                    highlightType,
-                    highlightElement.textRange?.shiftRight(-expression.startOffset),
-                    fix
+                    expression, message, isOnTheFly, highlightType, highlightElement.textRange?.shiftRight(-expression.startOffset), fix
                 )
             }
 
         }
 
+    private class LiftReturnOutFix(private val keyword: String) : LocalQuickFix {
+        override fun getName() = KotlinBundle.message("lift.return.out.fix.text.0", keyword)
+
+        override fun getFamilyName() = name
+
+        override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+            val element = descriptor.psiElement as? KtExpression ?: return
+            val replaced = BranchedFoldingUtils.foldToReturn(element)
+            replaced.findExistingEditor()?.caretModel?.moveToOffset(replaced.startOffset)
+        }
+    }
+
     private class LiftAssignmentOutFix(private val keyword: String) : LocalQuickFix {
         override fun getName() = KotlinBundle.message("lift.assignment.out.fix.text.0", keyword)
 
         override fun getFamilyName() = name
 
         override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
-            BranchedFoldingUtils.tryFoldToAssignment(descriptor.psiElement as KtExpression)
+            val element = descriptor.psiElement as? KtExpression ?: return
+            BranchedFoldingUtils.tryFoldToAssignment(element)
         }
     }
 
     private fun KtAnalysisSession.getStateForWhenOrTry(expression: KtExpression, keyword: PsiElement): List<LiftState>? {
-        if (expression.getLineCount() > LINES_LIMIT) return null
+        if (skipLongExpressions && expression.getLineCount() > LINES_LIMIT) return null
         if (expression.parent.node.elementType == KtNodeTypes.ELSE) return null
 
+        val foldableReturns = getFoldableReturnsFromBranches(expression)
+        if (foldableReturns.isNotEmpty()) {
+            val returns = foldableReturns.returnExpressions
+            val hasOtherReturns = expression.anyDescendantOfType<KtReturnExpression> { it !in returns }
+            val isSerious = !hasOtherReturns && returns.size > 1
+            return returns.map {
+                LiftState(keyword, isSerious, LiftType.LIFT_RETURN_OUT, it, INFORMATION)
+            } + LiftState(keyword, isSerious, LiftType.LIFT_RETURN_OUT)
+        }
+
         val assignmentNumber = getNumberOfFoldableAssignmentsOrNull(expression) ?: return null
         if (assignmentNumber > 0) {
             val isSerious = assignmentNumber > 1
@@ -140,13 +171,9 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
         }
     }
 
-    /**
-     * Types of lift.
-     *
-     * TODO: Add LIFT_RETURN_OUT and handle the lift-return case.
-     */
     enum class LiftType {
-        LIFT_ASSIGNMENT_OUT
+        LIFT_RETURN_OUT,
+        LIFT_ASSIGNMENT_OUT,
     }
 
     data class LiftState(
@@ -156,4 +183,4 @@ class LiftReturnOrAssignmentInspection : AbstractKotlinInspection() {
         val highlightElement: PsiElement = keyword,
         val highlightType: ProblemHighlightType = if (isSerious) GENERIC_ERROR_OR_WARNING else INFORMATION
     )
-}
+}

plugins/kotlin/code-insight/inspections-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/inspections/branchedTransformations/BranchedFoldingUtils.kt

@@ -8,6 +8,7 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.types.KtType
 import org.jetbrains.kotlin.analysis.api.types.KtTypeNullability
+import org.jetbrains.kotlin.idea.base.psi.replaced
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
@@ -122,7 +123,7 @@ object BranchedFoldingUtils {
             is KtWhenExpression -> {
                 val entries = e.entries
                 // When the KtWhenExpression has missing cases with an else branch, we cannot fold it.
-                if (!KtPsiUtil.checkWhenExpressionHasSingleElse(e) && e.getMissingCases().isNotEmpty()) false
+                if (e.hasMissingCases()) false
                 else entries.isNotEmpty() && entries.all { entry ->
                     val assignment = getFoldableBranchedAssignment(entry.expression)?.run { assignments.add(this) }
                     assignment != null || collectAssignmentsAndCheck(entry.expression?.lastBlockStatementOrThis())
@@ -224,4 +225,139 @@ object BranchedFoldingUtils {
         return nonNullableRightTypeOfFirst isEqualTo nonNullableRightTypeOfSecond ||
                 (first.operationToken == KtTokens.EQ && nonNullableRightTypeOfSecond isSubTypeOf leftType)
     }
-}
+
+    /**
+     * A function to lift `return` from return expressions in branches of if, when, or try [expression].
+     */
+    fun foldToReturn(expression: KtExpression): KtExpression {
+        fun KtReturnExpression.replaceWithReturned() {
+            returnedExpression?.let { replace(it) }
+        }
+
+        fun lift(e: KtExpression?) {
+            when (e) {
+                is KtWhenExpression -> e.entries.forEach { entry ->
+                    val entryExpr = entry.expression
+                    getFoldableBranchedReturn(entryExpr)?.replaceWithReturned() ?: lift(entryExpr?.lastBlockStatementOrThis())
+                }
+                is KtIfExpression -> e.branches.forEach { branch ->
+                    getFoldableBranchedReturn(branch)?.replaceWithReturned() ?: lift(branch?.lastBlockStatementOrThis())
+                }
+                is KtTryExpression -> e.tryBlockAndCatchBodies().forEach {
+                    getFoldableBranchedReturn(it)?.replaceWithReturned() ?: lift(it?.lastBlockStatementOrThis())
+                }
+            }
+        }
+        lift(expression)
+        return expression.replaced(KtPsiFactory(expression).createExpressionByPattern("return $0", expression))
+    }
+
+    /**
+     * Returns a return-expression inside [branch] or itself when the returned expression in the return-expression can be lifted.
+     * Otherwise, returns null.
+     *
+     * For example,
+     *     if (foo) {
+     *       return bar   // can be lifted -> this function will return `return bar`
+     *     } else {
+     *       return       // cannot be lifted because of the null returned expression -> this function will return `null`
+     *     }
+     */
+    private fun getFoldableBranchedReturn(branch: KtExpression?): KtReturnExpression? =
+        (branch?.lastBlockStatementOrThis() as? KtReturnExpression)?.takeIf {
+            it.returnedExpression != null &&
+                    it.returnedExpression !is KtLambdaExpression &&
+                    it.getTargetLabel() == null
+        }
+
+    data class FoldableReturns(val returnExpressions : List<KtReturnExpression>, val isFoldable : Boolean) {
+        fun isNotEmpty(): Boolean = isFoldable && returnExpressions.isNotEmpty()
+
+        companion object {
+            val NotFoldable = FoldableReturns(emptyList(), false)
+        }
+    }
+
+    /**
+     * Returns a list of return-expressions inside expressions [branches] that are branches of a if, when, or try expression.
+     * If there is any branch that we cannot lift its returned expression, this function returns an empty list with 'false' for
+     * `isFoldable`.
+     *
+     * For example, it returns `return bar` and `return zoo` expressions for the following code:
+     *     if (foo) {
+     *       return bar
+     *     } else {
+     *       return zoo
+     *     }
+     *
+     * It returns an empty list with 'false' for `isFoldable` for the following code because we cannot lift the return expression in the
+     * else-branch:
+     *     if (foo) {
+     *       return bar   // can be lifted
+     *     } else {
+     *       return       // cannot be lifted because of the null returned expression
+     *     }
+     */
+    private fun KtAnalysisSession.getFoldableReturnsFromBranches(branches: List<KtExpression?>): FoldableReturns {
+        val foldableReturns = mutableListOf<KtReturnExpression>()
+        for (branch in branches) {
+            val foldableBranchedReturn = getFoldableBranchedReturn(branch)
+            if (foldableBranchedReturn != null) {
+                foldableReturns.add(foldableBranchedReturn)
+            } else {
+                val currReturns = branch?.lastBlockStatementOrThis()?.let { getFoldableReturnsFromBranches(it) }
+                    ?: return FoldableReturns.NotFoldable
+                if (!currReturns.isFoldable) return FoldableReturns.NotFoldable
+                foldableReturns += currReturns.returnExpressions
+            }
+        }
+        return FoldableReturns(foldableReturns, true)
+    }
+
+    /**
+     * Returns a list of return-expressions that can be lifted from if, when, or try expression [expression].
+     *
+     * It returns an empty list with `isFoldable = false` if [expression] is one of if, when, and try expressions and
+     *  - [expression] doesn't have an else-branch, and it has a missing case, or
+     *  - there is any branch that we cannot lift its returned expression
+     *
+     * It returns an empty list with `isFoldable = true` if [expression] is one of [KtBreakExpression], [KtContinueExpression],
+     * [KtThrowExpression], and [KtCallExpression].
+     */
+    fun KtAnalysisSession.getFoldableReturnsFromBranches(expression: KtExpression): FoldableReturns = when (expression) {
+        is KtWhenExpression -> {
+            val entries = expression.entries
+            when {
+                expression.hasMissingCases() -> FoldableReturns.NotFoldable
+                entries.isEmpty() -> FoldableReturns.NotFoldable
+                else -> getFoldableReturnsFromBranches(entries.map { it.expression })
+            }
+        }
+        is KtIfExpression -> {
+            val branches = expression.branches
+            when {
+                branches.isEmpty() -> FoldableReturns.NotFoldable
+                branches.lastOrNull()?.getStrictParentOfType<KtIfExpression>()?.`else` == null -> FoldableReturns.NotFoldable
+                else -> getFoldableReturnsFromBranches(branches)
+            }
+        }
+        is KtTryExpression -> {
+            if (expression.finallyBlock?.finalExpression?.let { getFoldableReturnsFromBranches(listOf(it)) }?.isNotEmpty() == true)
+                FoldableReturns.NotFoldable
+            else
+                getFoldableReturnsFromBranches(expression.tryBlockAndCatchBodies())
+        }
+        is KtCallExpression -> {
+            if (expression.getKtType()?.isNothing == true) FoldableReturns(emptyList(), true) else FoldableReturns.NotFoldable
+        }
+        is KtBreakExpression, is KtContinueExpression, is KtThrowExpression -> FoldableReturns(emptyList(), true)
+        else -> FoldableReturns.NotFoldable
+    }
+
+    /**
+     * Returns true if the when-expression has a missing case with else-branch.
+     */
+    context(KtAnalysisSession)
+    private fun KtWhenExpression.hasMissingCases(): Boolean =
+        !KtPsiUtil.checkWhenExpressionHasSingleElse(this) && getMissingCases().isNotEmpty()
+}

plugins/kotlin/code-insight/inspections-k2/tests/test/org/jetbrains/kotlin/idea/k2/intentions/tests/K2LocalInspectionTestGenerated.java

@@ -924,6 +924,200 @@ public abstract class K2LocalInspectionTestGenerated extends AbstractK2LocalInsp
         }
     }
 
+    @RunWith(JUnit3RunnerWithInners.class)
+    @TestMetadata("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn")
+    public static class IfToReturn extends AbstractK2LocalInspectionTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        @TestMetadata("cascadeIf.kt")
+        public void testCascadeIf() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/cascadeIf.kt");
+        }
+
+        @TestMetadata("ifElseIf.kt")
+        public void testIfElseIf() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/ifElseIf.kt");
+        }
+
+        @TestMetadata("ifElseIfElse.kt")
+        public void testIfElseIfElse() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/ifElseIfElse.kt");
+        }
+
+        @TestMetadata("ifElseIfElseInconsistent.kt")
+        public void testIfElseIfElseInconsistent() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/ifElseIfElseInconsistent.kt");
+        }
+
+        @TestMetadata("ifVeryLong.kt")
+        public void testIfVeryLong() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/ifVeryLong.kt");
+        }
+
+        @TestMetadata("innerIfTransformed.kt")
+        public void testInnerIfTransformed() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/innerIfTransformed.kt");
+        }
+
+        @TestMetadata("onReturn.kt")
+        public void testOnReturn() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/onReturn.kt");
+        }
+
+        @TestMetadata("onReturn2.kt")
+        public void testOnReturn2() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/onReturn2.kt");
+        }
+
+        @TestMetadata("simpleIf.kt")
+        public void testSimpleIf() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/simpleIf.kt");
+        }
+
+        @TestMetadata("simpleIfWithBlocks.kt")
+        public void testSimpleIfWithBlocks() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/simpleIfWithBlocks.kt");
+        }
+    }
+
+    @RunWith(JUnit3RunnerWithInners.class)
+    @TestMetadata("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn")
+    public static class TryToReturn extends AbstractK2LocalInspectionTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        @TestMetadata("basic.kt")
+        public void testBasic() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/basic.kt");
+        }
+
+        @TestMetadata("block.kt")
+        public void testBlock() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/block.kt");
+        }
+
+        @TestMetadata("cascade.kt")
+        public void testCascade() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/cascade.kt");
+        }
+
+        @TestMetadata("finally.kt")
+        public void testFinally() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/finally.kt");
+        }
+
+        @TestMetadata("finallyWithCascadeReturn.kt")
+        public void testFinallyWithCascadeReturn() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/finallyWithCascadeReturn.kt");
+        }
+
+        @TestMetadata("finallyWithReturn.kt")
+        public void testFinallyWithReturn() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/finallyWithReturn.kt");
+        }
+
+        @TestMetadata("inner.kt")
+        public void testInner() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/inner.kt");
+        }
+
+        @TestMetadata("onReturn.kt")
+        public void testOnReturn() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/onReturn.kt");
+        }
+
+        @TestMetadata("withoutReturn.kt")
+        public void testWithoutReturn() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/withoutReturn.kt");
+        }
+    }
+
+    @RunWith(JUnit3RunnerWithInners.class)
+    @TestMetadata("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn")
+    public static class WhenToReturn extends AbstractK2LocalInspectionTest {
+        private void runTest(String testDataFilePath) throws Exception {
+            KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+        }
+
+        @TestMetadata("cascadeWhen.kt")
+        public void testCascadeWhen() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/cascadeWhen.kt");
+        }
+
+        @TestMetadata("innerWhenTransformed.kt")
+        public void testInnerWhenTransformed() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/innerWhenTransformed.kt");
+        }
+
+        @TestMetadata("insideLoop.kt")
+        public void testInsideLoop() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/insideLoop.kt");
+        }
+
+        @TestMetadata("localReturns.kt")
+        public void testLocalReturns() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/localReturns.kt");
+        }
+
+        @TestMetadata("onReturn.kt")
+        public void testOnReturn() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/onReturn.kt");
+        }
+
+        @TestMetadata("onReturn2.kt")
+        public void testOnReturn2() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/onReturn2.kt");
+        }
+
+        @TestMetadata("otherReturns.kt")
+        public void testOtherReturns() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/otherReturns.kt");
+        }
+
+        @TestMetadata("simpleWhen.kt")
+        public void testSimpleWhen() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/simpleWhen.kt");
+        }
+
+        @TestMetadata("simpleWhenWithBlocks.kt")
+        public void testSimpleWhenWithBlocks() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/simpleWhenWithBlocks.kt");
+        }
+
+        @TestMetadata("whenHasMissingCase.kt")
+        public void testWhenHasMissingCase() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/whenHasMissingCase.kt");
+        }
+
+        @TestMetadata("whenHasNoMissingCase.kt")
+        public void testWhenHasNoMissingCase() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/whenHasNoMissingCase.kt");
+        }
+
+        @TestMetadata("whenHasNoMissingCaseWithElse.kt")
+        public void testWhenHasNoMissingCaseWithElse() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/whenHasNoMissingCaseWithElse.kt");
+        }
+
+        @TestMetadata("whenOneReturn.kt")
+        public void testWhenOneReturn() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/whenOneReturn.kt");
+        }
+
+        @TestMetadata("whenThrowOnly.kt")
+        public void testWhenThrowOnly() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/whenThrowOnly.kt");
+        }
+
+        @TestMetadata("whenWithLambda.kt")
+        public void testWhenWithLambda() throws Exception {
+            runTest("../../../idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/whenWithLambda.kt");
+        }
+    }
+
     @RunWith(JUnit3RunnerWithInners.class)
     @TestMetadata("testData/inspectionsLocal")
     public abstract static class InspectionsLocal extends AbstractK2LocalInspectionTest {

plugins/kotlin/idea/tests/testData/inspectionsLocal/liftOut/ifToReturn/.k2Inspection

@@ -0,0 +1 @@
+org.jetbrains.kotlin.idea.k2.codeinsight.inspections.LiftReturnOrAssignmentInspection

plugins/kotlin/idea/tests/testData/inspectionsLocal/liftOut/tryToReturn/.k2Inspection

@@ -0,0 +1 @@
+org.jetbrains.kotlin.idea.k2.codeinsight.inspections.LiftReturnOrAssignmentInspection

plugins/kotlin/idea/tests/testData/inspectionsLocal/liftOut/whenToReturn/.k2Inspection

@@ -0,0 +1 @@
+org.jetbrains.kotlin.idea.k2.codeinsight.inspections.LiftReturnOrAssignmentInspection

plugins/kotlin/util/test-generator-fir/test/org/jetbrains/kotlin/fir/testGenerator/codeinsight/GenerateK2InspectionTests.kt

@@ -29,6 +29,9 @@ internal fun MutableTWorkspace.generateK2InspectionTests() {
             model("${idea}/inspectionsLocal/liftOut/ifToAssignment")
             model("${idea}/inspectionsLocal/liftOut/tryToAssignment")
             model("${idea}/inspectionsLocal/liftOut/whenToAssignment")
+            model("${idea}/inspectionsLocal/liftOut/ifToReturn")
+            model("${idea}/inspectionsLocal/liftOut/tryToReturn")
+            model("${idea}/inspectionsLocal/liftOut/whenToReturn")
             model("code-insight/inspections-k2/tests/testData/inspectionsLocal", pattern = pattern)
         }