[REPL] Fix completion for function arguments

plugins/scripting/scripting-ide-services-test/test/org/jetbrains/kotlin/scripting/ide_services/ReplCompletionAndErrorsAnalysisTest.kt

@@ -113,10 +113,13 @@ class ReplCompletionAndErrorsAnalysisTest : TestCase() {
     fun testFunctionArgumentNames() = test {
         run {
             doCompile
-            code = """fun _sf(_someInt: Int = 42, _someString: String = "s") = 1"""
+            code = """
+                fun _sf(_someInt: Int = 42, _someString: String = "s") = 1
+                fun String.f(_bar: Int) = _bar
+                class C(val _xyz: Int)
+            """.trimIndent()
         }
         run {
-            doComplete
             code = """_sf(_s"""
             cursor = code.length
             expect {
@@ -125,6 +128,38 @@ class ReplCompletionAndErrorsAnalysisTest : TestCase() {
                 addCompletion("_someString = ", "_someString", "String", "parameter")
             }
         }
+        run {
+            code = """ "my string".f(_b"""
+            cursor = code.length
+            expect {
+                addCompletion("_bar = ", "_bar", "Int", "parameter")
+            }
+        }
+        run {
+            code = "C(_x"
+            cursor = code.length
+            expect {
+                addCompletion("_xyz = ", "_xyz", "Int", "parameter")
+            }
+        }
+    }
+
+    @Test
+    fun testCompletionInsideFunctions() = test {
+        run {
+            val statement = "val a = _f"
+            code = """
+                fun dontCompleteMe(_foo: Int, bar: String) {
+                    val _foo2 = ""
+                    $statement
+                }
+            """.trimIndent()
+            cursor = code.indexOf(statement) + statement.length
+            expect {
+                addCompletion("_foo2", "_foo2", "String", "property")
+                addCompletion("_foo", "_foo", "Int", "parameter")
+            }
+        }
     }
 
     @Test

plugins/scripting/scripting-ide-services/src/org/jetbrains/kotlin/scripting/ide_services/compiler/impl/KJvmReplCompleter.kt

@@ -31,7 +31,6 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.annotations.argumentValue
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.resolve.scopes.LexicalScope
 import org.jetbrains.kotlin.resolve.scopes.MemberScope.Companion.ALL_NAME_FILTER
 import org.jetbrains.kotlin.scripting.ide_services.compiler.completion
 import org.jetbrains.kotlin.scripting.ide_services.compiler.filterOutShadowedDescriptors
@@ -41,8 +40,6 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.asFlexibleType
 import org.jetbrains.kotlin.types.isFlexible
 import java.io.File
-import java.lang.IllegalArgumentException
-import java.util.*
 import kotlin.script.experimental.api.ScriptCompilationConfiguration
 import kotlin.script.experimental.api.SourceCodeCompletionVariant
 
@@ -72,6 +69,12 @@ fun getKJvmCompletion(
 fun prepareCodeForCompletion(code: String, cursor: Int) =
     code.substring(0, cursor) + KJvmReplCompleter.INSERTED_STRING + code.substring(cursor)
 
+private inline fun <reified T> PsiElement.thisOrParent() = when {
+    this is T -> this
+    this.parent is T -> (this.parent as T)
+    else -> null
+}
+
 private class KJvmReplCompleter(
     private val ktScript: KtFile,
     private val bindingContext: BindingContext,
@@ -89,148 +92,162 @@ private class KJvmReplCompleter(
         return element
     }
 
-    fun getCompletion() = sequence<SourceCodeCompletionVariant> gen@{
+    private val getDescriptorsQualified = ResultGetter { element, options ->
-        val filterOutShadowedDescriptors = configuration[ScriptCompilationConfiguration.completion.filterOutShadowedDescriptors]!!
+        val expression = element.thisOrParent<KtQualifiedExpression>() ?: return@ResultGetter null
-        val nameFilter = configuration[ScriptCompilationConfiguration.completion.nameFilter]!!
 
-        val element = getElementAt(cursor)
+        val receiverExpression = expression.receiverExpression
+        val expressionType = bindingContext.get(
+            BindingContext.EXPRESSION_TYPE_INFO,
+            receiverExpression
+        )?.type
 
-        var descriptors: Collection<DeclarationDescriptor>? = null
+        DescriptorsResult(targetElement = expression).apply {
-        var isTipsManagerCompletion = true
+            if (expressionType != null) {
-        var isSortNeeded = true
+                sortNeeded = false
+                descriptors.addAll(
+                    getVariantsHelper { true }
+                        .getReferenceVariants(
+                            receiverExpression,
+                            CallTypeAndReceiver.DOT(receiverExpression),
+                            DescriptorKindFilter.ALL,
+                            ALL_NAME_FILTER,
+                            filterOutShadowed = options.filterOutShadowedDescriptors,
+                        )
+                )
+            }
+        }
+    }
 
-        if (element == null)
+    private val getDescriptorsSimple = ResultGetter { element, options ->
-            return@gen
+        val expression = element.thisOrParent<KtSimpleNameExpression>() ?: return@ResultGetter null
 
-        val simpleExpression = when {
+        val result = DescriptorsResult(targetElement = expression)
-            element is KtSimpleNameExpression -> element
+        val inDescriptor: DeclarationDescriptor = expression.getResolutionScope(bindingContext, resolutionFacade).ownerDescriptor
-            element.parent is KtSimpleNameExpression -> element.parent as KtSimpleNameExpression
+        val prefix = element.text.substring(0, cursor - element.startOffset)
-            else -> null
+
+        val elementParent = element.parent
+        if (prefix.isEmpty() && elementParent is KtBinaryExpression) {
+            val parentChildren = elementParent.children
+            if (parentChildren.size == 3 &&
+                parentChildren[1] is KtOperationReferenceExpression &&
+                parentChildren[1].text == INSERTED_STRING
+            ) return@ResultGetter result
         }
 
-        if (simpleExpression != null) {
+        val containingArgument = expression.thisOrParent<KtValueArgument>()
-            val inDescriptor: DeclarationDescriptor = simpleExpression.getResolutionScope(bindingContext, resolutionFacade).ownerDescriptor
+        val containingCall = containingArgument?.getParentOfType<KtCallExpression>(true)
-            val prefix = element.text.substring(0, cursor - element.startOffset)
+        val containingQualifiedExpression = containingCall?.parent as? KtDotQualifiedExpression
+        val containingCallId = containingCall?.calleeExpression?.text
+        fun Name.test(checkAgainstContainingCall: Boolean): Boolean {
+            if (isSpecial) return false
+            if (options.nameFilter(identifier, prefix)) return true
+            return checkAgainstContainingCall && containingCallId?.let { options.nameFilter(identifier, it) } == true
+        }
 
-            val elementParent = element.parent
+        DescriptorsResult(targetElement = element).apply {
-            if (prefix.isEmpty() && elementParent is KtBinaryExpression) {
+            sortNeeded = false
-                val parentChildren = elementParent.children
-                if (parentChildren.size == 3 &&
-                    parentChildren[1] is KtOperationReferenceExpression &&
-                    parentChildren[1].text == INSERTED_STRING
-                ) return@gen
-            }
 
-            val containingCallId = simpleExpression.getParentOfType<KtCallExpression>(true)?.calleeExpression?.text
+            descriptors.apply {
-            fun Name.test(checkAgainstContainingCall: Boolean): Boolean {
+                fun addParameters(descriptor: DeclarationDescriptor) {
-                if (isSpecial) return false
+                    if (containingCallId == descriptor.name.identifier) {
-                if (nameFilter(identifier, prefix)) return true
+                        val params = when (descriptor) {
-                return checkAgainstContainingCall && containingCallId?.let { nameFilter(identifier, it) } == true
+                            is CallableDescriptor -> descriptor.valueParameters
-            }
+                            is ClassDescriptor -> descriptor.constructors.flatMap { it.valueParameters }
+                            else -> emptyList()
+                        }
+                        val valueParams = params.filter { it.name.test(false) }
+                        addAll(valueParams)
+                        containingCallParameters.addAll(valueParams)
+                    }
+                }
 
-            isSortNeeded = false
+                getVariantsHelper(
-            descriptors = ArrayList<DeclarationDescriptor>().also { result ->
-                ReferenceVariantsHelper(
-                    bindingContext,
-                    resolutionFacade,
-                    moduleDescriptor,
                     VisibilityFilter(inDescriptor)
                 ).getReferenceVariants(
-                    simpleExpression,
+                    expression,
                     DescriptorKindFilter.ALL,
                     { it.test(true) },
                     filterOutJavaGettersAndSetters = true,
-                    filterOutShadowed = filterOutShadowedDescriptors, // setting to true makes it slower up to 4 times
+                    filterOutShadowed = options.filterOutShadowedDescriptors, // setting to true makes it slower up to 4 times
                     excludeNonInitializedVariable = true,
                     useReceiverType = null
                 ).forEach { descriptor ->
-                    if (descriptor.name.test(false)) result.add(descriptor)
+                    if (descriptor.name.test(false)) add(descriptor)
-                    if (descriptor is CallableDescriptor && containingCallId == descriptor.name.identifier) {
+                    addParameters(descriptor)
-                        descriptor.valueParameters.filterTo(result) { it.name.test(false) }
-                    }
                 }
-            }
 
-        } else if (element is KtStringTemplateExpression) {
+                if (containingQualifiedExpression != null) {
-            if (element.hasInterpolation()) {
+                    val receiverExpression = containingQualifiedExpression.receiverExpression
-                return@gen
+                    getVariantsHelper { true }
-            }
+                        .getReferenceVariants(
-
+                            receiverExpression,
-            val stringVal = element.entries.joinToString("") {
+                            CallTypeAndReceiver.DOT(receiverExpression),
-                val t = it.text
+                            DescriptorKindFilter.CALLABLES,
-                if (it.startOffset <= cursor && cursor <= it.endOffset) {
+                            ALL_NAME_FILTER,
-                    val s = cursor - it.startOffset
+                            filterOutShadowed = options.filterOutShadowedDescriptors,
-                    val e = s + INSERTED_STRING.length
+                        )
-                    t.substring(0, s) + t.substring(e)
+                        .forEach { descriptor ->
-                } else t
+                            addParameters(descriptor)
-            }
+                        }
-
-            val separatorIndex = stringVal.lastIndexOfAny(charArrayOf('/', '\\'))
-            val dir = if (separatorIndex != -1) {
-                stringVal.substring(0, separatorIndex + 1)
-            } else {
-                "."
-            }
-            val namePrefix = stringVal.substring(separatorIndex + 1)
-
-            val file = File(dir)
-
-            file.listFiles { p, f -> p == file && f.startsWith(namePrefix, true) }?.forEach {
-                yield(SourceCodeCompletionVariant(it.name, it.name, "file", "file"))
-            }
-
-            return@gen
-
-        } else {
-            isTipsManagerCompletion = false
-            val resolutionScope: LexicalScope?
-            val parent = element.parent
-
-            val qualifiedExpression = when {
-                element is KtQualifiedExpression -> {
-                    isTipsManagerCompletion = true
-                    element
                 }
-                parent is KtQualifiedExpression -> parent
-                else -> null
-            }
-
-            if (qualifiedExpression != null) {
-                val receiverExpression = qualifiedExpression.receiverExpression
-                val expressionType = bindingContext.get(
-                    BindingContext.EXPRESSION_TYPE_INFO,
-                    receiverExpression
-                )?.type
-                if (expressionType != null) {
-                    isSortNeeded = false
-                    descriptors = ReferenceVariantsHelper(
-                        bindingContext,
-                        resolutionFacade,
-                        moduleDescriptor,
-                        { true }
-                    ).getReferenceVariants(
-                        receiverExpression,
-                        CallTypeAndReceiver.DOT(receiverExpression),
-                        DescriptorKindFilter.ALL,
-                        ALL_NAME_FILTER,
-                        filterOutShadowed = filterOutShadowedDescriptors,
-                    )
-                }
-            } else {
-                resolutionScope = bindingContext.get(
-                    BindingContext.LEXICAL_SCOPE,
-                    element as KtExpression?
-                )
-                descriptors = (resolutionScope?.getContributedDescriptors(
-                    DescriptorKindFilter.ALL,
-                    ALL_NAME_FILTER
-                )
-                    ?: return@gen)
             }
         }
+    }
 
-        if (descriptors != null) {
+    private val getDescriptorsString = ResultGetter { element, _ ->
-            val targetElement = if (isTipsManagerCompletion) element else element.parent
+        if (element !is KtStringTemplateExpression) return@ResultGetter null
+
+        val stringVal = element.entries.joinToString("") {
+            val t = it.text
+            if (it.startOffset <= cursor && cursor <= it.endOffset) {
+                val s = cursor - it.startOffset
+                val e = s + INSERTED_STRING.length
+                t.substring(0, s) + t.substring(e)
+            } else t
+        }
+
+        val separatorIndex = stringVal.lastIndexOfAny(charArrayOf('/', '\\'))
+        val dir = if (separatorIndex != -1) {
+            stringVal.substring(0, separatorIndex + 1)
+        } else {
+            "."
+        }
+        val namePrefix = stringVal.substring(separatorIndex + 1)
+
+        val file = File(dir)
+        DescriptorsResult(targetElement = element).also { result ->
+            result.variants = sequence {
+                file.listFiles { p, f -> p == file && f.startsWith(namePrefix, true) }?.forEach {
+                    yield(SourceCodeCompletionVariant(it.name, it.name, "file", "file"))
+                }
+            }
+        }
+    }
+
+    private val getDescriptorsDefault = ResultGetter { element, _ ->
+        val resolutionScope = bindingContext.get(
+            BindingContext.LEXICAL_SCOPE,
+            element as KtExpression?
+        )
+        DescriptorsResult(targetElement = element).also { result ->
+            resolutionScope?.getContributedDescriptors(
+                DescriptorKindFilter.ALL,
+                ALL_NAME_FILTER
+            )?.let { descriptors ->
+                result.descriptors.addAll(descriptors)
+            }
+        }
+    }
+
+    private fun renderResult(
+        element: PsiElement,
+        options: DescriptorsOptions,
+        result: DescriptorsResult?
+    ): Sequence<SourceCodeCompletionVariant> {
+        if (result == null) return emptySequence()
+        result.variants?.let { return it }
+
+        with(result) {
             val prefixEnd = cursor - targetElement.startOffset
             var prefix = targetElement.text.substring(0, prefixEnd)
 
@@ -243,67 +260,109 @@ private class KJvmReplCompleter(
                 prefix.substring(0, cursorWithinElement)
             }
 
-            if (descriptors !is ArrayList<*>) {
+            return sequence {
-                descriptors = ArrayList(descriptors)
+                descriptors
-            }
+                    .map {
-
+                        val presentation =
-            (descriptors as ArrayList<DeclarationDescriptor>)
+                            getPresentation(
-                .map {
+                                it, result.containingCallParameters
-                    val presentation =
-                        getPresentation(
-                            it
-                        )
-                    Triple(it, presentation, (presentation.presentableText + presentation.tailText).lowercase())
-                }
-                .let {
-                    if (isSortNeeded) it.sortedBy { descTriple -> descTriple.third } else it
-                }
-                .forEach { resultTriple ->
-                    val descriptor = resultTriple.first
-                    val (rawName, presentableText, tailText, completionText) = resultTriple.second
-                    if (nameFilter(rawName, prefix)) {
-                        val fullName: String =
-                            formatName(
-                                presentableText
                             )
-                        val deprecationLevel = descriptor.annotations
+                        Triple(it, presentation, (presentation.presentableText + presentation.tailText).lowercase())
-                            .findAnnotation(FqName("kotlin.Deprecated"))
+                    }
-                            ?.let { annotationDescriptor ->
+                    .let {
-                                val valuePair = annotationDescriptor.argumentValue("level")?.value as? Pair<*, *>
+                        if (sortNeeded) it.sortedBy { descTriple -> descTriple.third } else it
-                                val valueClass = (valuePair?.first as? ClassId)?.takeIf { DeprecationLevel::class.classId == it }
+                    }
-                                val valueName = (valuePair?.second as? Name)?.identifier
+                    .forEach { resultTriple ->
-                                if (valueClass == null || valueName == null) return@let DeprecationLevel.WARNING
+                        val descriptor = resultTriple.first
-                                DeprecationLevel.valueOf(valueName)
+                        val (rawName, presentableText, tailText, completionText) = resultTriple.second
-                            }
+                        if (options.nameFilter(rawName, prefix)) {
-                        yield(
+                            val fullName: String =
-                            SourceCodeCompletionVariant(
+                                formatName(
-                                completionText,
+                                    presentableText
-                                fullName,
+                                )
-                                tailText,
+                            val deprecationLevel = descriptor.annotations
-                                getIconFromDescriptor(
+                                .findAnnotation(FqName("kotlin.Deprecated"))
-                                    descriptor
+                                ?.let { annotationDescriptor ->
-                                ),
+                                    val valuePair = annotationDescriptor.argumentValue("level")?.value as? Pair<*, *>
-                                deprecationLevel,
+                                    val valueClass = (valuePair?.first as? ClassId)?.takeIf { DeprecationLevel::class.classId == it }
-                            )
+                                    val valueName = (valuePair?.second as? Name)?.identifier
-                        )
+                                    if (valueClass == null || valueName == null) return@let DeprecationLevel.WARNING
+                                    DeprecationLevel.valueOf(valueName)
+                                }
+                            yield(
+                                SourceCodeCompletionVariant(
+                                    completionText,
+                                    fullName,
+                                    tailText,
+                                    getIconFromDescriptor(
+                                        descriptor
+                                    ),
+                                    deprecationLevel,
+                                )
+                            )
+                        }
                     }
-                }
 
-            yieldAll(
+                yieldAll(
-                keywordsCompletionVariants(
+                    keywordsCompletionVariants(
-                    KtTokens.KEYWORDS,
+                        KtTokens.KEYWORDS,
-                    prefix
+                        prefix
+                    )
                 )
-            )
+                yieldAll(
-            yieldAll(
+                    keywordsCompletionVariants(
-                keywordsCompletionVariants(
+                        KtTokens.SOFT_KEYWORDS,
-                    KtTokens.SOFT_KEYWORDS,
+                        prefix
-                    prefix
+                    )
                 )
-            )
+            }
         }
     }
 
+    fun getCompletion(): Sequence<SourceCodeCompletionVariant> {
+        val filterOutShadowedDescriptors = configuration[ScriptCompilationConfiguration.completion.filterOutShadowedDescriptors]!!
+        val nameFilter = configuration[ScriptCompilationConfiguration.completion.nameFilter]!!
+        val options = DescriptorsOptions(
+            nameFilter, filterOutShadowedDescriptors
+        )
+
+        val element = getElementAt(cursor) ?: return emptySequence()
+
+        val descriptorsGetters = listOf(
+            getDescriptorsSimple,
+            getDescriptorsString,
+            getDescriptorsQualified,
+            getDescriptorsDefault,
+        )
+
+        val result = descriptorsGetters.firstNotNullOfOrNull { it.get(element, options) }
+        return renderResult(element, options, result)
+    }
+
+    private fun getVariantsHelper(visibilityFilter: (DeclarationDescriptor) -> Boolean) = ReferenceVariantsHelper(
+        bindingContext,
+        resolutionFacade,
+        moduleDescriptor,
+        visibilityFilter,
+    )
+
+    private fun interface ResultGetter {
+        fun get(element: PsiElement, options: DescriptorsOptions): DescriptorsResult?
+    }
+
+    private class DescriptorsResult(
+        val descriptors: MutableList<DeclarationDescriptor> = mutableListOf(),
+        var variants: Sequence<SourceCodeCompletionVariant>? = null,
+        var sortNeeded: Boolean = true,
+        var targetElement: PsiElement,
+        val containingCallParameters: MutableList<ValueParameterDescriptor> = mutableListOf(),
+    )
+
+    private class DescriptorsOptions(
+        val nameFilter: (String, String) -> Boolean,
+        val filterOutShadowedDescriptors: Boolean,
+    )
+
     private class VisibilityFilter(
         private val inDescriptor: DeclarationDescriptor
     ) : (DeclarationDescriptor) -> Boolean {
@@ -396,7 +455,10 @@ private class KJvmReplCompleter(
             val completionText: String
         )
 
-        fun getPresentation(descriptor: DeclarationDescriptor): DescriptorPresentation {
+        fun getPresentation(
+            descriptor: DeclarationDescriptor,
+            callParameters: Collection<ValueParameterDescriptor>
+        ): DescriptorPresentation {
             val rawDescriptorName = descriptor.name.asString()
             val descriptorName = rawDescriptorName.quoteIfNeeded()
             var presentableText = descriptorName
@@ -426,7 +488,10 @@ private class KJvmReplCompleter(
                 val outType =
                     descriptor.type
                 typeText = RENDERER.renderType(outType)
-                if (descriptor is ValueParameterDescriptor) {
+                if (
+                    descriptor is ValueParameterDescriptor &&
+                    callParameters.contains(descriptor)
+                ) {
                     completionText = "$rawDescriptorName = "
                 }
             } else if (descriptor is ClassDescriptor) {