From ad4547af51943aa931ffce65ffbd5dfbf62f4484 Mon Sep 17 00:00:00 2001 From: Frederik Haselmeier Date: Wed, 21 Jan 2026 13:03:43 +0100 Subject: [PATCH] [kotlin] KTIJ-36991 added tail handling to K2 smart completion ^KTIJ-36991 fixed ^KTIJ-37243 fixed Space-RevId: f92fa35d7c43584c92abd456bab9a30d3cb35959 GitOrigin-RevId: d29bd4602da1bc69fd205794f08b68c6bd5e4888 --- .../impl/k2/K2CompletionContributor.kt | 11 +- .../k2/handlers/SmartCompletionTailHandler.kt | 262 ++++++++++++++++++ 2 files changed, 272 insertions(+), 1 deletion(-) create mode 100644 plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/handlers/SmartCompletionTailHandler.kt diff --git a/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/K2CompletionContributor.kt b/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/K2CompletionContributor.kt index 7bc3c9dee97e..086b3e30c2ed 100644 --- a/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/K2CompletionContributor.kt +++ b/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/K2CompletionContributor.kt @@ -1,6 +1,7 @@ // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package org.jetbrains.kotlin.idea.completion.impl.k2 +import com.intellij.codeInsight.completion.CompletionType import com.intellij.codeInsight.completion.PrefixMatcher import com.intellij.codeInsight.lookup.LookupElement import com.intellij.codeInsight.lookup.LookupElementDecorator @@ -19,6 +20,7 @@ import org.jetbrains.kotlin.idea.completion.KotlinFirCompletionParameters import org.jetbrains.kotlin.idea.completion.checkers.CompletionVisibilityChecker import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.evaluateRuntimeKaType import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.WrapSingleStringTemplateEntryWithBracesInsertHandler +import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.addSmartCompletionTailInsertHandler import org.jetbrains.kotlin.idea.completion.implCommon.handlers.CompletionCharInsertHandler import org.jetbrains.kotlin.idea.completion.implCommon.stringTemplates.InsertStringTemplateBracesInsertHandler import org.jetbrains.kotlin.idea.completion.isAtFunctionLiteralStart @@ -290,10 +292,17 @@ internal abstract class K2CompletionContributor

( else -> WrapSingleStringTemplateEntryWithBracesInsertHandler } - return LookupElementDecorator.withDelegateInsertHandler( + var element = element + if (context.completionContext.parameters.completionType == CompletionType.SMART) { + element = element.addSmartCompletionTailInsertHandler() + } + + element = LookupElementDecorator.withDelegateInsertHandler( LookupElementDecorator.withDelegateInsertHandler(element, bracesInsertHandler), CompletionCharInsertHandler(context.parameters.delegate.isAutoPopup), ) + + return element } } diff --git a/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/handlers/SmartCompletionTailHandler.kt b/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/handlers/SmartCompletionTailHandler.kt new file mode 100644 index 000000000000..7cf62a880351 --- /dev/null +++ b/plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/handlers/SmartCompletionTailHandler.kt @@ -0,0 +1,262 @@ +// 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.completion.impl.k2.handlers + +import com.intellij.codeInsight.lookup.LookupElement +import com.intellij.codeInsight.lookup.LookupElementDecorator +import com.intellij.psi.PsiElement +import org.jetbrains.kotlin.analysis.api.KaSession +import org.jetbrains.kotlin.analysis.api.components.isFunctionType +import org.jetbrains.kotlin.analysis.api.components.resolveToCallCandidates +import org.jetbrains.kotlin.analysis.api.resolution.KaCall +import org.jetbrains.kotlin.analysis.api.resolution.KaFunctionCall +import org.jetbrains.kotlin.analysis.api.signatures.KaVariableSignature +import org.jetbrains.kotlin.analysis.api.symbols.KaValueParameterSymbol +import org.jetbrains.kotlin.idea.completion.handlers.WithTailInsertHandler +import org.jetbrains.kotlin.idea.completion.impl.k2.K2CompletionSectionContext +import org.jetbrains.kotlin.idea.util.positionContext.KotlinExpressionNameReferencePositionContext +import org.jetbrains.kotlin.name.Name +import org.jetbrains.kotlin.psi.* +import org.jetbrains.kotlin.psi.psiUtil.getQualifiedExpressionForSelectorOrThis + +/** + * Tails are the tail of what should be inserted by smart completion after choosing a completion item to + * move the user's caret to the place where they will likely need to invoke completion next. + * In some cases, the tail is already in the code, in which case rather than inserting it, + * smart completion moves past the existing tail and places the caret after it. + * Otherwise, the tail is inserted and the caret is moved to be after the tail. + * + * For example, in the following code: + * ``` + * fun foo(a: Int, b: Int) {} + * fun bar(a: Int, b: Int) { + * foo() + * } + * ``` + * Smart completion knows that after completing the first parameter for `foo`, another argument will be + * required so the tail will be `, ` and the resulting code will be: + * ``` + * fun foo(a: Int, b: Int) {} + * fun bar(a: Int, b: Int) { + * foo(a, ) + * } + * ``` + */ +private enum class Tail { + COMMA, // Additional argument to call required + RPARENTH, // Call needs no more arguments, move past the closing parenthesis + RBRACKET, // Similar as RPARENTH but for array access + ELSE, // When an if expression is completed and an `else` branch is required + RBRACE // For lambda parameters +} + +/** + * Returns the tail for a named argument + */ +private fun namedArgumentTail( + argumentName: Name, + call: KaFunctionCall<*>, +): Tail? { + val usedParameterNames = (call.argumentMapping.values.map { it.name } + listOf(argumentName)).toSet() + val notUsedParameters = call.signature.valueParameters.filter { it.name !in usedParameterNames } + return when { + notUsedParameters.isEmpty() -> Tail.RPARENTH // named arguments no supported for [] + notUsedParameters.all { it.symbol.hasDefaultValue } -> null + else -> Tail.COMMA + } +} + +/** + * Returns the tail for this [call] belonging to the [argumentExpression]. + * Returns a single tail if a tail could be determined, or null otherwise. + */ +context(_: KaSession) +private fun calculateTailForCall(argumentExpression: KtExpression, call: KaCall): Tail? { + if (call !is KaFunctionCall<*>) return null + if (call.argumentMapping.isEmpty()) return null + + val argument = argumentExpression.parent ?: return null + val argumentName = (argument as? KtValueArgument)?.getArgumentName()?.asName + val isFunctionLiteralArgument = argument is LambdaArgument + + var parameter = call.argumentMapping[argumentExpression] + + val isArrayAccess = argument.parent is KtArrayAccessExpression + val rparenthTail = if (isArrayAccess) Tail.RBRACKET else Tail.RPARENTH + + var parameters = call.signature.valueParameters.toList() + if (isArrayAccess && call.argumentMapping.size == 2) { + // last parameter in set is used for value assigned + if (parameter == parameters.last()) { + parameter = null + } + parameters = parameters.dropLast(1) + } + + if (parameter == null) { + return null + } + + fun needCommaForParameter(parameter: KaVariableSignature): Boolean { + if (parameter.symbol.hasDefaultValue) return false // parameter is optional + if (parameter.symbol.isVararg) return false // vararg arguments list can be empty + // last parameter of functional type can be placed outside parenthesis: + if (!isArrayAccess && parameter == parameters.last() && parameter.returnType.isFunctionType) return false + return true + } + + val tail = if (argumentName == null) { + when { + parameter == parameters.last() -> rparenthTail + parameters.dropWhile { it != parameter }.drop(1).any { needCommaForParameter(it) } -> Tail.COMMA + else -> null + } + } else { + namedArgumentTail(argumentName, call) + } + + val alreadyHasStar = (argument as? KtValueArgument)?.getSpreadElement() != null + return if (parameter.symbol.isVararg) { + if (isFunctionLiteralArgument) return null + + val varargTail = if (argumentName == null && tail == rparenthTail) { + null /* even if it's the last parameter, there can be more arguments for the same parameter */ + } else tail + + varargTail?.takeIf { !alreadyHasStar } + } else { + if (alreadyHasStar) return null + tail.takeIf { !isFunctionLiteralArgument } + } +} + +/** + * Calculates all tails based on all possible call candidates at the position. + */ +context(_: KaSession, _: K2CompletionSectionContext<*>) +private fun calculateTailForCalls( + expression: KtExpression, + argument: KtElement, + argumentParent: PsiElement, +): Set { + // The way of getting the callCandidates depends on what kind of position we are in + val callCandidates = when { + argument is KtValueArgument -> { + val argumentList = argument.parent as? KtValueArgumentList ?: return emptySet() + val call = argumentList.parent as? KtCallElement ?: return emptySet() + call.resolveToCallCandidates() + } + + argumentParent is KtArrayAccessExpression -> { + argumentParent.resolveToCallCandidates() + } + + else -> emptyList() + } + + if (callCandidates.isEmpty()) return emptySet() + + return callCandidates.mapNotNullTo(mutableSetOf()) { + calculateTailForCall(expression, it.candidate) + } +} + +context(_: KaSession, context: K2CompletionSectionContext<*>) +private fun calculateTailForIfExpression( + argumentExpression: KtExpression, + ifExpression: KtIfExpression +): Set { + return when (argumentExpression) { + // If we are inside the condition, we need to close the condition's parenthesis. + ifExpression.condition -> setOf(Tail.RPARENTH) + // If we are done with the `then` block, add `else` and move into it + ifExpression.then -> setOf(Tail.ELSE) + // If we completed the `else`, the `if` is effectively done and we use the tail that should + // be used for the surrounding if expression. + ifExpression.`else` -> ifExpression.calculateTails() + else -> emptySet() + } +} + +context(_: KaSession, context: K2CompletionSectionContext<*>) +private fun calculateTailForBinaryExpression( + argumentExpression: KtExpression, + binaryExpression: KtBinaryExpression, +): Set { + if (argumentExpression != binaryExpression.right) return emptySet() + // If we completed the binary expression, we are done with it, and we use the + // tail from the surrounding expression to move to the correct place. + return binaryExpression.calculateTails() +} + +context(_: KaSession, _: K2CompletionSectionContext<*>) +private fun calculateTailForBlockExpression( + blockExpression: KtBlockExpression, +): Set { + if (blockExpression.parent !is KtFunctionLiteral) return emptySet() + // We completed the last expression inside a function literal, which + // means we are done with the function literal and can move past it. + return setOf(Tail.RBRACE) +} + + +/** + * Calculates the expected tail to be used in smart completion for the given expression. + * See [Tail] for details. + * + * The majority of the tail logic has been ported from K1 code from: + * `org.jetbrains.kotlin.idea.core.ExpectedInfos` + */ +context(_: KaSession, context: K2CompletionSectionContext<*>) +private fun KtExpression.calculateTails(): Set { + val argumentExpression = getQualifiedExpressionForSelectorOrThis() + val argument = argumentExpression.parent as? KtElement ?: return emptySet() + val argumentParent = argument.parent ?: return emptySet() + + return when { + argumentParent is KtIfExpression -> { + calculateTailForIfExpression(argumentExpression, argumentParent) + } + + argument is KtBinaryExpression -> { + calculateTailForBinaryExpression(argumentExpression, argument) + } + + argument is KtBlockExpression -> { + calculateTailForBlockExpression(argument) + } + + else -> { + calculateTailForCalls(argumentExpression, argument, argumentParent) + } + } +} + +private fun Tail.getInsertHandler(): WithTailInsertHandler = when (this) { + Tail.COMMA -> WithTailInsertHandler.COMMA + Tail.RPARENTH -> WithTailInsertHandler.RPARENTH + Tail.RBRACKET -> WithTailInsertHandler.RBRACKET + Tail.ELSE -> WithTailInsertHandler.ELSE + Tail.RBRACE -> WithTailInsertHandler.RBRACE +} + +/** + * Adds an insertion handler to the element that inserts or moves past (if it already exists) + * a tail depending on the calls available at the position. + * + * See [Tail] for more detailed documentation. + */ +context(_: KaSession, context: K2CompletionSectionContext<*>) +internal fun LookupElement.addSmartCompletionTailInsertHandler(): LookupElement { + val positionContext = context.positionContext + if (positionContext !is KotlinExpressionNameReferencePositionContext) return this + + // There might be multiple tails depending on all the available signatures. + // If all call signatures require the same tail (i.e., the set is a singleton), we can use this single tail. + val singleTail = positionContext.nameExpression.calculateTails().singleOrNull() + + return if (singleTail != null) { + LookupElementDecorator.withDelegateInsertHandler(this, singleTail.getInsertHandler()) + } else { + this + } +} \ No newline at end of file