[kotlin] KTIJ-37603 added replacement offset handling for smart completion

This is also required for KTIJ-35604 for replacing multiple arguments correctly

^KTIJ-37603 fixed

GitOrigin-RevId: f7b553d5a4f9af3db5c41365b9b973c7aec56ab8

plugins/kotlin/completion/api/src/org/jetbrains/kotlin/idea/completion/api/serialization/lookup/LookupModelConverter.kt

@@ -129,6 +129,7 @@ object LookupModelConverter {
                 put(
                     keyName,
                     when (value) {
+                        is Unit -> UserDataValueModel.UnitModel
                         is Boolean -> UserDataValueModel.BooleanModel(value)
                         is String -> UserDataValueModel.StringModel(value)
                         is Enum<*> -> UserDataValueModel.EnumModel(value.ordinal, value::class.java.name)
@@ -170,6 +171,7 @@ object LookupModelConverter {
         for ((k, v) in data) {
             val key = Key.findKeyByName(k) as Key<Any>
             val value = when (v) {
+                is UserDataValueModel.UnitModel -> Unit
                 is UserDataValueModel.BooleanModel -> v.value
                 is UserDataValueModel.StringModel -> v.value
                 is UserDataValueModel.EnumModel -> Class.forName(v.enumClass).enumConstants[v.ordinal]

plugins/kotlin/completion/api/src/org/jetbrains/kotlin/idea/completion/api/serialization/lookup/model/UserDataValueModel.kt

@@ -12,6 +12,9 @@ import org.jetbrains.kotlin.name.Name
 @Serializable
 @ApiStatus.Internal
 sealed class UserDataValueModel {
+    @Serializable
+    data object UnitModel : UserDataValueModel()
+
     @Serializable
     data class BooleanModel(val value: Boolean) : UserDataValueModel()
 

plugins/kotlin/completion/impl-k2/resources/intellij.kotlin.completion.impl.xml

@@ -19,6 +19,10 @@
     <applicationService serviceImplementation="org.jetbrains.kotlin.idea.completion.impl.k2.FirCompletionDummyIdentifierProviderService"
                         serviceInterface="org.jetbrains.kotlin.idea.completion.api.CompletionDummyIdentifierProviderService"/>
 
+    <applicationService
+            serviceInterface="org.jetbrains.kotlin.idea.completion.handlers.SmartCompletionTailOffsetProvider"
+            serviceImplementation="org.jetbrains.kotlin.idea.completion.impl.k2.handlers.K2SmartCompletionTailOffsetProviderImpl"/>
+
     <registryKey key="kotlin.k2.chain.completion.enabled"
                  description="Allow chained completion items in Kotlin"
                  defaultValue="true"

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/K2CompletionContributor.kt

@@ -18,8 +18,12 @@ import org.jetbrains.kotlin.idea.base.analysis.api.utils.KtSymbolFromIndexProvid
 import org.jetbrains.kotlin.idea.base.codeInsight.contributorClass
 import org.jetbrains.kotlin.idea.completion.impl.k2.checkers.CompletionVisibilityChecker
 import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.evaluateRuntimeKaType
+import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.KEEP_OLD_ARGUMENT_LIST_ON_TAB_KEY
+import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.SmartCompletionReplaceExistingArgumentHandler
 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.impl.k2.handlers.keepOldArgumentListOnTab
+import org.jetbrains.kotlin.idea.completion.impl.k2.lookups.factories.FunctionCallLookupObject
 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
@@ -286,13 +290,30 @@ internal abstract class K2CompletionContributor<P : KotlinRawPositionContext>(
             element.suppressItemSelectionByCharsOnTyping = true
         }
 
+        val lookupObject = element.`object`
+        if (lookupObject is FunctionCallLookupObject && lookupObject.inputValueArgumentsAreRequired) {
+            // This is required to not replace arguments already passed to existing function calls,
+            // if the replacing function also requires value arguments
+            element.keepOldArgumentListOnTab()
+        }
+
         val bracesInsertHandler = when (context.parameters.type) {
             KotlinFirCompletionParameters.CorrectionType.BRACES_FOR_STRING_TEMPLATE -> InsertStringTemplateBracesInsertHandler
             else -> WrapSingleStringTemplateEntryWithBracesInsertHandler
         }
 
         var element = element
+
         if (context.completionContext.parameters.completionType == CompletionType.SMART) {
+            if (element.getUserData(KEEP_OLD_ARGUMENT_LIST_ON_TAB_KEY) == null) {
+                // In smart completion we want to replace entire arguments (e.g., including their parentheses)
+                // if replacement completion is used
+                element = LookupElementDecorator.withDelegateInsertHandler(
+                    element,
+                    SmartCompletionReplaceExistingArgumentHandler()
+                )
+            }
+
             element = element.addSmartCompletionTailInsertHandler()
         }
 

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/KotlinFirCompletionContributor.kt

@@ -21,6 +21,7 @@ import com.intellij.util.ProcessingContext
 import com.intellij.util.applyIf
 import org.jetbrains.kotlin.idea.completion.KDocTagCompletionProvider
 import org.jetbrains.kotlin.idea.completion.api.CompletionDummyIdentifierProviderService
+import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.K2SmartCompletionTailOffsetProviderImpl
 import org.jetbrains.kotlin.idea.completion.impl.k2.jfr.CompletionEvent
 import org.jetbrains.kotlin.idea.completion.impl.k2.jfr.CompletionSetupEvent
 import org.jetbrains.kotlin.idea.completion.impl.k2.jfr.timeEvent
@@ -85,6 +86,9 @@ class KotlinFirCompletionContributor : CompletionContributor() {
         context.dummyIdentifier = identifierProviderService.provideDummyIdentifier(context)
 
         identifierProviderService.correctPositionForParameter(context)
+
+        // Mark replacement offsets for (smart) completion
+        K2SmartCompletionTailOffsetProviderImpl.calculateReplacementOffsets(context)
     }
 }
 

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/SerializableInsertionHandlerSerializersModule.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 kotlinx.serialization.builtins.serializer
 import kotlinx.serialization.modules.SerializersModule
 import kotlinx.serialization.modules.polymorphic
 import org.jetbrains.annotations.ApiStatus
@@ -16,6 +17,7 @@ import org.jetbrains.kotlin.idea.completion.handlers.LineAdjusterInsertionHandle
 import org.jetbrains.kotlin.idea.completion.handlers.WithTailInsertHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.ChainedInsertHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.CompoundInsertionHandler
+import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.MultipleArgumentsInsertHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.SuperLookupObject
 import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.TailTextInsertHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.TypeParameterInWhenClauseILookupObject
@@ -28,6 +30,7 @@ import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.AnonymousObjectInse
 import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.BracketOperatorInsertionHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.InsertRequiredTypeArgumentsInsertHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.QualifyContextSensitiveResolutionHandler
+import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.SmartCompletionReplaceExistingArgumentHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.TrailingLambdaInsertionHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.WithImportInsertionHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.WrapSingleStringTemplateEntryWithBracesInsertHandler
@@ -114,5 +117,7 @@ val serializableInsertionHandlerSerializersModule: SerializersModule = Serialize
         subclass(TrailingLambdaInsertionHandler::class, serializer = TrailingLambdaInsertionHandler.serializer())
         subclass(InsertRequiredTypeArgumentsInsertHandler::class, serializer = InsertRequiredTypeArgumentsInsertHandler.serializer())
         subclass(QualifyContextSensitiveResolutionHandler::class, serializer = QualifyContextSensitiveResolutionHandler.serializer())
+        subclass(MultipleArgumentsInsertHandler::class, serializer = MultipleArgumentsInsertHandler.serializer())
+        subclass(SmartCompletionReplaceExistingArgumentHandler::class, serializer = SmartCompletionReplaceExistingArgumentHandler.serializer())
     }
 }

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/contributors/K2MultipleArgumentContributor.kt

@@ -1,6 +1,10 @@
 // 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.contributors
 
+import com.intellij.codeInsight.completion.InsertionContext
+import com.intellij.codeInsight.lookup.Lookup
+import com.intellij.codeInsight.lookup.LookupElement
+import kotlinx.serialization.Serializable
 import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.analysis.api.components.KaScopeKind
 import org.jetbrains.kotlin.analysis.api.components.resolveToCallCandidates
@@ -10,9 +14,12 @@ import org.jetbrains.kotlin.analysis.api.signatures.KaFunctionSignature
 import org.jetbrains.kotlin.analysis.api.symbols.KaVariableSymbol
 import org.jetbrains.kotlin.analysis.api.types.KaType
 import org.jetbrains.kotlin.idea.base.analysis.api.utils.isPossiblySubTypeOf
+import org.jetbrains.kotlin.idea.completion.api.serialization.SerializableInsertHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.K2CompletionSectionContext
 import org.jetbrains.kotlin.idea.completion.impl.k2.K2SimpleCompletionContributor
+import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.K2SmartCompletionTailOffsetProviderImpl
 import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.Tail
+import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.tryGetOffset
 import org.jetbrains.kotlin.idea.completion.impl.k2.lookups.factories.KotlinFirLookupElementFactory
 import org.jetbrains.kotlin.idea.util.positionContext.KotlinNameReferencePositionContext
 import org.jetbrains.kotlin.name.Name
@@ -189,5 +196,25 @@ internal class K2MultipleArgumentContributor : K2SimpleCompletionContributor<Kot
             addElement(KotlinFirLookupElementFactory.createMultipleArgumentsLookupElement(data.matchingVariables, data.tail))
         }
     }
-
 }
+
+/**
+ * This insertion handler is responsible for replacing the existing arguments
+ * if replacement completion is invoked with a multi-argument completion item.
+ *
+ * e.g.: foo(<caret>a, b) -> foo(c, d)<caret>
+ */
+@Serializable
+internal class MultipleArgumentsInsertHandler : SerializableInsertHandler {
+    override fun handleInsert(
+        context: InsertionContext,
+        item: LookupElement
+    ) {
+        if (context.completionChar != Lookup.REPLACE_SELECT_CHAR) return
+        val offset = context.offsetMap.tryGetOffset(K2SmartCompletionTailOffsetProviderImpl.MULTIPLE_ARGUMENTS_REPLACEMENT_OFFSET)
+        if (offset != null) {
+            context.document.deleteString(context.tailOffset, offset)
+        }
+    }
+
+}

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/handlers/K2SmartCompletionTailOffsetProviderImpl.kt

@@ -0,0 +1,156 @@
+// 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.completion.CompletionInitializationContext
+import com.intellij.codeInsight.completion.CompletionType
+import com.intellij.codeInsight.completion.InsertionContext
+import com.intellij.codeInsight.completion.OffsetKey
+import com.intellij.codeInsight.completion.OffsetMap
+import com.intellij.codeInsight.lookup.Lookup
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.openapi.editor.Document
+import com.intellij.openapi.util.Key
+import kotlinx.serialization.Serializable
+import org.jetbrains.kotlin.idea.completion.api.CompletionDummyIdentifierProviderService
+import org.jetbrains.kotlin.idea.completion.api.serialization.SerializableInsertHandler
+import org.jetbrains.kotlin.idea.completion.handlers.SmartCompletionTailOffsetProvider
+import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.K2SmartCompletionTailOffsetProviderImpl.Companion.OLD_ARGUMENTS_REPLACEMENT_OFFSET
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.KtQualifiedExpression
+import org.jetbrains.kotlin.psi.KtValueArgument
+import org.jetbrains.kotlin.psi.KtValueArgumentList
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import kotlin.math.max
+
+/**
+ * This class and the remaining file is responsible for handling the replacement offsets in smart completion
+ * when a user decides to use the `\t` replacement character.
+ *
+ * The code is ported entirely from the corresponding K1 implementation:
+ * [org.jetbrains.kotlin.idea.completion.handlers.SmartCompletionTailOffsetProviderFE10Impl]
+ */
+internal class K2SmartCompletionTailOffsetProviderImpl : SmartCompletionTailOffsetProvider() {
+
+    companion object {
+        internal val OLD_ARGUMENTS_REPLACEMENT_OFFSET: OffsetKey = OffsetKey.create("nonFunctionReplacementOffset")
+        internal val MULTIPLE_ARGUMENTS_REPLACEMENT_OFFSET: OffsetKey = OffsetKey.create("multipleArgumentsReplacementOffset")
+
+        internal fun calculateReplacementOffsets(context: CompletionInitializationContext) {
+            val offset = context.startOffset
+            val tokenAt = context.file.findElementAt(max(0, offset)) ?: return
+            /* do not use parent expression if we are at the end of line - it's probably parsed incorrectly */
+            if (context.completionType == CompletionType.SMART && !isAtEndOfLine(offset, context.editor.document)) {
+                var parent = tokenAt.parent
+                if (parent is KtExpression && parent !is KtBlockExpression) {
+                    // search expression to be replaced - go up while we are the first child of parent expression
+                    var expression: KtExpression = parent
+                    parent = expression.parent
+                    while (parent is KtExpression && parent.getFirstChild() == expression) {
+                        expression = parent
+                        parent = expression.parent
+                    }
+
+                    val suggestedReplacementOffset = replacementOffsetByExpression(expression)
+                    if (suggestedReplacementOffset > context.replacementOffset) {
+                        context.replacementOffset = suggestedReplacementOffset
+                    }
+
+                    context.offsetMap.addOffset(OLD_ARGUMENTS_REPLACEMENT_OFFSET, expression.endOffset)
+
+                    val argumentList = (expression.parent as? KtValueArgument)?.parent as? KtValueArgumentList
+                    if (argumentList != null) {
+                        context.offsetMap.addOffset(
+                            MULTIPLE_ARGUMENTS_REPLACEMENT_OFFSET,
+                            argumentList.rightParenthesis?.textRange?.startOffset ?: argumentList.endOffset
+                        )
+                    }
+                }
+            }
+            CompletionDummyIdentifierProviderService.getInstance().correctPositionForParameter(context)
+        }
+    }
+
+    override fun getTailOffset(
+        context: InsertionContext,
+        item: LookupElement
+    ): Int {
+        val completionChar = context.completionChar
+        var tailOffset = context.tailOffset
+        if (completionChar == Lookup.REPLACE_SELECT_CHAR && item.getUserData(KEEP_OLD_ARGUMENT_LIST_ON_TAB_KEY) != null) {
+            context.offsetMap.tryGetOffset(OLD_ARGUMENTS_REPLACEMENT_OFFSET)
+                ?.let { tailOffset = it }
+        }
+        return tailOffset
+    }
+}
+
+internal val KEEP_OLD_ARGUMENT_LIST_ON_TAB_KEY = Key<Unit>("KEEP_OLD_ARGUMENT_LIST_ON_TAB_KEY")
+
+internal fun LookupElement.keepOldArgumentListOnTab(): LookupElement {
+    putUserData(KEEP_OLD_ARGUMENT_LIST_ON_TAB_KEY, Unit)
+    return this
+}
+
+/**
+ * This insertion handler is responsible for replacing the entire following argument if
+ * replacement completion is used (e.g., using the `\t` key).
+ * It uses the [OLD_ARGUMENTS_REPLACEMENT_OFFSET] to know the end offset of the argument we are currently replacing.
+ */
+@Serializable
+internal class SmartCompletionReplaceExistingArgumentHandler : SerializableInsertHandler {
+    override fun handleInsert(
+        context: InsertionContext,
+        item: LookupElement
+    ) {
+        if (context.completionChar == Lookup.REPLACE_SELECT_CHAR) {
+            val offset = context.offsetMap.tryGetOffset(OLD_ARGUMENTS_REPLACEMENT_OFFSET)
+            if (offset != null) {
+                context.document.deleteString(context.tailOffset, offset)
+            }
+        }
+        item.handleInsert(context)
+    }
+
+}
+
+internal fun OffsetMap.tryGetOffset(key: OffsetKey): Int? {
+    try {
+        if (!containsOffset(key)) return null
+        return getOffset(key).takeIf { it != -1 } // prior to IDEA 2016.3 getOffset() returned -1 if not found, now it throws exception
+    } catch (_: Exception) {
+        return null
+    }
+}
+
+private fun isAtEndOfLine(offset: Int, document: Document): Boolean {
+    var i = offset
+    val chars = document.charsSequence
+    while (i < chars.length) {
+        val c = chars[i]
+        if (c == '\n') return true
+        if (!Character.isWhitespace(c)) return false
+        i++
+    }
+    return true
+}
+
+private fun replacementOffsetByExpression(expression: KtExpression): Int {
+    when (expression) {
+        is KtCallExpression -> {
+            val calleeExpression = expression.calleeExpression
+            if (calleeExpression != null) {
+                return calleeExpression.textRange!!.endOffset
+            }
+        }
+
+        is KtQualifiedExpression -> {
+            val selector = expression.selectorExpression
+            if (selector != null) {
+                return replacementOffsetByExpression(selector)
+            }
+        }
+    }
+    return expression.textRange!!.endOffset
+}

plugins/kotlin/completion/impl-k2/src/org/jetbrains/kotlin/idea/completion/impl/k2/lookups/factories/VariableLookupElementFactory.kt

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.idea.base.analysis.api.utils.shortenReferencesInRang
 import org.jetbrains.kotlin.idea.base.analysis.withRootPrefixIfNeeded
 import org.jetbrains.kotlin.idea.base.codeInsight.KotlinIconProvider.getIconFor
 import org.jetbrains.kotlin.idea.base.serialization.names.KotlinNameSerializer
+import org.jetbrains.kotlin.idea.completion.impl.k2.contributors.MultipleArgumentsInsertHandler
 import org.jetbrains.kotlin.idea.completion.impl.k2.handlers.Tail
 import org.jetbrains.kotlin.idea.completion.impl.k2.lookups.CallableInsertionOptions
 import org.jetbrains.kotlin.idea.completion.impl.k2.lookups.CallableInsertionStrategy
@@ -126,6 +127,7 @@ internal object VariableLookupElementFactory {
         val lookupObject = MultipleArgumentsLookupObject(names.first(), renderedText, tail)
         val lookupElement = LookupElementBuilder.create(lookupObject, renderedText)
             .withIcon(compoundIcon)
+            .withInsertHandler(MultipleArgumentsInsertHandler())
 
         lookupElement.matchesExpectedType = ExpectedTypeWeigher.MatchesExpectedType.MATCHES_PREFERRED