[kotlin] KTIJ-36879 Added sequence exit points highlighting & tests

GitOrigin-RevId: 136415abd5a5bd27dbf5397f1d6e2cd45f5bc292
This commit is contained in:
Piotr Kasprowicz
2026-01-13 12:33:58 +00:00
committed by intellij-monorepo-bot
parent 8541a8d63c
commit 901cea64df
14 changed files with 384 additions and 32 deletions
@@ -61,6 +61,10 @@ object StandardKotlinNames {
@JvmField val Sequence: FqName = BASE_SEQUENCES_PACKAGE + "Sequence"
@JvmField val sequence: FqName = BASE_SEQUENCES_PACKAGE + "sequence"
@JvmField val yield: FqName = BASE_SEQUENCES_PACKAGE + "SequenceScope.yield"
@JvmField val yieldAll: FqName = BASE_SEQUENCES_PACKAGE + "SequenceScope.yieldAll"
@JvmField val terminations: List<FqName> =
collectionTerminationFunctionNames.map { BASE_SEQUENCES_PACKAGE + it }
@@ -1,4 +1,4 @@
// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package org.jetbrains.kotlin.idea.k2.highlighting;
@@ -305,6 +305,56 @@ public class K2HighlightExitPointsTestGenerated extends AbstractK2HighlightExitP
runTest("../../idea/tests/testData/exitPoints/returnWhenStatement.kt");
}
@TestMetadata("sequenceInLambda.kt")
public void testSequenceInLambda() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceInLambda.kt");
}
@TestMetadata("sequenceNested.kt")
public void testSequenceNested() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceNested.kt");
}
@TestMetadata("sequenceNested2.kt")
public void testSequenceNested2() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceNested2.kt");
}
@TestMetadata("sequenceNestedMixedSyntax.kt")
public void testSequenceNestedMixedSyntax() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceNestedMixedSyntax.kt");
}
@TestMetadata("sequenceNestedParenthesized.kt")
public void testSequenceNestedParenthesized() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceNestedParenthesized.kt");
}
@TestMetadata("sequenceSimple.kt")
public void testSequenceSimple() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceSimple.kt");
}
@TestMetadata("sequenceTryCatch.kt")
public void testSequenceTryCatch() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceTryCatch.kt");
}
@TestMetadata("sequenceWithControlFlow.kt")
public void testSequenceWithControlFlow() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceWithControlFlow.kt");
}
@TestMetadata("sequenceWithLambdaContinue.kt")
public void testSequenceWithLambdaContinue() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceWithLambdaContinue.kt");
}
@TestMetadata("sequenceWithReturn.kt")
public void testSequenceWithReturn() throws Exception {
runTest("../../idea/tests/testData/exitPoints/sequenceWithReturn.kt");
}
@TestMetadata("simple.kt")
public void testSimple() throws Exception {
runTest("../../idea/tests/testData/exitPoints/simple.kt");
@@ -17,8 +17,12 @@ import com.intellij.psi.search.LocalSearchScope
import com.intellij.psi.util.PsiTreeUtil
import com.intellij.psi.util.elementType
import com.intellij.util.Consumer
import org.jetbrains.kotlin.analysis.api.analyze
import org.jetbrains.kotlin.analysis.api.resolution.successfulFunctionCallOrNull
import org.jetbrains.kotlin.idea.codeinsight.utils.StandardKotlinNames
import org.jetbrains.kotlin.idea.codeinsight.utils.doesBelongToLoop
import org.jetbrains.kotlin.idea.codeinsight.utils.findRelevantLoopForExpression
import org.jetbrains.kotlin.name.FqName
import org.jetbrains.kotlin.idea.references.unwrappedTargets
import org.jetbrains.kotlin.lexer.KtTokens
import org.jetbrains.kotlin.psi.*
@@ -32,15 +36,14 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
is KtPropertyAccessor -> parent
is KtReturnExpression, is KtThrowExpression -> parent
is KtFunctionLiteral -> parent.takeIf { with((target as? ASTNode)?.elementType) { this == KtTokens.LBRACE || this == KtTokens.RBRACE } }
is KtLabelReferenceExpression ->
PsiTreeUtil.getParentOfType(
target, KtReturnExpression::class.java, KtThrowExpression::class.java, KtFunction::class.java
)?.takeUnless {
it is KtFunction
}
is KtLabelReferenceExpression -> PsiTreeUtil.getParentOfType(
target, KtReturnExpression::class.java, KtThrowExpression::class.java, KtFunction::class.java
)?.takeUnless {
it is KtFunction
}
else -> null
} as? KtExpression ?: return null
} ?: return null
return OnExitUsagesHandler(editor, file, null, expression, false)
}
@@ -54,9 +57,22 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
return OnLoopUsagesHandler(editor, file, expression)
}
private fun getOnSequenceUsageHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
val expression = when (val parent = target.parent) {
is KtNameReferenceExpression -> parent.takeIf {
target.elementType == KtTokens.IDENTIFIER && parent.text in SEQUENCE_KEYWORDS
}?.let { it.parent as? KtCallExpression }
else -> null
} as? KtExpression ?: return null
return OnSequenceUsagesHandler(editor, file, expression)
}
override fun createHighlightUsagesHandler(editor: Editor, file: PsiFile, target: PsiElement): HighlightUsagesHandlerBase<*>? {
return getOnReturnOrThrowOrLambdaUsageHandler(editor, file, target)
?: getOnBreakOrContinueUsageHandler(editor, file, target)
?: getOnSequenceUsageHandler(editor, file, target)
}
protected abstract fun getRelevantReturnDeclaration(returnExpression: KtReturnExpression): KtDeclarationWithBody?
@@ -95,8 +111,7 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
val referenceExpression: KtNameReferenceExpression?,
val target: KtExpression,
val highlightReferences: Boolean
) :
HighlightUsagesHandlerBase<PsiElement>(editor, file) {
) : HighlightUsagesHandlerBase<PsiElement>(editor, file) {
override fun getTargets(): List<KtExpression> = listOf(target)
@@ -105,13 +120,12 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
}
override fun computeUsages(targets: List<PsiElement>) {
val relevantFunction: KtDeclarationWithBody? =
when (target) {
is KtFunctionLiteral -> target
is KtPropertyAccessor -> target
is KtNamedFunction -> target
else -> getRelevantDeclaration(target)
}
val relevantFunction: KtDeclarationWithBody? = when (target) {
is KtFunctionLiteral -> target
is KtPropertyAccessor -> target
is KtNamedFunction -> target
else -> getRelevantDeclaration(target)
}
var targetOccurrenceAdded = false
if (target is KtReturnExpression || target is KtThrowExpression || target is KtNamedFunction || target is KtPropertyAccessor) {
@@ -135,11 +149,9 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
}
val lastStatementExpressions =
if ((relevantFunction is KtFunctionLiteral && hasNonUnitReturnType(relevantFunction)) ||
(relevantFunction is KtNamedFunction && relevantFunction.bodyBlockExpression == null)
) {
if ((relevantFunction is KtFunctionLiteral && hasNonUnitReturnType(relevantFunction)) || (relevantFunction is KtNamedFunction && relevantFunction.bodyBlockExpression == null)) {
val lastStatements = mutableSetOf<PsiElement>(relevantFunction)
relevantFunction.acceptChildren(object : KtVisitorVoid(),PsiRecursiveVisitor {
relevantFunction.acceptChildren(object : KtVisitorVoid(), PsiRecursiveVisitor {
override fun visitKtElement(element: KtElement) {
ProgressIndicatorProvider.checkCanceled()
element.acceptChildren(this)
@@ -202,7 +214,7 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
emptySet()
}
relevantFunction?.accept(object : KtVisitorVoid(),PsiRecursiveVisitor {
relevantFunction?.accept(object : KtVisitorVoid(), PsiRecursiveVisitor {
override fun visitKtElement(element: KtElement) {
ProgressIndicatorProvider.checkCanceled()
element.acceptChildren(this)
@@ -258,7 +270,7 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
when (expression.returnedExpression) {
is KtIfExpression, is KtWhenExpression, is KtTryExpression -> {
addOccurrence(expression.returnKeyword)
expression.acceptChildren(object : KtVisitorVoid(),PsiRecursiveVisitor {
expression.acceptChildren(object : KtVisitorVoid(), PsiRecursiveVisitor {
override fun visitKtElement(element: KtElement) {
ProgressIndicatorProvider.checkCanceled()
element.acceptChildren(this)
@@ -279,10 +291,9 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
}
}
is KtWhenExpression ->
expression.entries.forEach { whenEntry ->
whenEntry.expression?.let { visitExpression(it) }
}
is KtWhenExpression -> expression.entries.forEach { whenEntry ->
whenEntry.expression?.let { visitExpression(it) }
}
else -> addOccurrence(expression)
}
@@ -309,9 +320,11 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
(target as? PsiNameIdentifierOwner)?.nameIdentifier.takeIf { target.containingFile == containingFile }
?.let(::addOccurrence)
val handler: FindUsagesHandler? = (FindManager.getInstance(relevantFunction.project) as FindManagerImpl)
.findUsagesManager
.getFindUsagesHandler(target, true)
val handler: FindUsagesHandler? =
(FindManager.getInstance(relevantFunction.project) as FindManagerImpl).findUsagesManager.getFindUsagesHandler(
target,
true
)
handler?.findReferencesToHighlight(target, LocalSearchScope(containingFile)).let { ref ->
ref?.forEach { addOccurrence(it.element) }
}
@@ -322,7 +335,7 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
override fun highlightReferences(): Boolean = highlightReferences
}
private inner class OnLoopUsagesHandler(editor: Editor, file: PsiFile, val target: KtExpression) :
private class OnLoopUsagesHandler(editor: Editor, file: PsiFile, val target: KtExpression) :
HighlightUsagesHandlerBase<PsiElement>(editor, file) {
override fun getTargets(): List<KtExpression> = listOf(target)
@@ -359,6 +372,7 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
if (loopLabelName == null) return
}
}
is KtBreakExpression, is KtContinueExpression -> {
val expressionLabelName = (expression as? KtExpressionWithLabel)?.getLabelName()
if (expressionLabelName != null && expressionLabelName == loopLabelName) {
@@ -378,6 +392,94 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
}
private inner class OnSequenceUsagesHandler(editor: Editor, file: PsiFile, val target: KtExpression) :
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: List<PsiElement>) {
val sequenceCall = findSequenceCall(target) ?: return
// Handles both trailing lambda syntax: sequence { } and parenthesized: sequence({ })
val sequenceLambda =
sequenceCall.lambdaArguments.firstOrNull()?.getLambdaExpression() ?: (sequenceCall.valueArguments.firstOrNull()
?.getArgumentExpression() as? KtLambdaExpression) ?: return
// Highlight "sequence" keyword
sequenceCall.calleeExpression?.let { addOccurrence(it) }
// Find all yield/yieldAll for current sequence lambda
sequenceLambda.accept(object : KtVisitorVoid(), PsiRecursiveVisitor {
override fun visitKtElement(element: KtElement) {
element.acceptChildren(this)
}
override fun visitCallExpression(expression: KtCallExpression) {
super.visitCallExpression(expression)
// Check if this call belongs to our sequence (not nested)
if (belongsToSequence(expression, sequenceLambda) && isYieldCall(expression)) {
addOccurrence(expression)
}
}
})
}
}
private fun findSequenceCall(expression: KtExpression): KtCallExpression? {
// If clicked the "sequence"
if (expression is KtCallExpression && isSequenceBuilderCall(expression)) {
return expression
}
// Clicked yield/yieldAll -> Find enclosing sequence call
return expression.parents.filterIsInstance<KtLambdaExpression>().firstNotNullOfOrNull { lambdaExpr ->
val lambdaArg = lambdaExpr.parent
val call = when (lambdaArg) {
is KtLambdaArgument -> lambdaArg.parent as? KtCallExpression
is KtValueArgument -> (lambdaArg.parent as? KtValueArgumentList)?.parent as? KtCallExpression
else -> null
}
call?.takeIf { isSequenceBuilderCall(it) }
}
}
private fun belongsToSequence(yieldCall: KtCallExpression, targetSequenceLambda: KtLambdaExpression): Boolean {
// Check if yield doesn't belong to a nested sequence
for (parent in yieldCall.parents) {
if (parent == targetSequenceLambda) return true
if (parent is KtLambdaExpression && parent != targetSequenceLambda) {
val call = when (val lambdaParent = parent.parent) {
is KtLambdaArgument -> lambdaParent.parent as? KtCallExpression // sequence { }
is KtValueArgument -> (lambdaParent.parent as? KtValueArgumentList)?.parent as? KtCallExpression // sequence({ })
else -> null
}
if (call != null && isSequenceBuilderCall(call)) {
// a new sequence
return false
}
}
}
return false
}
private fun isSequenceBuilderCall(call: KtCallExpression): Boolean {
if (call.calleeExpression?.text != "sequence") return false
return call.resolvesToFqName(StandardKotlinNames.Sequences.sequence)
}
private fun isYieldCall(call: KtCallExpression): Boolean {
val calleeName = call.calleeExpression?.text
if (calleeName != "yield" && calleeName != "yieldAll") return false
return call.resolvesToFqName(
StandardKotlinNames.Sequences.yield, StandardKotlinNames.Sequences.yieldAll
)
}
private fun MutableSet<PsiElement>.addIfNotNullAndNotBlock(element: PsiElement?) {
addIfNotNull(element.takeUnless { it is KtBlockExpression })
}
@@ -392,4 +494,19 @@ abstract class AbstractKotlinHighlightExitPointsHandlerFactory : HighlightUsages
return expression
}
private fun KtCallExpression.resolvesToFqName(vararg expectedFqNames: FqName): Boolean {
analyze(this) {
val resolvedCall = resolveToCall()?.successfulFunctionCallOrNull() ?: return false
val symbol = resolvedCall.partiallyAppliedSymbol.signature.symbol
val callableId = symbol.callableId ?: return false
val fqName = callableId.asSingleFqName()
return fqName in expectedFqNames
}
}
companion object {
private val SEQUENCE_KEYWORDS = setOf("yield", "yieldAll", "sequence")
}
}
@@ -1,4 +1,4 @@
// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
// Copyright 2000-2026 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package org.jetbrains.kotlin.idea.highlighter;
@@ -305,6 +305,56 @@ public class HighlightExitPointsTestGenerated extends AbstractCustomHighlightUsa
runTest("testData/exitPoints/returnWhenStatement.kt");
}
@TestMetadata("sequenceInLambda.kt")
public void testSequenceInLambda() throws Exception {
runTest("testData/exitPoints/sequenceInLambda.kt");
}
@TestMetadata("sequenceNested.kt")
public void testSequenceNested() throws Exception {
runTest("testData/exitPoints/sequenceNested.kt");
}
@TestMetadata("sequenceNested2.kt")
public void testSequenceNested2() throws Exception {
runTest("testData/exitPoints/sequenceNested2.kt");
}
@TestMetadata("sequenceNestedMixedSyntax.kt")
public void testSequenceNestedMixedSyntax() throws Exception {
runTest("testData/exitPoints/sequenceNestedMixedSyntax.kt");
}
@TestMetadata("sequenceNestedParenthesized.kt")
public void testSequenceNestedParenthesized() throws Exception {
runTest("testData/exitPoints/sequenceNestedParenthesized.kt");
}
@TestMetadata("sequenceSimple.kt")
public void testSequenceSimple() throws Exception {
runTest("testData/exitPoints/sequenceSimple.kt");
}
@TestMetadata("sequenceTryCatch.kt")
public void testSequenceTryCatch() throws Exception {
runTest("testData/exitPoints/sequenceTryCatch.kt");
}
@TestMetadata("sequenceWithControlFlow.kt")
public void testSequenceWithControlFlow() throws Exception {
runTest("testData/exitPoints/sequenceWithControlFlow.kt");
}
@TestMetadata("sequenceWithLambdaContinue.kt")
public void testSequenceWithLambdaContinue() throws Exception {
runTest("testData/exitPoints/sequenceWithLambdaContinue.kt");
}
@TestMetadata("sequenceWithReturn.kt")
public void testSequenceWithReturn() throws Exception {
runTest("testData/exitPoints/sequenceWithReturn.kt");
}
@TestMetadata("simple.kt")
public void testSimple() throws Exception {
runTest("testData/exitPoints/simple.kt");
@@ -0,0 +1,9 @@
// WITH_STDLIB
fun f() {
listOf(1, 2, 3).map { n ->
<info descr="null">sequence</info> { // Highlighted
<info descr="null">~yield(n)</info> // Highlighted
<info descr="null">yield(n * 2)</info> // Highlighted
}
}
}
@@ -0,0 +1,14 @@
// WITH_STDLIB
fun f() {
val outer = <info descr="null">~sequence</info> { // Caret
<info descr="null">yield(1)</info> // Highlighted
val inner = sequence { // Nested sequence
yield(99) // NOT highlighted
yield(100) // NOT highlighted
}
<info descr="null">yield(2)</info> // Highlighted
<info descr="null">yieldAll(inner)</info> // Highlighted
}
}
@@ -0,0 +1,21 @@
// WITH_STDLIB
fun f() {
val nested0 = sequence { // Level 1
yield(1)
val nested1 = <info descr="null">sequence</info> { // Level 2
<info descr="null">~yield(10)</info>
val nested2 = sequence { // Level 3
yield(100)
yield(200)
}
<info descr="null">yieldAll(nested2)</info>
<info descr="null">yield(20)</info>
}
yieldAll(nested1)
yield(2)
}
}
@@ -0,0 +1,12 @@
// WITH_STDLIB
fun f() {
val outer = sequence { // Trailing lambda
yield(1)
val inner = <info descr="null">sequence</info>({ // Parenthesized
<info descr="null">~yield(99)</info>
})
yield(2)
}
}
@@ -0,0 +1,12 @@
// WITH_STDLIB
fun f() {
val outer = <info descr="null">sequence</info>({ // Parenthesized outer
<info descr="null">yield(1)</info>
val inner = sequence({ // Parenthesized inner
yield(99)
})
<info descr="null">~yield(2)</info>
})
}
@@ -0,0 +1,8 @@
// WITH_STDLIB
fun f() {
val seq = <info descr="null">~sequence</info> { // Caret
<info descr="null">yield(1)</info> // Highlighted
<info descr="null">yield(2)</info> // Highlighted
<info descr="null">yieldAll(listOf(3, 4))</info> // Highlighted
}
}
@@ -0,0 +1,12 @@
// WITH_STDLIB
fun f() {
val seq = <info descr="null">~sequence</info> { // Caret
try {
<info descr="null">yield(1)</info> // Highlighted
} catch (e: Exception) {
<info descr="null">yield(-1)</info> // Highlighted
} finally {
// yield(0) // not allowed in finally
}
}
}
@@ -0,0 +1,19 @@
// WITH_STDLIB
fun f() {
val seq = <info descr="null">~sequence</info> { // Caret
for (i in 1..5) {
<info descr="null">yield(i)</info> // Highlighted
}
if (true) {
<info descr="null">yield(100)</info> // Highlighted
} else {
<info descr="null">yieldAll(emptyList())</info> // Highlighted
}
when (1) {
1 -> <info descr="null">yield(200)</info> // Highlighted
else -> <info descr="null">yield(300)</info> // Highlighted
}
}
}
@@ -0,0 +1,11 @@
// WITH_STDLIB
fun f(){
val seq = <info descr="null">sequence</info> {
outer@ for (i in 1..5) {
for (j in 1..5){
if (j == 2) continue@outer
<info descr="null">~yield(i * j)</info> // Highlighted
}
}
}
}
@@ -0,0 +1,13 @@
// WITH_STDLIB
fun f(): Sequence<Int> {
return <info descr="null">~sequence</info> {
<info descr="null">yield(1)</info>
if(true) {
return @sequence
}
<info descr="null">yield(2)</info>
}
// return 0
}