Keyword completion: use proper prefix matcher for keyword completion

fixes #KT-32519

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/BasicCompletionSession.kt

@@ -556,7 +556,7 @@ class BasicCompletionSession(
 
             val keywordsPrefix = prefix.substringBefore('@') // if there is '@' in the prefix - use shorter prefix to not loose 'this' etc
             val isUseSiteAnnotationTarget = position.prevLeaf()?.node?.elementType == KtTokens.AT
-            KeywordCompletion.complete(expression ?: parameters.position, keywordsPrefix, isJvmModule) { lookupElement ->
+            KeywordCompletion.complete(expression ?: parameters.position, resultSet.prefixMatcher, isJvmModule) { lookupElement ->
                 val keyword = lookupElement.lookupString
                 if (keyword in keywordsToSkip) return@complete
 

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/KeywordCompletion.kt

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.idea.completion
 
 import com.intellij.codeInsight.completion.InsertHandler
 import com.intellij.codeInsight.completion.InsertionContext
+import com.intellij.codeInsight.completion.PrefixMatcher
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.openapi.module.ModuleUtilCore
@@ -95,7 +96,7 @@ object KeywordCompletion {
         SET_KEYWORD
     ).map { it.value } + "companion object"
 
-    fun complete(position: PsiElement, prefix: String, isJvmModule: Boolean, consumer: (LookupElement) -> Unit) {
+    fun complete(position: PsiElement, prefixMatcher: PrefixMatcher, isJvmModule: Boolean, consumer: (LookupElement) -> Unit) {
         if (!GENERAL_FILTER.isAcceptable(position, position)) return
 
         val parserFilter = buildFilter(position)
@@ -122,8 +123,7 @@ object KeywordCompletion {
 
             if (keywordToken == DYNAMIC_KEYWORD && isJvmModule) continue // not supported for JVM
 
-            // we use simple matching by prefix, not prefix matcher from completion
-            if (!keyword.startsWith(prefix) && keywordToken !in KEYWORDS_TO_IGNORE_PREFIX) continue
+            if (keywordToken !in KEYWORDS_TO_IGNORE_PREFIX && !prefixMatcher.isStartMatch(keyword)) continue
 
             if (!parserFilter(keywordToken)) continue
 

idea/idea-completion/testData/keywords/FixingLayoutAndReturn.kt

@@ -0,0 +1,5 @@
+fun foo() {
+    куе<caret>
+}
+
+// EXIST: { lookupString: "return", itemText: "return", tailText: null, attributes: "bold" }

idea/idea-completion/tests/org/jetbrains/kotlin/idea/completion/test/KeywordCompletionTestGenerated.java

@@ -149,6 +149,11 @@ public class KeywordCompletionTestGenerated extends AbstractKeywordCompletionTes
         runTest("idea/idea-completion/testData/keywords/FileKeyword.kt");
     }
 
+    @TestMetadata("FixingLayoutAndReturn.kt")
+    public void testFixingLayoutAndReturn() throws Exception {
+        runTest("idea/idea-completion/testData/keywords/FixingLayoutAndReturn.kt");
+    }
+
     @TestMetadata("GlobalPropertyAccessors.kt")
     public void testGlobalPropertyAccessors() throws Exception {
         runTest("idea/idea-completion/testData/keywords/GlobalPropertyAccessors.kt");