[kotlin] Highlight loop exit points

#KTIJ-18471 Fixed


Merge-request: IJ-MR-112687
Merged-by: Vladimir Dolzhenko <Vladimir.Dolzhenko@jetbrains.com>

GitOrigin-RevId: 573f56bc4e721c754f453f11308ef1acdf1ed153

plugins/kotlin/highlighting/highlighting-k2/test/org/jetbrains/kotlin/idea/k2/highlighting/K2HighlightExitPointsTestGenerated.java

@@ -33,6 +33,36 @@ public class K2HighlightExitPointsTestGenerated extends AbstractK2HighlightExitP
         runTest("../../idea/tests/testData/exitPoints/anonymousFunction2.kt");
     }
 
+    @TestMetadata("forLoopBreak.kt")
+    public void testForLoopBreak() throws Exception {
+        runTest("../../idea/tests/testData/exitPoints/forLoopBreak.kt");
+    }
+
+    @TestMetadata("forLoopNestedBreak.kt")
+    public void testForLoopNestedBreak() throws Exception {
+        runTest("../../idea/tests/testData/exitPoints/forLoopNestedBreak.kt");
+    }
+
+    @TestMetadata("forLoopNestedBreak2.kt")
+    public void testForLoopNestedBreak2() throws Exception {
+        runTest("../../idea/tests/testData/exitPoints/forLoopNestedBreak2.kt");
+    }
+
+    @TestMetadata("forLoopNestedLabeledBreak.kt")
+    public void testForLoopNestedLabeledBreak() throws Exception {
+        runTest("../../idea/tests/testData/exitPoints/forLoopNestedLabeledBreak.kt");
+    }
+
+    @TestMetadata("forLoopNestedLabeledBreakOnLoop.kt")
+    public void testForLoopNestedLabeledBreakOnLoop() throws Exception {
+        runTest("../../idea/tests/testData/exitPoints/forLoopNestedLabeledBreakOnLoop.kt");
+    }
+
+    @TestMetadata("forLoopNestedLabeledBreakOnLoop2.kt")
+    public void testForLoopNestedLabeledBreakOnLoop2() throws Exception {
+        runTest("../../idea/tests/testData/exitPoints/forLoopNestedLabeledBreakOnLoop2.kt");
+    }
+
     @TestMetadata("funKeyword.kt")
     public void testFunKeyword() throws Exception {
         runTest("../../idea/tests/testData/exitPoints/funKeyword.kt");

plugins/kotlin/highlighting/highlighting-shared/src/org/jetbrains/kotlin/idea/base/highlighting/AbstractKotlinHighlightExitPointsHandlerFactory.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.parents
 import org.jetbrains.kotlin.resolve.inline.InlineUtil
 import org.jetbrains.kotlin.utils.addIfNotNull
+import java.util.*
 
 abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsagesHandlerFactoryBase() {
     private fun getOnReturnOrThrowUsageHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
@@ -35,6 +36,15 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
         return OnExitUsagesHandler(editor, file, expression, false)
     }
 
+    private fun getOnBreakOrContinueUsageHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
+        val expression = when (val parent = target.parent) {
+            is KtBreakExpression, is KtContinueExpression -> parent
+            is KtLoopExpression -> parent
+            else -> null
+        } as? KtExpression ?: return null
+        return OnLoopUsagesHandler(editor, file, expression, false)
+    }
+
     private fun getOnLambdaCallUsageHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
         if (target !is LeafPsiElement
             || target.elementType !is KtToken // do not trigger loading of KtTokens in Java
@@ -57,6 +67,7 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
 
     override fun createHighlightUsagesHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
         return getOnReturnOrThrowUsageHandler(editor, file, target)
+            ?: getOnBreakOrContinueUsageHandler(editor, file, target)
             ?: getOnLambdaCallUsageHandler(editor, file, target)
     }
 
@@ -91,7 +102,7 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
     private inner class OnExitUsagesHandler(editor: Editor, file: PsiFile, val target: KtExpression, val highlightReferences: Boolean) :
         HighlightUsagesHandlerBase<PsiElement>(editor, file) {
 
-        override fun getTargets() = listOf(target)
+        override fun getTargets(): List<KtExpression> = listOf(target)
 
         override fun selectTargets(targets: List<PsiElement>, selectionConsumer: Consumer<in List<PsiElement>>) {
             selectionConsumer.consume(targets)
@@ -290,6 +301,89 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
         override fun highlightReferences(): Boolean = highlightReferences
     }
 
+    private inner class OnLoopUsagesHandler(editor: Editor, file: PsiFile, val target: KtExpression, val highlightReferences: Boolean) :
+        HighlightUsagesHandlerBase<PsiElement>(editor, file) {
+        override fun getTargets(): List<KtExpression> = listOf(target)
+
+        override fun selectTargets(targets: List<PsiElement>, selectionConsumer: Consumer<in List<PsiElement>>) {
+            selectionConsumer.consume(targets)
+        }
+
+        override fun computeUsages(targets: MutableList<out PsiElement>) {
+            val labelName = when (target) {
+                is KtExpressionWithLabel -> target.getLabelName()
+                is KtLoopExpression -> (target.parent as? KtLabeledExpression)?.getLabelName()
+                else -> null
+            }
+            val relevantLoop: KtLoopExpression = when (target) {
+                is KtLoopExpression -> target
+                else -> {
+                    var element: PsiElement? = target
+                    var targetLoop: KtLoopExpression? = null
+                    while (element != null) {
+                        val parent = element.parent
+                        if (element is KtLoopExpression && (labelName == null || (parent as? KtLabeledExpression)?.getLabelName() == labelName)) {
+                            targetLoop = element
+                            break
+                        }
+                        element = parent
+                    }
+                    targetLoop
+                }
+            } ?: return
+
+            when(relevantLoop) {
+                is KtForExpression -> addOccurrence(relevantLoop.forKeyword)
+                is KtDoWhileExpression -> relevantLoop.node.findChildByType(KtTokens.DO_KEYWORD)?.psi?.let(::addOccurrence)
+                is KtWhileExpression -> relevantLoop.node.findChildByType(KtTokens.WHILE_KEYWORD)?.psi?.let(::addOccurrence)
+            }
+
+
+
+            relevantLoop.accept(object : KtVisitorVoid() {
+                var nestedLoopExpressions = Stack<KtLoopExpression>()
+
+                override fun visitKtElement(element: KtElement) {
+                    ProgressIndicatorProvider.checkCanceled()
+                    element.acceptChildren(this)
+                }
+
+                override fun visitExpression(expression: KtExpression) {
+                    val nestedLoopFound = if (expression != relevantLoop && expression is KtLoopExpression) {
+                        val loopLabelName = (expression.parent as? KtLabeledExpression)?.getLabelName()
+                        // no reasons to step into another loop with the same label name or no label name
+                        if (labelName == null || labelName == loopLabelName) return
+
+                        nestedLoopExpressions.push(expression)
+                        true
+                    } else {
+                        false
+                    }
+
+                    if (expression is KtBreakExpression || expression is KtContinueExpression) {
+                        val expressionLabelName = (expression as? KtExpressionWithLabel)?.getLabelName()
+                        if (nestedLoopExpressions.isEmpty()) {
+                            if (expressionLabelName == null || expressionLabelName == labelName) {
+                                addOccurrence(expression)
+                            }
+                        } else if (expressionLabelName == labelName) {
+                            addOccurrence(expression)
+                        }
+                    }
+
+                    try {
+                        super.visitExpression(expression)
+                    } finally {
+                        if (nestedLoopFound) {
+                            nestedLoopExpressions.pop()
+                        }
+                    }
+                }
+            })
+        }
+
+    }
+
     private fun MutableSet<PsiElement>.addIfNotNullAndNotBlock(element: PsiElement?) {
         addIfNotNull(element.takeUnless { it is KtBlockExpression })
     }

plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/highlighter/HighlightExitPointsTestGenerated.java

@@ -33,6 +33,36 @@ public class HighlightExitPointsTestGenerated extends AbstractHighlightExitPoint
         runTest("testData/exitPoints/anonymousFunction2.kt");
     }
 
+    @TestMetadata("forLoopBreak.kt")
+    public void testForLoopBreak() throws Exception {
+        runTest("testData/exitPoints/forLoopBreak.kt");
+    }
+
+    @TestMetadata("forLoopNestedBreak.kt")
+    public void testForLoopNestedBreak() throws Exception {
+        runTest("testData/exitPoints/forLoopNestedBreak.kt");
+    }
+
+    @TestMetadata("forLoopNestedBreak2.kt")
+    public void testForLoopNestedBreak2() throws Exception {
+        runTest("testData/exitPoints/forLoopNestedBreak2.kt");
+    }
+
+    @TestMetadata("forLoopNestedLabeledBreak.kt")
+    public void testForLoopNestedLabeledBreak() throws Exception {
+        runTest("testData/exitPoints/forLoopNestedLabeledBreak.kt");
+    }
+
+    @TestMetadata("forLoopNestedLabeledBreakOnLoop.kt")
+    public void testForLoopNestedLabeledBreakOnLoop() throws Exception {
+        runTest("testData/exitPoints/forLoopNestedLabeledBreakOnLoop.kt");
+    }
+
+    @TestMetadata("forLoopNestedLabeledBreakOnLoop2.kt")
+    public void testForLoopNestedLabeledBreakOnLoop2() throws Exception {
+        runTest("testData/exitPoints/forLoopNestedLabeledBreakOnLoop2.kt");
+    }
+
     @TestMetadata("funKeyword.kt")
     public void testFunKeyword() throws Exception {
         runTest("testData/exitPoints/funKeyword.kt");

plugins/kotlin/idea/tests/testData/exitPoints/forLoopBreak.kt

@@ -0,0 +1,5 @@
+fun some() {
+    <info descr="null">for</info> (i in 0..5) {
+        <info descr="null">break</info>~
+    }
+}

plugins/kotlin/idea/tests/testData/exitPoints/forLoopNestedBreak.kt

@@ -0,0 +1,7 @@
+fun some() {
+    for (j in 0..5) {
+        <info descr="null">for</info> (i in 0..5) {
+            <info descr="null">break</info>~
+        }
+    }
+}

plugins/kotlin/idea/tests/testData/exitPoints/forLoopNestedBreak2.kt

@@ -0,0 +1,12 @@
+fun some() {
+    <info descr="null">for</info>~ (j in 0..5) {
+        if (j == 1) <info descr="null">break</info>
+        for (i in 0..5) {
+            break
+        }
+        loop@ for (i in 0..5) {
+            break
+        }
+        if (j == 2) <info descr="null">continue</info>
+    }
+}

plugins/kotlin/idea/tests/testData/exitPoints/forLoopNestedLabeledBreak.kt

@@ -0,0 +1,13 @@
+fun some(list: List<String>) {
+    a@ <info descr="null">for</info> (j in 0..5) {
+        for (i in 0..5) {
+            <info descr="null">break@a</info>
+        }
+        if (j == 0) <info descr="null">continue</info>
+        if (j == 2) <info descr="null">break</info>
+        b@ for (i in 0..5) {
+            break@b
+            <info descr="null">brea~k@a</info>
+        }
+    }
+}

plugins/kotlin/idea/tests/testData/exitPoints/forLoopNestedLabeledBreakOnLoop.kt

@@ -0,0 +1,12 @@
+fun some(list: List<String>) {
+    loop@ <info descr="null">for</info>~ (s in l) {
+        if (s == "a") <info descr="null">break@loop</info>
+        if (s == "b") <info descr="null">continue</info>
+
+        for (s1 in l) {
+            if (s1 == "a") break
+            if (s1 == "b") continue
+            if (s1 == "c") <info descr="null">break@loop</info>
+        }
+    }
+}

plugins/kotlin/idea/tests/testData/exitPoints/forLoopNestedLabeledBreakOnLoop2.kt

@@ -0,0 +1,34 @@
+fun some(list: List<String>) {
+    loop@ <info descr="null">for</info>~ (s in l) {
+        if (s == "a") <info descr="null">break@loop</info>
+        if (s == "b") <info descr="null">continue</info>
+
+        loop@ for (s1 in l) {
+            if (s1 == "a") break
+            if (s1 == "b") continue
+            if (s1 == "c") break@loop
+        }
+
+        loop2@ for (s1 in l) {
+            if (s1 == "a") break
+            if (s1 == "b") continue
+            if (s1 == "c") <info descr="null">break@loop</info>
+            if (s1 == "d") break
+        }
+
+        loop1@ for (s1 in l) {
+            if (s1 == "a") break
+            if (s1 == "b") continue
+            loop2@ for (s1 in l) {
+                if (s1 == "a") break
+                if (s1 == "b") continue
+                if (s1 == "c") <info descr="null">break@loop</info>
+            }
+            if (s1 == "c") break
+        }
+
+        if (s == "c") <info descr="null">continue</info>
+        if (s == "d") <info descr="null">break</info>
+        if (s == "e") <info descr="null">break@loop</info>
+    }
+}